orientator/internal/detector/detector.go

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
}