94 lines
2.6 KiB
Go
94 lines
2.6 KiB
Go
|
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)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|