Initial orientator

internal/detector/assets/README.txt

@@ -0,0 +1,2 @@
+Place dedoc_orientation.onnx here.
+ONNX runtime shared libs go in onnxruntime/<goos>-<goarch>/.

internal/detector/assets_darwin_arm64.go

@@ -0,0 +1,10 @@
+//go:build darwin && arm64
+
+package detector
+
+import _ "embed"
+
+//go:embed assets/onnxruntime/darwin-arm64/libonnxruntime.dylib
+var runtimeLibBytes []byte
+
+const runtimeLibFileName = "libonnxruntime.dylib"

internal/detector/assets_linux_amd64.go

@@ -0,0 +1,10 @@
+//go:build linux && amd64
+
+package detector
+
+import _ "embed"
+
+//go:embed assets/onnxruntime/linux-amd64/libonnxruntime.so
+var runtimeLibBytes []byte
+
+const runtimeLibFileName = "libonnxruntime.so"

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
+}

internal/detector/preprocess.go

@@ -0,0 +1,53 @@
+package detector
+
+import (
+	"image"
+	"image/color"
+	stddraw "image/draw"
+
+	xdraw "golang.org/x/image/draw"
+)
+
+const ImgSize = 1200
+
+// Preprocess implements the Python reference exactly:
+// 1. resize keeping aspect ratio to fit in 1200x1200 (Lanczos-equivalent CatmullRom)
+// 2. paste onto white 1200x1200 canvas at (0,0)
+// 3. normalize: (x/255 - 0.5) / 0.5
+// 4. HWC -> CHW with batch dim -> [1, 3, 1200, 1200]
+func Preprocess(img image.Image) []float32 {
+	bounds := img.Bounds()
+	w, h := bounds.Dx(), bounds.Dy()
+
+	scale := float64(ImgSize) / float64(w)
+	if s := float64(ImgSize) / float64(h); s < scale {
+		scale = s
+	}
+	newW := int(float64(w) * scale)
+	newH := int(float64(h) * scale)
+	if newW < 1 {
+		newW = 1
+	}
+	if newH < 1 {
+		newH = 1
+	}
+
+	canvas := image.NewRGBA(image.Rect(0, 0, ImgSize, ImgSize))
+	stddraw.Draw(canvas, canvas.Bounds(), &image.Uniform{C: color.RGBA{255, 255, 255, 255}}, image.Point{}, stddraw.Src)
+	xdraw.CatmullRom.Scale(canvas, image.Rect(0, 0, newW, newH), img, bounds, xdraw.Over, nil)
+
+	out := make([]float32, 3*ImgSize*ImgSize)
+	plane := ImgSize * ImgSize
+	pix := canvas.Pix
+	for y := 0; y < ImgSize; y++ {
+		row := y * canvas.Stride
+		base := y * ImgSize
+		for x := 0; x < ImgSize; x++ {
+			i := row + x*4
+			out[base+x] = float32(pix[i+0])/127.5 - 1
+			out[plane+base+x] = float32(pix[i+1])/127.5 - 1
+			out[2*plane+base+x] = float32(pix[i+2])/127.5 - 1
+		}
+	}
+	return out
+}