64 lines
1.8 KiB
Go
64 lines
1.8 KiB
Go
package spike
|
|
|
|
import (
|
|
"fmt"
|
|
|
|
ort "github.com/yalue/onnxruntime_go"
|
|
)
|
|
|
|
const (
|
|
poseInputSize = 640
|
|
poseInputCount = 1 * 3 * poseInputSize * poseInputSize
|
|
poseOutputCount = 1 * 56 * 8400
|
|
)
|
|
|
|
// 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)
|
|
}
|
|
|
|
ort.SetSharedLibraryPath(runtimeDLLPath)
|
|
if err := ort.InitializeEnvironment(); err != nil {
|
|
return nil, fmt.Errorf("initialize ONNX Runtime: %w", err)
|
|
}
|
|
defer func() { _ = ort.DestroyEnvironment() }()
|
|
|
|
inputTensor, err := ort.NewTensor(ort.NewShape(1, 3, poseInputSize, poseInputSize), input)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("create pose input tensor: %w", err)
|
|
}
|
|
defer func() { _ = inputTensor.Destroy() }()
|
|
|
|
outputTensor, err := ort.NewEmptyTensor[float32](ort.NewShape(1, 56, 8400))
|
|
if err != nil {
|
|
return nil, fmt.Errorf("create pose output tensor: %w", err)
|
|
}
|
|
defer func() { _ = outputTensor.Destroy() }()
|
|
|
|
session, err := ort.NewAdvancedSession(
|
|
modelPath,
|
|
[]string{"images"},
|
|
[]string{"output0"},
|
|
[]ort.Value{inputTensor},
|
|
[]ort.Value{outputTensor},
|
|
nil,
|
|
)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("open pose ONNX session: %w", err)
|
|
}
|
|
defer func() { _ = session.Destroy() }()
|
|
|
|
if err := session.Run(); err != nil {
|
|
return nil, fmt.Errorf("run pose ONNX session: %w", err)
|
|
}
|
|
|
|
output := outputTensor.GetData()
|
|
if len(output) != poseOutputCount {
|
|
return nil, fmt.Errorf("pose output length = %d, want %d", len(output), poseOutputCount)
|
|
}
|
|
return append([]float32(nil), output...), nil
|
|
}
|