silver_pose/v2/internal/spike/ort.go

package spike

import (
	"fmt"

	ort "github.com/yalue/onnxruntime_go"
)

const (
	poseInputSize   = 640
	poseInputCount  = 1 * 3 * poseInputSize * poseInputSize
	poseOutputCount = 1 * 56 * 8400
)

// Runtime owns one reusable ONNX session.  V2 video replay must keep the
// model session alive across frames; recreating it for every frame would make
// event-latency measurements meaningless.
type Runtime struct {
	input       *ort.Tensor[float32]
	output      *ort.Tensor[float32]
	session     *ort.AdvancedSession
	initialized bool
	closed      bool
}

func OpenRuntime(modelPath, runtimeDLLPath string) (*Runtime, error) {
	ort.SetSharedLibraryPath(runtimeDLLPath)
	if err := ort.InitializeEnvironment(); err != nil {
		return nil, fmt.Errorf("initialize ONNX Runtime: %w", err)
	}
	runtime := &Runtime{initialized: true}
	var err error
	runtime.input, err = ort.NewEmptyTensor[float32](ort.NewShape(1, 3, poseInputSize, poseInputSize))
	if err != nil {
		runtime.Close()
		return nil, fmt.Errorf("create pose input tensor: %w", err)
	}
	runtime.output, err = ort.NewEmptyTensor[float32](ort.NewShape(1, 56, 8400))
	if err != nil {
		runtime.Close()
		return nil, fmt.Errorf("create pose output tensor: %w", err)
	}
	runtime.session, err = ort.NewAdvancedSession(
		modelPath,
		[]string{"images"},
		[]string{"output0"},
		[]ort.Value{runtime.input},
		[]ort.Value{runtime.output},
		nil,
	)
	if err != nil {
		runtime.Close()
		return nil, fmt.Errorf("open pose ONNX session: %w", err)
	}
	return runtime, nil
}

func (runtime *Runtime) Run(input []float32) ([]float32, error) {
	if runtime == nil || runtime.closed {
		return nil, fmt.Errorf("ONNX Runtime session is closed")
	}
	if len(input) != poseInputCount {
		return nil, fmt.Errorf("pose input length = %d, want %d", len(input), poseInputCount)
	}
	copy(runtime.input.GetData(), input)
	if err := runtime.session.Run(); err != nil {
		return nil, fmt.Errorf("run pose ONNX session: %w", err)
	}
	output := runtime.output.GetData()
	if len(output) != poseOutputCount {
		return nil, fmt.Errorf("pose output length = %d, want %d", len(output), poseOutputCount)
	}
	return append([]float32(nil), output...), nil
}

func (runtime *Runtime) Close() {
	if runtime == nil || runtime.closed {
		return
	}
	if runtime.session != nil {
		_ = runtime.session.Destroy()
		runtime.session = nil
	}
	if runtime.output != nil {
		_ = runtime.output.Destroy()
		runtime.output = nil
	}
	if runtime.input != nil {
		_ = runtime.input.Destroy()
		runtime.input = nil
	}
	if runtime.initialized {
		_ = ort.DestroyEnvironment()
		runtime.initialized = false
	}
	runtime.closed = true
}

// RunPose executes the locked YOLO pose ONNX graph once using the CPU runtime.
// The caller supplies absolute model and DLL paths so no camera credential or
// machine-specific path is retained in V2 configuration or source.
func RunPose(modelPath, runtimeDLLPath string, input []float32) ([]float32, error) {
	if len(input) != poseInputCount {
		return nil, fmt.Errorf("pose input length = %d, want %d", len(input), poseInputCount)
	}
	runtime, err := OpenRuntime(modelPath, runtimeDLLPath)
	if err != nil {
		return nil, err
	}
	defer runtime.Close()
	return runtime.Run(input)
}