46 lines
1.7 KiB
Go
46 lines
1.7 KiB
Go
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
|
|
}
|
|
}
|