feat(v2): add fall event regression engine
This commit is contained in:
@@ -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
|
||||
}
|
||||
Reference in New Issue
Block a user