feat(v2): add fall event regression engine
This commit is contained in:
@@ -0,0 +1,118 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"flag"
|
||||
"fmt"
|
||||
"io"
|
||||
"os"
|
||||
"os/exec"
|
||||
"time"
|
||||
|
||||
"silverpose/v2/internal/fall"
|
||||
"silverpose/v2/internal/pose"
|
||||
"silverpose/v2/internal/spike"
|
||||
)
|
||||
|
||||
func main() {
|
||||
ffmpegPath := flag.String("ffmpeg", "ffmpeg", "path to ffmpeg.exe")
|
||||
videoPath := flag.String("video", "..\\demo\\1.mp4", "local non-sensitive validation video")
|
||||
width := flag.Int("width", 848, "source video frame width")
|
||||
height := flag.Int("height", 480, "source video frame height")
|
||||
fps := flag.Float64("fps", 30, "replay timestamp frames per second")
|
||||
maxFrames := flag.Int("max-frames", 0, "maximum frames to process; zero reads the complete video")
|
||||
modelPath := flag.String("onnx", "assets\\best.onnx", "locked ONNX pose model path")
|
||||
runtimeDLLPath := flag.String("ort-dll", "", "absolute onnxruntime.dll path")
|
||||
flag.Parse()
|
||||
|
||||
if *runtimeDLLPath == "" || *width <= 0 || *height <= 0 || *fps <= 0 || *maxFrames < 0 {
|
||||
fail("--ort-dll is required; width, height and fps must be positive; max-frames cannot be negative")
|
||||
}
|
||||
engine, err := fall.NewEngine(fall.EngineConfig{
|
||||
KeypointConfidenceThreshold: 0.4,
|
||||
SuspectWindowSeconds: 0.5,
|
||||
ConfirmWindowSeconds: 1.8,
|
||||
RecoveryWindowSeconds: 2.0,
|
||||
CooldownSeconds: 10.0,
|
||||
RequireRapidDrop: false,
|
||||
RequireLowerBody: false,
|
||||
HorizontalAngleThresholdDegrees: 45.0,
|
||||
ConfigVersion: "v1-default-20260722",
|
||||
SessionID: "video-regression",
|
||||
})
|
||||
if err != nil {
|
||||
fail("create event engine: %v", err)
|
||||
}
|
||||
runtime, err := spike.OpenRuntime(*modelPath, *runtimeDLLPath)
|
||||
if err != nil {
|
||||
fail("open ONNX Pose runtime: %v", err)
|
||||
}
|
||||
defer runtime.Close()
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Minute)
|
||||
defer cancel()
|
||||
command := exec.CommandContext(ctx, *ffmpegPath, "-v", "error", "-i", *videoPath, "-f", "rawvideo", "-pix_fmt", "bgr24", "-")
|
||||
stdout, err := command.StdoutPipe()
|
||||
if err != nil {
|
||||
fail("open FFmpeg output: %v", err)
|
||||
}
|
||||
if err := command.Start(); err != nil {
|
||||
fail("start FFmpeg: %v", err)
|
||||
}
|
||||
|
||||
frameBytes := *width * *height * 3
|
||||
frame := make([]byte, frameBytes)
|
||||
frameCount := 0
|
||||
personFrames := 0
|
||||
personCount := 0
|
||||
events := make([]fall.Event, 0)
|
||||
limited := false
|
||||
for *maxFrames == 0 || frameCount < *maxFrames {
|
||||
_, err := io.ReadFull(stdout, frame)
|
||||
if errors.Is(err, io.EOF) || errors.Is(err, io.ErrUnexpectedEOF) {
|
||||
break
|
||||
}
|
||||
if err != nil {
|
||||
fail("read FFmpeg BGR frame: %v", err)
|
||||
}
|
||||
input, transform, err := pose.PreprocessBGR(frame, *width, *height, 640)
|
||||
if err != nil {
|
||||
fail("preprocess frame %d: %v", frameCount, err)
|
||||
}
|
||||
output, err := runtime.Run(input)
|
||||
if err != nil {
|
||||
fail("run ONNX Pose for frame %d: %v", frameCount, err)
|
||||
}
|
||||
people, err := pose.ParseYOLOv8Pose(output, transform, 0.25, 0.70)
|
||||
if err != nil {
|
||||
fail("parse ONNX Pose for frame %d: %v", frameCount, err)
|
||||
}
|
||||
result := engine.Process(fall.Frame{Timestamp: float64(frameCount) / *fps, Width: *width, Height: *height, Poses: people})
|
||||
if len(people) > 0 {
|
||||
personFrames++
|
||||
personCount += len(people)
|
||||
}
|
||||
events = append(events, result.Events...)
|
||||
frameCount++
|
||||
}
|
||||
if *maxFrames > 0 && frameCount == *maxFrames {
|
||||
limited = true
|
||||
cancel()
|
||||
}
|
||||
if err := command.Wait(); err != nil {
|
||||
if !limited {
|
||||
fail("FFmpeg ended with an error: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
fmt.Printf("REGRESSION_OK frames=%d person_frames=%d people=%d events=%d\n", frameCount, personFrames, personCount, len(events))
|
||||
for _, event := range events {
|
||||
fmt.Printf("EVENT id=%s track=%s latency=%.6f state=%s\n", event.EventID, event.TrackID, event.LatencySeconds, event.State)
|
||||
}
|
||||
}
|
||||
|
||||
func fail(format string, args ...any) {
|
||||
fmt.Fprintf(os.Stderr, "regression: "+format+"\n", args...)
|
||||
os.Exit(1)
|
||||
}
|
||||
+10
-2
@@ -8,6 +8,7 @@ import (
|
||||
"os/exec"
|
||||
"time"
|
||||
|
||||
"silverpose/v2/internal/pose"
|
||||
"silverpose/v2/internal/spike"
|
||||
)
|
||||
|
||||
@@ -41,7 +42,7 @@ func main() {
|
||||
fail("decoded BGR frame has %d bytes, want %d; verify --width/--height", len(raw), expectedBytes)
|
||||
}
|
||||
|
||||
input, err := spike.LetterboxBGRToNCHW(raw, *width, *height, 640)
|
||||
input, transform, err := pose.PreprocessBGR(raw, *width, *height, 640)
|
||||
if err != nil {
|
||||
fail("preprocess frame: %v", err)
|
||||
}
|
||||
@@ -49,6 +50,10 @@ func main() {
|
||||
if err != nil {
|
||||
fail("run ONNX Pose: %v", err)
|
||||
}
|
||||
people, err := pose.ParseYOLOv8Pose(output, transform, 0.25, 0.70)
|
||||
if err != nil {
|
||||
fail("parse ONNX Pose: %v", err)
|
||||
}
|
||||
|
||||
max := float32(0)
|
||||
for _, value := range output {
|
||||
@@ -56,7 +61,10 @@ func main() {
|
||||
max = value
|
||||
}
|
||||
}
|
||||
fmt.Printf("SPIKE_OK frame=%dx%d input=%d output=%d max=%.4f\n", *width, *height, len(input), len(output), max)
|
||||
fmt.Printf("SPIKE_OK frame=%dx%d input=%d output=%d people=%d max=%.4f\n", *width, *height, len(input), len(output), len(people), max)
|
||||
for index, person := range people {
|
||||
fmt.Printf("PERSON index=%d conf=%.6f box=%.3f,%.3f,%.3f,%.3f\n", index, person.Confidence, person.Box.Left, person.Box.Top, person.Box.Right, person.Box.Bottom)
|
||||
}
|
||||
}
|
||||
|
||||
func fail(format string, args ...any) {
|
||||
|
||||
@@ -0,0 +1,62 @@
|
||||
package fall
|
||||
|
||||
import (
|
||||
"math"
|
||||
|
||||
"silverpose/v2/internal/pose"
|
||||
)
|
||||
|
||||
var torsoJoints = [...]int{5, 6, 11, 12}
|
||||
var lowerJoints = [...]int{13, 14, 15, 16}
|
||||
|
||||
func assessPoseQuality(person pose.PersonPose, threshold float32, requireLowerBody bool) PoseQuality {
|
||||
if threshold < 0 || threshold > 1 {
|
||||
return PoseQuality{Reason: "invalid_threshold"}
|
||||
}
|
||||
visible := 0
|
||||
for _, keypoint := range person.Keypoints {
|
||||
if keypoint.Confidence >= threshold {
|
||||
visible++
|
||||
}
|
||||
}
|
||||
for _, index := range torsoJoints {
|
||||
if person.Keypoints[index].Confidence < threshold {
|
||||
return PoseQuality{Reason: "required_joint_low_confidence", VisibleJointCount: visible}
|
||||
}
|
||||
}
|
||||
if requireLowerBody {
|
||||
for _, index := range lowerJoints {
|
||||
if person.Keypoints[index].Confidence < threshold {
|
||||
return PoseQuality{Reason: "required_joint_low_confidence", VisibleJointCount: visible}
|
||||
}
|
||||
}
|
||||
}
|
||||
return PoseQuality{Accepted: true, Reason: "accepted", VisibleJointCount: visible}
|
||||
}
|
||||
|
||||
func extractEvidence(person pose.PersonPose, quality PoseQuality, previous *PoseEvidence, horizontalAngleThreshold float32) PoseEvidence {
|
||||
if !quality.Accepted {
|
||||
return PoseEvidence{Reason: quality.Reason}
|
||||
}
|
||||
shoulderX, shoulderY := midpoint(person, 5, 6)
|
||||
hipX, hipY := midpoint(person, 11, 12)
|
||||
vectorX := hipX - shoulderX
|
||||
vectorY := hipY - shoulderY
|
||||
torsoLength := float32(math.Hypot(float64(vectorX), float64(vectorY)))
|
||||
if torsoLength == 0 {
|
||||
return PoseEvidence{HipCenterY: hipY, HasHipCenterY: true, TorsoLength: 0, Reason: "degenerate_torso"}
|
||||
}
|
||||
cosine := math.Max(-1, math.Min(1, math.Abs(float64(vectorX))/float64(torsoLength)))
|
||||
angle := float32(math.Acos(cosine) * 180 / math.Pi)
|
||||
rapid := previous != nil && previous.Accepted && previous.HasHipCenterY && hipY-previous.HipCenterY >= 0.5*torsoLength
|
||||
return PoseEvidence{
|
||||
Accepted: true, HorizontalPose: angle <= horizontalAngleThreshold,
|
||||
RapidVerticalChange: rapid, HorizontalAngleDegree: angle, HasAngle: true,
|
||||
HipCenterY: hipY, HasHipCenterY: true, TorsoLength: torsoLength, Reason: "accepted",
|
||||
}
|
||||
}
|
||||
|
||||
func midpoint(person pose.PersonPose, first, second int) (float32, float32) {
|
||||
return (person.Keypoints[first].X + person.Keypoints[second].X) / 2,
|
||||
(person.Keypoints[first].Y + person.Keypoints[second].Y) / 2
|
||||
}
|
||||
@@ -0,0 +1,94 @@
|
||||
package fall
|
||||
|
||||
import "fmt"
|
||||
|
||||
type Engine struct {
|
||||
config EngineConfig
|
||||
tracker *tracker
|
||||
policy *policy
|
||||
stateMachine *stateMachine
|
||||
previousEvidence map[string]PoseEvidence
|
||||
activeTrackIDs map[string]bool
|
||||
}
|
||||
|
||||
func NewEngine(config EngineConfig) (*Engine, error) {
|
||||
if config.KeypointConfidenceThreshold < 0 || config.KeypointConfidenceThreshold > 1 {
|
||||
return nil, fmt.Errorf("keypoint confidence threshold must be between 0 and 1")
|
||||
}
|
||||
if config.SuspectWindowSeconds < 0 {
|
||||
return nil, fmt.Errorf("suspect window seconds must be non-negative")
|
||||
}
|
||||
if config.HorizontalAngleThresholdDegrees < 0 || config.HorizontalAngleThresholdDegrees > 90 {
|
||||
return nil, fmt.Errorf("horizontal angle threshold degrees must be between 0 and 90")
|
||||
}
|
||||
tracker, err := newTracker(0.2, 2.0)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
stateMachine, err := newStateMachine(config)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &Engine{
|
||||
config: config, tracker: tracker, policy: newPolicy(config.SuspectWindowSeconds, config.RequireRapidDrop),
|
||||
stateMachine: stateMachine, previousEvidence: make(map[string]PoseEvidence), activeTrackIDs: make(map[string]bool),
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (engine *Engine) StateOf(trackID string) State {
|
||||
return engine.stateMachine.stateOf(trackID)
|
||||
}
|
||||
|
||||
func (engine *Engine) Process(frame Frame) FrameResult {
|
||||
tracked, err := engine.tracker.update(frame.Poses, frame.Timestamp, frame.Width, frame.Height)
|
||||
if err != nil {
|
||||
return FrameResult{}
|
||||
}
|
||||
currentIDs := make(map[string]bool, len(tracked))
|
||||
for _, trackedPose := range tracked {
|
||||
currentIDs[trackedPose.TrackID] = true
|
||||
}
|
||||
events := engine.rejectMissingTracks(currentIDs, frame.Timestamp)
|
||||
people := make([]PersonAnalysis, 0, len(tracked))
|
||||
for _, trackedPose := range tracked {
|
||||
quality := assessPoseQuality(trackedPose.Pose, engine.config.KeypointConfidenceThreshold, engine.config.RequireLowerBody)
|
||||
var previous *PoseEvidence
|
||||
if candidate, found := engine.previousEvidence[trackedPose.TrackID]; found {
|
||||
previous = &candidate
|
||||
}
|
||||
poseEvidence := extractEvidence(trackedPose.Pose, quality, previous, engine.config.HorizontalAngleThresholdDegrees)
|
||||
stateBefore := engine.stateMachine.stateOf(trackedPose.TrackID)
|
||||
evidence := engine.policy.evaluate(trackedPose.TrackID, poseEvidence, frame.Timestamp, stateBefore)
|
||||
newEvents, err := engine.stateMachine.update(trackedPose.TrackID, evidence, frame.Timestamp)
|
||||
if err == nil {
|
||||
events = append(events, newEvents...)
|
||||
}
|
||||
if poseEvidence.Accepted {
|
||||
engine.previousEvidence[trackedPose.TrackID] = poseEvidence
|
||||
} else {
|
||||
delete(engine.previousEvidence, trackedPose.TrackID)
|
||||
}
|
||||
people = append(people, PersonAnalysis{
|
||||
TrackedPose: trackedPose, PoseEvidence: poseEvidence, Evidence: evidence,
|
||||
State: engine.stateMachine.stateOf(trackedPose.TrackID),
|
||||
})
|
||||
}
|
||||
engine.activeTrackIDs = currentIDs
|
||||
return FrameResult{People: people, Events: events}
|
||||
}
|
||||
|
||||
func (engine *Engine) rejectMissingTracks(currentIDs map[string]bool, now float64) []Event {
|
||||
events := make([]Event, 0)
|
||||
for trackID := range engine.activeTrackIDs {
|
||||
if currentIDs[trackID] {
|
||||
continue
|
||||
}
|
||||
delete(engine.previousEvidence, trackID)
|
||||
engine.policy.evaluate(trackID, PoseEvidence{}, now, engine.stateMachine.stateOf(trackID))
|
||||
newEvents, err := engine.stateMachine.update(trackID, Evidence{}, now)
|
||||
if err == nil {
|
||||
events = append(events, newEvents...)
|
||||
}
|
||||
}
|
||||
return events
|
||||
}
|
||||
@@ -0,0 +1,105 @@
|
||||
package fall
|
||||
|
||||
import (
|
||||
"math"
|
||||
"testing"
|
||||
|
||||
"silverpose/v2/internal/pose"
|
||||
)
|
||||
|
||||
func TestEngineConfirmsPersistentHorizontalPoseAtV1Delay(t *testing.T) {
|
||||
engine, err := NewEngine(EngineConfig{
|
||||
KeypointConfidenceThreshold: 0.4,
|
||||
SuspectWindowSeconds: 0.5,
|
||||
ConfirmWindowSeconds: 1.8,
|
||||
RecoveryWindowSeconds: 2.0,
|
||||
CooldownSeconds: 10.0,
|
||||
RequireRapidDrop: false,
|
||||
RequireLowerBody: false,
|
||||
HorizontalAngleThresholdDegrees: 45.0,
|
||||
ConfigVersion: "cfg-v1-contract",
|
||||
SessionID: "regression",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("NewEngine returned an error: %v", err)
|
||||
}
|
||||
|
||||
engine.Process(Frame{Timestamp: 0.0, Width: 180, Height: 180, Poses: []pose.PersonPose{testPose(false)}})
|
||||
engine.Process(Frame{Timestamp: 0.1, Width: 180, Height: 180, Poses: []pose.PersonPose{testPose(true)}})
|
||||
result := engine.Process(Frame{Timestamp: 1.91, Width: 180, Height: 180, Poses: []pose.PersonPose{testPose(true)}})
|
||||
|
||||
if len(result.Events) != 1 {
|
||||
t.Fatalf("events = %+v, want one confirmation", result.Events)
|
||||
}
|
||||
event := result.Events[0]
|
||||
if event.TrackID != "P-0001" || math.Abs(event.LatencySeconds-1.81) > 1e-9 || event.ConfigVersion != "cfg-v1-contract" {
|
||||
t.Fatalf("event = %+v", event)
|
||||
}
|
||||
if got := engine.StateOf("P-0001"); got != Confirmed {
|
||||
t.Fatalf("state = %s, want CONFIRMED", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEngineDoesNotConfirmWhenTrackIsMissing(t *testing.T) {
|
||||
engine, err := NewEngine(testConfig())
|
||||
if err != nil {
|
||||
t.Fatalf("NewEngine returned an error: %v", err)
|
||||
}
|
||||
|
||||
engine.Process(Frame{Timestamp: 0.0, Width: 180, Height: 180, Poses: []pose.PersonPose{testPose(false)}})
|
||||
engine.Process(Frame{Timestamp: 0.1, Width: 180, Height: 180, Poses: []pose.PersonPose{testPose(true)}})
|
||||
engine.Process(Frame{Timestamp: 0.2, Width: 180, Height: 180})
|
||||
result := engine.Process(Frame{Timestamp: 1.9, Width: 180, Height: 180, Poses: []pose.PersonPose{testPose(true)}})
|
||||
|
||||
if len(result.Events) != 0 {
|
||||
t.Fatalf("events = %+v, want none after a missing track", result.Events)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEngineDoesNotAlarmForUprightSequence(t *testing.T) {
|
||||
engine, err := NewEngine(testConfig())
|
||||
if err != nil {
|
||||
t.Fatalf("NewEngine returned an error: %v", err)
|
||||
}
|
||||
|
||||
for _, timestamp := range []float64{0.0, 0.6, 1.2, 1.8, 2.4} {
|
||||
result := engine.Process(Frame{Timestamp: timestamp, Width: 180, Height: 180, Poses: []pose.PersonPose{testPose(false)}})
|
||||
if len(result.Events) != 0 {
|
||||
t.Fatalf("timestamp %v emitted %+v for upright pose", timestamp, result.Events)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func testConfig() EngineConfig {
|
||||
return EngineConfig{
|
||||
KeypointConfidenceThreshold: 0.4,
|
||||
SuspectWindowSeconds: 0.5,
|
||||
ConfirmWindowSeconds: 1.8,
|
||||
RecoveryWindowSeconds: 2.0,
|
||||
CooldownSeconds: 10.0,
|
||||
RequireRapidDrop: false,
|
||||
RequireLowerBody: false,
|
||||
HorizontalAngleThresholdDegrees: 45.0,
|
||||
ConfigVersion: "cfg-v1-contract",
|
||||
SessionID: "regression",
|
||||
}
|
||||
}
|
||||
|
||||
func testPose(horizontal bool) pose.PersonPose {
|
||||
person := pose.PersonPose{Box: pose.Box{Left: 20, Top: 20, Right: 160, Bottom: 160}, Confidence: 0.9}
|
||||
for index := range person.Keypoints {
|
||||
person.Keypoints[index] = pose.Keypoint{X: float32(index), Y: float32(index), Confidence: 0.9}
|
||||
}
|
||||
if horizontal {
|
||||
person.Keypoints[5] = pose.Keypoint{X: 20, Y: 80, Confidence: 0.9}
|
||||
person.Keypoints[6] = pose.Keypoint{X: 30, Y: 80, Confidence: 0.9}
|
||||
person.Keypoints[11] = pose.Keypoint{X: 70, Y: 100, Confidence: 0.9}
|
||||
person.Keypoints[12] = pose.Keypoint{X: 80, Y: 100, Confidence: 0.9}
|
||||
} else {
|
||||
person.Keypoints[5] = pose.Keypoint{X: 30, Y: 10, Confidence: 0.9}
|
||||
person.Keypoints[6] = pose.Keypoint{X: 40, Y: 10, Confidence: 0.9}
|
||||
person.Keypoints[11] = pose.Keypoint{X: 30, Y: 30, Confidence: 0.9}
|
||||
person.Keypoints[12] = pose.Keypoint{X: 40, Y: 30, Confidence: 0.9}
|
||||
}
|
||||
return person
|
||||
}
|
||||
@@ -0,0 +1,37 @@
|
||||
package fall
|
||||
|
||||
type policy struct {
|
||||
suspectWindowSeconds float64
|
||||
requireRapidDrop bool
|
||||
rapidDropAt map[string]float64
|
||||
}
|
||||
|
||||
func newPolicy(suspectWindowSeconds float64, requireRapidDrop bool) *policy {
|
||||
return &policy{
|
||||
suspectWindowSeconds: suspectWindowSeconds,
|
||||
requireRapidDrop: requireRapidDrop,
|
||||
rapidDropAt: make(map[string]float64),
|
||||
}
|
||||
}
|
||||
|
||||
func (policy *policy) evaluate(trackID string, poseEvidence PoseEvidence, now float64, state State) Evidence {
|
||||
if !poseEvidence.Accepted {
|
||||
delete(policy.rapidDropAt, trackID)
|
||||
return Evidence{}
|
||||
}
|
||||
if poseEvidence.RapidVerticalChange {
|
||||
policy.rapidDropAt[trackID] = now
|
||||
}
|
||||
candidate := false
|
||||
if state == Suspect {
|
||||
candidate = poseEvidence.HorizontalPose
|
||||
} else if poseEvidence.HorizontalPose {
|
||||
if !policy.requireRapidDrop {
|
||||
candidate = true
|
||||
} else if droppedAt, found := policy.rapidDropAt[trackID]; found {
|
||||
candidate = now-droppedAt <= policy.suspectWindowSeconds
|
||||
}
|
||||
}
|
||||
recovery := (state == Confirmed || state == Recovering) && !poseEvidence.HorizontalPose && !poseEvidence.RapidVerticalChange
|
||||
return Evidence{Accepted: true, IsFallCandidate: candidate, IsRecoveryCandidate: recovery}
|
||||
}
|
||||
@@ -0,0 +1,142 @@
|
||||
package fall
|
||||
|
||||
import "fmt"
|
||||
|
||||
type stateRecord struct {
|
||||
state State
|
||||
|
||||
suspectStartedAt float64
|
||||
hasSuspectStartedAt bool
|
||||
confirmedAt float64
|
||||
hasConfirmedAt bool
|
||||
recoveryStartedAt float64
|
||||
hasRecoveryStartedAt bool
|
||||
lastUpdatedAt float64
|
||||
hasLastUpdatedAt bool
|
||||
}
|
||||
|
||||
type stateMachine struct {
|
||||
confirmWindowSeconds float64
|
||||
recoveryWindowSeconds float64
|
||||
cooldownSeconds float64
|
||||
configVersion string
|
||||
sessionID string
|
||||
records map[string]*stateRecord
|
||||
nextEventNumber int
|
||||
}
|
||||
|
||||
func newStateMachine(config EngineConfig) (*stateMachine, error) {
|
||||
if config.ConfirmWindowSeconds < 1 || config.ConfirmWindowSeconds > 3 {
|
||||
return nil, fmt.Errorf("confirm window seconds must be between 1 and 3")
|
||||
}
|
||||
if config.RecoveryWindowSeconds <= 0 {
|
||||
return nil, fmt.Errorf("recovery window seconds must be positive")
|
||||
}
|
||||
if config.CooldownSeconds < 0 {
|
||||
return nil, fmt.Errorf("cooldown seconds must be non-negative")
|
||||
}
|
||||
if config.ConfigVersion == "" {
|
||||
return nil, fmt.Errorf("config version must be non-empty")
|
||||
}
|
||||
return &stateMachine{
|
||||
confirmWindowSeconds: config.ConfirmWindowSeconds, recoveryWindowSeconds: config.RecoveryWindowSeconds,
|
||||
cooldownSeconds: config.CooldownSeconds, configVersion: config.ConfigVersion, sessionID: config.SessionID,
|
||||
records: make(map[string]*stateRecord), nextEventNumber: 1,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (machine *stateMachine) stateOf(trackID string) State {
|
||||
if record, found := machine.records[trackID]; found {
|
||||
return record.state
|
||||
}
|
||||
return Normal
|
||||
}
|
||||
|
||||
func (machine *stateMachine) update(trackID string, evidence Evidence, now float64) ([]Event, error) {
|
||||
if trackID == "" {
|
||||
return nil, fmt.Errorf("track ID must be non-empty")
|
||||
}
|
||||
record, found := machine.records[trackID]
|
||||
if !found {
|
||||
record = &stateRecord{state: Normal}
|
||||
machine.records[trackID] = record
|
||||
}
|
||||
if record.hasLastUpdatedAt && now < record.lastUpdatedAt {
|
||||
return nil, fmt.Errorf("timestamps must be monotonic per track")
|
||||
}
|
||||
record.lastUpdatedAt, record.hasLastUpdatedAt = now, true
|
||||
|
||||
if !evidence.Accepted {
|
||||
machine.rejectEvidence(record)
|
||||
return nil, nil
|
||||
}
|
||||
switch record.state {
|
||||
case Normal:
|
||||
if evidence.IsFallCandidate {
|
||||
record.state, record.suspectStartedAt, record.hasSuspectStartedAt = Suspect, now, true
|
||||
}
|
||||
return nil, nil
|
||||
case Suspect:
|
||||
if !evidence.IsFallCandidate {
|
||||
machine.setNormal(record)
|
||||
return nil, nil
|
||||
}
|
||||
if now-record.suspectStartedAt >= machine.confirmWindowSeconds {
|
||||
record.state, record.confirmedAt, record.hasConfirmedAt = Confirmed, now, true
|
||||
event := machine.newEvent(trackID, record.suspectStartedAt, now)
|
||||
record.hasSuspectStartedAt = false
|
||||
return []Event{event}, nil
|
||||
}
|
||||
return nil, nil
|
||||
case Confirmed:
|
||||
if evidence.IsRecoveryCandidate && now-record.confirmedAt >= machine.cooldownSeconds {
|
||||
record.state, record.recoveryStartedAt, record.hasRecoveryStartedAt = Recovering, now, true
|
||||
}
|
||||
return nil, nil
|
||||
case Recovering:
|
||||
if evidence.IsFallCandidate {
|
||||
record.state, record.suspectStartedAt, record.hasSuspectStartedAt = Suspect, now, true
|
||||
record.hasRecoveryStartedAt = false
|
||||
return nil, nil
|
||||
}
|
||||
if !evidence.IsRecoveryCandidate {
|
||||
record.state, record.hasRecoveryStartedAt = Confirmed, false
|
||||
return nil, nil
|
||||
}
|
||||
if now-record.recoveryStartedAt >= machine.recoveryWindowSeconds {
|
||||
machine.setNormal(record)
|
||||
}
|
||||
return nil, nil
|
||||
default:
|
||||
return nil, fmt.Errorf("unknown fall state")
|
||||
}
|
||||
}
|
||||
|
||||
func (machine *stateMachine) rejectEvidence(record *stateRecord) {
|
||||
if record.state == Suspect {
|
||||
machine.setNormal(record)
|
||||
} else if record.state == Recovering {
|
||||
record.state, record.hasRecoveryStartedAt = Confirmed, false
|
||||
}
|
||||
}
|
||||
|
||||
func (machine *stateMachine) setNormal(record *stateRecord) {
|
||||
record.state = Normal
|
||||
record.hasSuspectStartedAt = false
|
||||
record.hasConfirmedAt = false
|
||||
record.hasRecoveryStartedAt = false
|
||||
}
|
||||
|
||||
func (machine *stateMachine) newEvent(trackID string, suspectedAt, confirmedAt float64) Event {
|
||||
prefix := "FALL-"
|
||||
if machine.sessionID != "" {
|
||||
prefix += machine.sessionID + "-"
|
||||
}
|
||||
event := Event{
|
||||
EventID: fmt.Sprintf("%s%06d", prefix, machine.nextEventNumber), TrackID: trackID,
|
||||
ConfigVersion: machine.configVersion, SuspectedAtMonotonic: suspectedAt,
|
||||
ConfirmedAtMonotonic: confirmedAt, LatencySeconds: confirmedAt - suspectedAt, State: Confirmed,
|
||||
}
|
||||
machine.nextEventNumber++
|
||||
return event
|
||||
}
|
||||
@@ -0,0 +1,93 @@
|
||||
package fall
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"math"
|
||||
"sort"
|
||||
|
||||
"silverpose/v2/internal/pose"
|
||||
)
|
||||
|
||||
type track struct {
|
||||
centerX float32
|
||||
centerY float32
|
||||
lastSeenAt float64
|
||||
}
|
||||
|
||||
type tracker struct {
|
||||
maxMatchDistanceRatio float64
|
||||
maxAgeSeconds float64
|
||||
tracks map[string]track
|
||||
nextTrackNumber int
|
||||
}
|
||||
|
||||
func newTracker(maxMatchDistanceRatio, maxAgeSeconds float64) (*tracker, error) {
|
||||
if maxMatchDistanceRatio <= 0 || maxMatchDistanceRatio > 1 {
|
||||
return nil, fmt.Errorf("max match distance ratio must be in (0, 1]")
|
||||
}
|
||||
if maxAgeSeconds <= 0 {
|
||||
return nil, fmt.Errorf("max age seconds must be positive")
|
||||
}
|
||||
return &tracker{
|
||||
maxMatchDistanceRatio: maxMatchDistanceRatio,
|
||||
maxAgeSeconds: maxAgeSeconds,
|
||||
tracks: make(map[string]track),
|
||||
nextTrackNumber: 1,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (tracker *tracker) update(poses []pose.PersonPose, now float64, width, height int) ([]TrackedPose, error) {
|
||||
if width <= 0 || height <= 0 {
|
||||
return nil, fmt.Errorf("frame size must contain positive width and height")
|
||||
}
|
||||
tracker.expire(now)
|
||||
availableIDs := make(map[string]bool, len(tracker.tracks))
|
||||
for trackID := range tracker.tracks {
|
||||
availableIDs[trackID] = true
|
||||
}
|
||||
tracked := make([]TrackedPose, 0, len(poses))
|
||||
for _, person := range poses {
|
||||
centerX := (person.Box.Left + person.Box.Right) / 2
|
||||
centerY := (person.Box.Top + person.Box.Bottom) / 2
|
||||
trackID := tracker.nearestAvailable(centerX, centerY, availableIDs, width, height)
|
||||
if trackID == "" {
|
||||
trackID = fmt.Sprintf("P-%04d", tracker.nextTrackNumber)
|
||||
tracker.nextTrackNumber++
|
||||
} else {
|
||||
delete(availableIDs, trackID)
|
||||
}
|
||||
tracker.tracks[trackID] = track{centerX: centerX, centerY: centerY, lastSeenAt: now}
|
||||
tracked = append(tracked, TrackedPose{TrackID: trackID, Pose: person})
|
||||
}
|
||||
return tracked, nil
|
||||
}
|
||||
|
||||
func (tracker *tracker) nearestAvailable(centerX, centerY float32, availableIDs map[string]bool, width, height int) string {
|
||||
ids := make([]string, 0, len(availableIDs))
|
||||
for trackID := range availableIDs {
|
||||
ids = append(ids, trackID)
|
||||
}
|
||||
sort.Strings(ids)
|
||||
closestID := ""
|
||||
closestDistance := math.Inf(1)
|
||||
for _, trackID := range ids {
|
||||
previous := tracker.tracks[trackID]
|
||||
distance := math.Hypot(
|
||||
float64(centerX-previous.centerX)/float64(width),
|
||||
float64(centerY-previous.centerY)/float64(height),
|
||||
)
|
||||
if distance <= tracker.maxMatchDistanceRatio && distance < closestDistance {
|
||||
closestID = trackID
|
||||
closestDistance = distance
|
||||
}
|
||||
}
|
||||
return closestID
|
||||
}
|
||||
|
||||
func (tracker *tracker) expire(now float64) {
|
||||
for trackID, current := range tracker.tracks {
|
||||
if now-current.lastSeenAt > tracker.maxAgeSeconds {
|
||||
delete(tracker.tracks, trackID)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,83 @@
|
||||
package fall
|
||||
|
||||
import "silverpose/v2/internal/pose"
|
||||
|
||||
type State string
|
||||
|
||||
const (
|
||||
Normal State = "NORMAL"
|
||||
Suspect State = "SUSPECT"
|
||||
Confirmed State = "CONFIRMED"
|
||||
Recovering State = "RECOVERING"
|
||||
)
|
||||
|
||||
type Evidence struct {
|
||||
Accepted bool
|
||||
IsFallCandidate bool
|
||||
IsRecoveryCandidate bool
|
||||
}
|
||||
|
||||
type PoseQuality struct {
|
||||
Accepted bool
|
||||
Reason string
|
||||
VisibleJointCount int
|
||||
}
|
||||
|
||||
type PoseEvidence struct {
|
||||
Accepted bool
|
||||
HorizontalPose bool
|
||||
RapidVerticalChange bool
|
||||
HorizontalAngleDegree float32
|
||||
HasAngle bool
|
||||
HipCenterY float32
|
||||
HasHipCenterY bool
|
||||
TorsoLength float32
|
||||
Reason string
|
||||
}
|
||||
|
||||
type Event struct {
|
||||
EventID string
|
||||
TrackID string
|
||||
ConfigVersion string
|
||||
SuspectedAtMonotonic float64
|
||||
ConfirmedAtMonotonic float64
|
||||
LatencySeconds float64
|
||||
State State
|
||||
}
|
||||
|
||||
type Frame struct {
|
||||
Timestamp float64
|
||||
Width int
|
||||
Height int
|
||||
Poses []pose.PersonPose
|
||||
}
|
||||
|
||||
type TrackedPose struct {
|
||||
TrackID string
|
||||
Pose pose.PersonPose
|
||||
}
|
||||
|
||||
type PersonAnalysis struct {
|
||||
TrackedPose TrackedPose
|
||||
PoseEvidence PoseEvidence
|
||||
Evidence Evidence
|
||||
State State
|
||||
}
|
||||
|
||||
type FrameResult struct {
|
||||
People []PersonAnalysis
|
||||
Events []Event
|
||||
}
|
||||
|
||||
type EngineConfig struct {
|
||||
KeypointConfidenceThreshold float32
|
||||
SuspectWindowSeconds float64
|
||||
ConfirmWindowSeconds float64
|
||||
RecoveryWindowSeconds float64
|
||||
CooldownSeconds float64
|
||||
RequireRapidDrop bool
|
||||
RequireLowerBody bool
|
||||
HorizontalAngleThresholdDegrees float32
|
||||
ConfigVersion string
|
||||
SessionID string
|
||||
}
|
||||
@@ -0,0 +1,125 @@
|
||||
package pose
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"sort"
|
||||
)
|
||||
|
||||
// ParseYOLOv8Pose turns the fixed [1,56,8400] YOLO Pose output into source
|
||||
// pixel coordinates, applying confidence filtering and class-agnostic NMS.
|
||||
func ParseYOLOv8Pose(output []float32, transform LetterboxTransform, confidenceThreshold, iouThreshold float32) ([]PersonPose, error) {
|
||||
if len(output) != YOLOPoseOutputValues {
|
||||
return nil, fmt.Errorf("YOLO Pose output length = %d, want %d", len(output), YOLOPoseOutputValues)
|
||||
}
|
||||
if transform.Scale <= 0 {
|
||||
return nil, fmt.Errorf("letterbox scale must be positive")
|
||||
}
|
||||
if confidenceThreshold < 0 || confidenceThreshold > 1 || iouThreshold < 0 || iouThreshold > 1 {
|
||||
return nil, fmt.Errorf("confidence and IoU thresholds must be between 0 and 1")
|
||||
}
|
||||
|
||||
candidates := make([]PersonPose, 0)
|
||||
for candidate := 0; candidate < YOLOPoseCandidateCount; candidate++ {
|
||||
confidence := valueAt(output, 4, candidate)
|
||||
if confidence < confidenceThreshold {
|
||||
continue
|
||||
}
|
||||
centerX := valueAt(output, 0, candidate)
|
||||
centerY := valueAt(output, 1, candidate)
|
||||
width := valueAt(output, 2, candidate)
|
||||
height := valueAt(output, 3, candidate)
|
||||
person := PersonPose{
|
||||
Box: Box{
|
||||
Left: transform.restoreX(centerX - width/2),
|
||||
Top: transform.restoreY(centerY - height/2),
|
||||
Right: transform.restoreX(centerX + width/2),
|
||||
Bottom: transform.restoreY(centerY + height/2),
|
||||
},
|
||||
Confidence: confidence,
|
||||
}
|
||||
for keypoint := 0; keypoint < YOLOPoseKeypointCount; keypoint++ {
|
||||
channel := 5 + keypoint*3
|
||||
person.Keypoints[keypoint] = Keypoint{
|
||||
X: transform.restoreX(valueAt(output, channel, candidate)),
|
||||
Y: transform.restoreY(valueAt(output, channel+1, candidate)),
|
||||
Confidence: valueAt(output, channel+2, candidate),
|
||||
}
|
||||
}
|
||||
candidates = append(candidates, person)
|
||||
}
|
||||
|
||||
sort.SliceStable(candidates, func(left, right int) bool {
|
||||
return candidates[left].Confidence > candidates[right].Confidence
|
||||
})
|
||||
selected := make([]PersonPose, 0, len(candidates))
|
||||
for _, candidate := range candidates {
|
||||
overlaps := false
|
||||
for _, accepted := range selected {
|
||||
if intersectionOverUnion(candidate.Box, accepted.Box) > iouThreshold {
|
||||
overlaps = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !overlaps {
|
||||
selected = append(selected, candidate)
|
||||
}
|
||||
}
|
||||
return selected, nil
|
||||
}
|
||||
|
||||
func valueAt(output []float32, channel, candidate int) float32 {
|
||||
return output[channel*YOLOPoseCandidateCount+candidate]
|
||||
}
|
||||
|
||||
func (transform LetterboxTransform) restoreX(value float32) float32 {
|
||||
return clampCoordinate((value-float32(transform.PadLeft))/transform.Scale, transform.OriginalWidth)
|
||||
}
|
||||
|
||||
func (transform LetterboxTransform) restoreY(value float32) float32 {
|
||||
return clampCoordinate((value-float32(transform.PadTop))/transform.Scale, transform.OriginalHeight)
|
||||
}
|
||||
|
||||
func clampCoordinate(value float32, size int) float32 {
|
||||
if size <= 0 {
|
||||
return value
|
||||
}
|
||||
if value < 0 {
|
||||
return 0
|
||||
}
|
||||
maximum := float32(size)
|
||||
if value > maximum {
|
||||
return maximum
|
||||
}
|
||||
return value
|
||||
}
|
||||
|
||||
func intersectionOverUnion(left, right Box) float32 {
|
||||
intersectionLeft := maxFloat(left.Left, right.Left)
|
||||
intersectionTop := maxFloat(left.Top, right.Top)
|
||||
intersectionRight := minFloat(left.Right, right.Right)
|
||||
intersectionBottom := minFloat(left.Bottom, right.Bottom)
|
||||
intersectionWidth := maxFloat(0, intersectionRight-intersectionLeft)
|
||||
intersectionHeight := maxFloat(0, intersectionBottom-intersectionTop)
|
||||
intersection := intersectionWidth * intersectionHeight
|
||||
leftArea := maxFloat(0, left.Right-left.Left) * maxFloat(0, left.Bottom-left.Top)
|
||||
rightArea := maxFloat(0, right.Right-right.Left) * maxFloat(0, right.Bottom-right.Top)
|
||||
union := leftArea + rightArea - intersection
|
||||
if union <= 0 {
|
||||
return 0
|
||||
}
|
||||
return intersection / union
|
||||
}
|
||||
|
||||
func minFloat(left, right float32) float32 {
|
||||
if left < right {
|
||||
return left
|
||||
}
|
||||
return right
|
||||
}
|
||||
|
||||
func maxFloat(left, right float32) float32 {
|
||||
if left > right {
|
||||
return left
|
||||
}
|
||||
return right
|
||||
}
|
||||
@@ -0,0 +1,45 @@
|
||||
package pose
|
||||
|
||||
import "testing"
|
||||
|
||||
func TestParseYOLOv8PoseRestoresCoordinatesAndSuppressesOverlap(t *testing.T) {
|
||||
output := make([]float32, YOLOPoseOutputValues)
|
||||
putDetection(output, 0, 320, 320, 100, 200, 0.90, 330, 340)
|
||||
putDetection(output, 1, 322, 321, 100, 200, 0.80, 332, 341)
|
||||
|
||||
people, err := ParseYOLOv8Pose(output, LetterboxTransform{Scale: 1}, 0.25, 0.70)
|
||||
if err != nil {
|
||||
t.Fatalf("ParseYOLOv8Pose returned an error: %v", err)
|
||||
}
|
||||
if len(people) != 1 {
|
||||
t.Fatalf("people = %d, want one NMS survivor", len(people))
|
||||
}
|
||||
person := people[0]
|
||||
if person.Box != (Box{Left: 270, Top: 220, Right: 370, Bottom: 420}) {
|
||||
t.Fatalf("box = %+v", person.Box)
|
||||
}
|
||||
if point := person.Keypoints[0]; point.X != 330 || point.Y != 340 || point.Confidence != 0.9 {
|
||||
t.Fatalf("first keypoint = %+v", point)
|
||||
}
|
||||
}
|
||||
|
||||
func TestParseYOLOv8PoseRejectsUnexpectedOutputLength(t *testing.T) {
|
||||
_, err := ParseYOLOv8Pose([]float32{0}, LetterboxTransform{Scale: 1}, 0.25, 0.70)
|
||||
if err == nil {
|
||||
t.Fatal("ParseYOLOv8Pose accepted a malformed output")
|
||||
}
|
||||
}
|
||||
|
||||
func putDetection(output []float32, candidate int, centerX, centerY, width, height, confidence, keypointX, keypointY float32) {
|
||||
output[0*YOLOPoseCandidateCount+candidate] = centerX
|
||||
output[1*YOLOPoseCandidateCount+candidate] = centerY
|
||||
output[2*YOLOPoseCandidateCount+candidate] = width
|
||||
output[3*YOLOPoseCandidateCount+candidate] = height
|
||||
output[4*YOLOPoseCandidateCount+candidate] = confidence
|
||||
for keypoint := 0; keypoint < YOLOPoseKeypointCount; keypoint++ {
|
||||
base := 5 + keypoint*3
|
||||
output[(base+0)*YOLOPoseCandidateCount+candidate] = keypointX
|
||||
output[(base+1)*YOLOPoseCandidateCount+candidate] = keypointY
|
||||
output[(base+2)*YOLOPoseCandidateCount+candidate] = 0.9
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,87 @@
|
||||
package pose
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"math"
|
||||
)
|
||||
|
||||
const letterboxPadding = byte(114)
|
||||
|
||||
// PreprocessBGR reproduces the fixed-shape Ultralytics letterbox contract:
|
||||
// BGR source pixels are resized with bilinear interpolation, padded in 114
|
||||
// gray, converted to RGB/CHW and normalised to [0, 1].
|
||||
func PreprocessBGR(frame []byte, width, height, target int) ([]float32, LetterboxTransform, error) {
|
||||
if width <= 0 || height <= 0 || target <= 0 {
|
||||
return nil, LetterboxTransform{}, fmt.Errorf("frame width, height and target must be positive")
|
||||
}
|
||||
if len(frame) != width*height*3 {
|
||||
return nil, LetterboxTransform{}, fmt.Errorf("BGR frame length = %d, want %d", len(frame), width*height*3)
|
||||
}
|
||||
|
||||
scale := math.Min(float64(target)/float64(width), float64(target)/float64(height))
|
||||
resizedWidth := int(math.Round(float64(width) * scale))
|
||||
resizedHeight := int(math.Round(float64(height) * scale))
|
||||
padWidth := target - resizedWidth
|
||||
padHeight := target - resizedHeight
|
||||
padLeft := int(math.Round(float64(padWidth)/2.0 - 0.1))
|
||||
padTop := int(math.Round(float64(padHeight)/2.0 - 0.1))
|
||||
transform := LetterboxTransform{
|
||||
OriginalWidth: width, OriginalHeight: height, InputSize: target,
|
||||
ResizedWidth: resizedWidth, ResizedHeight: resizedHeight,
|
||||
PadLeft: padLeft, PadTop: padTop, Scale: float32(scale),
|
||||
}
|
||||
|
||||
plane := target * target
|
||||
input := make([]float32, 3*plane)
|
||||
padding := float32(letterboxPadding) / 255.0
|
||||
for index := range input {
|
||||
input[index] = padding
|
||||
}
|
||||
for y := 0; y < resizedHeight; y++ {
|
||||
sourceY := clampFloat((float64(y)+0.5)*float64(height)/float64(resizedHeight)-0.5, 0, float64(height-1))
|
||||
y0 := int(math.Floor(sourceY))
|
||||
y1 := minInt(y0+1, height-1)
|
||||
yWeight := float32(sourceY - float64(y0))
|
||||
for x := 0; x < resizedWidth; x++ {
|
||||
sourceX := clampFloat((float64(x)+0.5)*float64(width)/float64(resizedWidth)-0.5, 0, float64(width-1))
|
||||
x0 := int(math.Floor(sourceX))
|
||||
x1 := minInt(x0+1, width-1)
|
||||
xWeight := float32(sourceX - float64(x0))
|
||||
leftTop := (y0*width + x0) * 3
|
||||
rightTop := (y0*width + x1) * 3
|
||||
leftBottom := (y1*width + x0) * 3
|
||||
rightBottom := (y1*width + x1) * 3
|
||||
destination := (padTop+y)*target + padLeft + x
|
||||
for sourceChannel := 0; sourceChannel < 3; sourceChannel++ {
|
||||
value := bilinear(
|
||||
frame[leftTop+sourceChannel], frame[rightTop+sourceChannel],
|
||||
frame[leftBottom+sourceChannel], frame[rightBottom+sourceChannel],
|
||||
xWeight, yWeight,
|
||||
)
|
||||
// BGR source maps to RGB tensor planes.
|
||||
tensorChannel := 2 - sourceChannel
|
||||
// cv2.resize writes uint8 pixels before Ultralytics converts the
|
||||
// image to float; round here to retain that observable contract.
|
||||
input[tensorChannel*plane+destination] = float32(math.Round(float64(value))) / 255.0
|
||||
}
|
||||
}
|
||||
}
|
||||
return input, transform, nil
|
||||
}
|
||||
|
||||
func bilinear(topLeft, topRight, bottomLeft, bottomRight byte, xWeight, yWeight float32) float32 {
|
||||
top := float32(topLeft)*(1-xWeight) + float32(topRight)*xWeight
|
||||
bottom := float32(bottomLeft)*(1-xWeight) + float32(bottomRight)*xWeight
|
||||
return top*(1-yWeight) + bottom*yWeight
|
||||
}
|
||||
|
||||
func clampFloat(value, minimum, maximum float64) float64 {
|
||||
return math.Max(minimum, math.Min(maximum, value))
|
||||
}
|
||||
|
||||
func minInt(left, right int) int {
|
||||
if left < right {
|
||||
return left
|
||||
}
|
||||
return right
|
||||
}
|
||||
@@ -0,0 +1,39 @@
|
||||
package pose
|
||||
|
||||
import "testing"
|
||||
|
||||
func TestPreprocessBGRUsesUltralyticsPaddingAndRGBCHW(t *testing.T) {
|
||||
// A 2x1 BGR frame letterboxes into a 4x4 tensor with one row of 114-gray
|
||||
// padding above and below. The rightmost source pixel is pure blue in BGR.
|
||||
frame := []byte{0, 0, 0, 255, 0, 0}
|
||||
|
||||
input, transform, err := PreprocessBGR(frame, 2, 1, 4)
|
||||
if err != nil {
|
||||
t.Fatalf("PreprocessBGR returned an error: %v", err)
|
||||
}
|
||||
if transform.ResizedWidth != 4 || transform.ResizedHeight != 2 || transform.PadTop != 1 {
|
||||
t.Fatalf("unexpected transform: %+v", transform)
|
||||
}
|
||||
|
||||
const plane = 16
|
||||
padding := float32(114.0 / 255.0)
|
||||
if got := input[2*plane+0]; got != padding {
|
||||
t.Fatalf("top padding blue = %v, want %v", got, padding)
|
||||
}
|
||||
if got := input[0*plane+1*4+3]; got != 0 {
|
||||
t.Fatalf("red channel = %v, want 0", got)
|
||||
}
|
||||
if got := input[1*plane+1*4+3]; got != 0 {
|
||||
t.Fatalf("green channel = %v, want 0", got)
|
||||
}
|
||||
if got := input[2*plane+1*4+3]; got != 1 {
|
||||
t.Fatalf("blue channel = %v, want 1", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPreprocessBGRRejectsTruncatedFrame(t *testing.T) {
|
||||
_, _, err := PreprocessBGR([]byte{0}, 2, 1, 4)
|
||||
if err == nil {
|
||||
t.Fatal("PreprocessBGR accepted a truncated BGR frame")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,39 @@
|
||||
package pose
|
||||
|
||||
const (
|
||||
YOLOPoseKeypointCount = 17
|
||||
YOLOPoseCandidateCount = 8400
|
||||
YOLOPoseChannelCount = 56
|
||||
YOLOPoseOutputValues = YOLOPoseChannelCount * YOLOPoseCandidateCount
|
||||
)
|
||||
|
||||
type Keypoint struct {
|
||||
X float32
|
||||
Y float32
|
||||
Confidence float32
|
||||
}
|
||||
|
||||
type Box struct {
|
||||
Left float32
|
||||
Top float32
|
||||
Right float32
|
||||
Bottom float32
|
||||
}
|
||||
|
||||
type PersonPose struct {
|
||||
Box Box
|
||||
Confidence float32
|
||||
Keypoints [YOLOPoseKeypointCount]Keypoint
|
||||
}
|
||||
|
||||
// LetterboxTransform maps fixed-square YOLO coordinates back to source pixels.
|
||||
type LetterboxTransform struct {
|
||||
OriginalWidth int
|
||||
OriginalHeight int
|
||||
InputSize int
|
||||
ResizedWidth int
|
||||
ResizedHeight int
|
||||
PadLeft int
|
||||
PadTop int
|
||||
Scale float32
|
||||
}
|
||||
+88
-39
@@ -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)
|
||||
}
|
||||
|
||||
@@ -14,3 +14,13 @@ func TestRunPoseRejectsWrongInputLengthBeforeLoadingRuntime(t *testing.T) {
|
||||
t.Fatalf("RunPose error = %q, want input length error", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenRuntimeReportsMissingDLL(t *testing.T) {
|
||||
_, err := OpenRuntime("missing.onnx", "missing.dll")
|
||||
if err == nil {
|
||||
t.Fatal("OpenRuntime accepted a missing ONNX Runtime DLL")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "initialize ONNX Runtime") {
|
||||
t.Fatalf("OpenRuntime error = %q", err)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,43 +1,14 @@
|
||||
package spike
|
||||
|
||||
import "fmt"
|
||||
import "silverpose/v2/internal/pose"
|
||||
|
||||
// LetterboxBGRToNCHW converts one packed BGR frame into the RGB, CHW,
|
||||
// float32 tensor expected by the locked 640-pixel YOLO pose ONNX model.
|
||||
//
|
||||
// It intentionally uses nearest-neighbour scaling for this Spike. T-303
|
||||
// must replace or validate this preprocessing against the V1 implementation
|
||||
// before it becomes the V2 production preprocessing path.
|
||||
// T-303 aligned the implementation with the fixed-shape Ultralytics
|
||||
// letterbox used by the V1/ONNX baseline. New V2 code should import the pose
|
||||
// package directly to retain the returned coordinate transform.
|
||||
func LetterboxBGRToNCHW(frame []byte, width, height, target int) ([]float32, error) {
|
||||
if width <= 0 || height <= 0 || target <= 0 {
|
||||
return nil, fmt.Errorf("frame width, height and target must be positive")
|
||||
}
|
||||
if len(frame) != width*height*3 {
|
||||
return nil, fmt.Errorf("BGR frame length = %d, want %d", len(frame), width*height*3)
|
||||
}
|
||||
|
||||
scaleWidth := target
|
||||
scaleHeight := height * target / width
|
||||
if scaleHeight > target {
|
||||
scaleHeight = target
|
||||
scaleWidth = width * target / height
|
||||
}
|
||||
padX := (target - scaleWidth) / 2
|
||||
padY := (target - scaleHeight) / 2
|
||||
plane := target * target
|
||||
input := make([]float32, 3*plane)
|
||||
|
||||
for y := 0; y < scaleHeight; y++ {
|
||||
sourceY := y * height / scaleHeight
|
||||
for x := 0; x < scaleWidth; x++ {
|
||||
sourceX := x * width / scaleWidth
|
||||
source := (sourceY*width + sourceX) * 3
|
||||
destination := (padY+y)*target + padX + x
|
||||
input[destination] = float32(frame[source+2]) / 255.0
|
||||
input[plane+destination] = float32(frame[source+1]) / 255.0
|
||||
input[2*plane+destination] = float32(frame[source]) / 255.0
|
||||
}
|
||||
}
|
||||
|
||||
return input, nil
|
||||
input, _, err := pose.PreprocessBGR(frame, width, height, target)
|
||||
return input, err
|
||||
}
|
||||
|
||||
@@ -16,8 +16,9 @@ func TestLetterboxBGRToNCHWPlacesPixelInScaledImage(t *testing.T) {
|
||||
if got := len(input); got != 3*plane {
|
||||
t.Fatalf("input length = %d, want %d", got, 3*plane)
|
||||
}
|
||||
if got := input[2*plane+159*640+320]; got != 0 {
|
||||
t.Fatalf("top padding blue channel = %v, want 0", got)
|
||||
padding := float32(114.0 / 255.0)
|
||||
if got := input[2*plane+159*640+320]; got != padding {
|
||||
t.Fatalf("top padding blue channel = %v, want %v", got, padding)
|
||||
}
|
||||
if got := input[0*plane+320*640+320]; got != 0 {
|
||||
t.Fatalf("red channel = %v, want 0", got)
|
||||
@@ -25,8 +26,8 @@ func TestLetterboxBGRToNCHWPlacesPixelInScaledImage(t *testing.T) {
|
||||
if got := input[1*plane+320*640+320]; got != 0 {
|
||||
t.Fatalf("green channel = %v, want 0", got)
|
||||
}
|
||||
if got := input[2*plane+320*640+320]; got != 1 {
|
||||
t.Fatalf("blue channel = %v, want 1", got)
|
||||
if got := input[2*plane+320*640+320]; got != float32(128.0/255.0) {
|
||||
t.Fatalf("blue channel = %v, want cv2 INTER_LINEAR value 128/255", got)
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user