Initial orientator

2026-05-13 19:52:51 +03:00
commit 6dc5080ee9
23 changed files with 1920 additions and 0 deletions
--- a/internal/detector/detector.go
+++ b/internal/detector/detector.go
@@ -0,0 +1,181 @@
+package detector
+
+import (
+	_ "embed"
+	"fmt"
+	"image"
+	"math"
+	"os"
+	"path/filepath"
+	"runtime"
+	"sync"
+
+	ort "github.com/yalue/onnxruntime_go"
+)
+
+// Angles map to model's [2:6] logits in the same order as the Python reference:
+// indices [0,1,2,3] -> degrees [0, 270, 180, 90].
+var dedocMap = [4]int{0, 270, 180, 90}
+
+//go:embed assets/dedoc_orientation.onnx
+var modelBytes []byte
+
+type Detector struct {
+	session *ort.AdvancedSession
+	input   *ort.Tensor[float32]
+	output  *ort.Tensor[float32]
+	mu      sync.Mutex // ONNX session is not goroutine-safe under concurrent Run
+}
+
+var (
+	once     sync.Once
+	instance *Detector
+	initErr  error
+)
+
+// Get returns the lazily-initialized singleton detector.
+func Get() (*Detector, error) {
+	once.Do(func() {
+		instance, initErr = newDetector()
+	})
+	return instance, initErr
+}
+
+func newDetector() (*Detector, error) {
+	libPath, err := extractRuntimeLib()
+	if err != nil {
+		return nil, fmt.Errorf("extract runtime: %w", err)
+	}
+	ort.SetSharedLibraryPath(libPath)
+	if err := ort.InitializeEnvironment(); err != nil {
+		return nil, fmt.Errorf("init ort env: %w", err)
+	}
+
+	modelPath, err := extractModel()
+	if err != nil {
+		return nil, fmt.Errorf("extract model: %w", err)
+	}
+
+	inputs, outputs, err := ort.GetInputOutputInfo(modelPath)
+	if err != nil {
+		return nil, fmt.Errorf("inspect model: %w", err)
+	}
+	if len(inputs) == 0 || len(outputs) == 0 {
+		return nil, fmt.Errorf("model has no inputs/outputs")
+	}
+
+	inShape := ort.NewShape(1, 3, ImgSize, ImgSize)
+	inTensor, err := ort.NewEmptyTensor[float32](inShape)
+	if err != nil {
+		return nil, err
+	}
+	outShape := ort.NewShape(1, 6)
+	outTensor, err := ort.NewEmptyTensor[float32](outShape)
+	if err != nil {
+		return nil, err
+	}
+
+	opts, err := ort.NewSessionOptions()
+	if err != nil {
+		return nil, err
+	}
+	defer opts.Destroy()
+	_ = opts.SetIntraOpNumThreads(runtime.NumCPU())
+
+	sess, err := ort.NewAdvancedSession(
+		modelPath,
+		[]string{inputs[0].Name},
+		[]string{outputs[0].Name},
+		[]ort.ArbitraryTensor{inTensor},
+		[]ort.ArbitraryTensor{outTensor},
+		opts,
+	)
+	if err != nil {
+		return nil, fmt.Errorf("create session: %w", err)
+	}
+
+	return &Detector{
+		session: sess,
+		input:   inTensor,
+		output:  outTensor,
+	}, nil
+}
+
+// Predict returns the rotation needed to right the image and the confidence.
+// Output angle ∈ {0, 90, 180, 270}; meaning matches the Python reference.
+func (d *Detector) Predict(img image.Image) (angle int, confidence float64, err error) {
+	if img == nil {
+		return 0, 0, fmt.Errorf("nil image")
+	}
+	data := Preprocess(img)
+
+	d.mu.Lock()
+	defer d.mu.Unlock()
+
+	dst := d.input.GetData()
+	copy(dst, data)
+
+	if err := d.session.Run(); err != nil {
+		return 0, 0, fmt.Errorf("ort run: %w", err)
+	}
+	logits := d.output.GetData() // shape [1,6]
+	// Use only orientation logits at index [2:6].
+	probs := softmax4(logits[2:6])
+	idx := 0
+	best := probs[0]
+	for i := 1; i < 4; i++ {
+		if probs[i] > best {
+			best = probs[i]
+			idx = i
+		}
+	}
+	return dedocMap[idx], float64(best), nil
+}
+
+func softmax4(x []float32) [4]float32 {
+	var max float32 = x[0]
+	for i := 1; i < 4; i++ {
+		if x[i] > max {
+			max = x[i]
+		}
+	}
+	var sum float32
+	var e [4]float32
+	for i := 0; i < 4; i++ {
+		e[i] = float32(math.Exp(float64(x[i] - max)))
+		sum += e[i]
+	}
+	for i := 0; i < 4; i++ {
+		e[i] /= sum
+	}
+	return e
+}
+
+func extractRuntimeLib() (string, error) {
+	return materialize("runtime", runtimeLibFileName, runtimeLibBytes)
+}
+
+func extractModel() (string, error) {
+	return materialize("model", "dedoc_orientation.onnx", modelBytes)
+}
+
+// materialize writes embedded bytes to the user cache dir, reusing the existing
+// file if its size already matches.
+func materialize(subdir, name string, data []byte) (string, error) {
+	cacheDir, err := os.UserCacheDir()
+	if err != nil {
+		cacheDir = os.TempDir()
+	}
+	dir := filepath.Join(cacheDir, "orientator", subdir, fmt.Sprintf("%s-%s", runtime.GOOS, runtime.GOARCH))
+	if err := os.MkdirAll(dir, 0o755); err != nil {
+		return "", err
+	}
+	out := filepath.Join(dir, name)
+	if st, err := os.Stat(out); err == nil && int(st.Size()) == len(data) {
+		return out, nil
+	}
+	if err := os.WriteFile(out, data, 0o644); err != nil {
+		return "", err
+	}
+	return out, nil
+}