113 lines
3.1 KiB
Go
113 lines
3.1 KiB
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)
|
|
}
|