feat(v2): add fall event regression engine

2026-07-22 16:18:16 +08:00
parent 258fa49259
commit 1298369054
22 changed files with 1214 additions and 92 deletions
@@ -12,6 +12,90 @@ const (
 	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.
@@ -19,45 +103,10 @@ func RunPose(modelPath, runtimeDLLPath string, input []float32) ([]float32, erro
 	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)
+	runtime, err := OpenRuntime(modelPath, runtimeDLLPath)
 	if err != nil {
-		return nil, fmt.Errorf("create pose input tensor: %w", err)
+		return nil, 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
+	defer runtime.Close()
+	return runtime.Run(input)
 }