480 lines
15 KiB
Go
480 lines
15 KiB
Go
package main
|
|
|
|
import (
|
|
"fmt"
|
|
"os"
|
|
"path/filepath"
|
|
"strings"
|
|
"testing"
|
|
)
|
|
|
|
func TestTrainingDetectorLabelContentAllowsExplicitNegative(t *testing.T) {
|
|
content, err := trainingDetectorLabelContent(nil, map[string]int{"person": 0}, true)
|
|
if err != nil {
|
|
t.Fatalf("negative label content returned error: %v", err)
|
|
}
|
|
if len(content) != 0 {
|
|
t.Fatalf("negative label content = %q, want empty", string(content))
|
|
}
|
|
}
|
|
|
|
func TestTrainingDetectorLabelContentRejectsAccidentalEmptySample(t *testing.T) {
|
|
_, err := trainingDetectorLabelContent(
|
|
[]TrainingBox{{Label: "unknown_label", X: 0, Y: 0, W: 1, H: 1}},
|
|
map[string]int{"person": 0},
|
|
false,
|
|
)
|
|
if err == nil {
|
|
t.Fatal("invalid non-negative sample should be rejected")
|
|
}
|
|
}
|
|
|
|
func TestTrainingDetectorLabelContentWritesYOLOBox(t *testing.T) {
|
|
content, err := trainingDetectorLabelContent(
|
|
[]TrainingBox{{Label: "person", X: 0.1, Y: 0.2, W: 0.4, H: 0.6}},
|
|
map[string]int{"person": 3},
|
|
false,
|
|
)
|
|
if err != nil {
|
|
t.Fatalf("positive label content returned error: %v", err)
|
|
}
|
|
|
|
got := strings.TrimSpace(string(content))
|
|
want := "3 0.300000 0.500000 0.400000 0.600000"
|
|
if got != want {
|
|
t.Fatalf("label content = %q, want %q", got, want)
|
|
}
|
|
}
|
|
|
|
func TestTrainingCountDetectorSamplesIncludesEmptyLabels(t *testing.T) {
|
|
root := t.TempDir()
|
|
imagesDir := filepath.Join(root, "images")
|
|
labelsDir := filepath.Join(root, "labels")
|
|
|
|
if err := os.MkdirAll(imagesDir, 0755); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if err := os.MkdirAll(labelsDir, 0755); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if err := os.WriteFile(filepath.Join(imagesDir, "negative.jpg"), []byte("image"), 0644); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if err := os.WriteFile(filepath.Join(labelsDir, "negative.txt"), []byte{}, 0644); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
if got := trainingCountDetectorSamples(imagesDir, labelsDir); got != 1 {
|
|
t.Fatalf("sample count = %d, want 1", got)
|
|
}
|
|
if got := trainingCountPositiveDetectorSamples(imagesDir, labelsDir); got != 0 {
|
|
t.Fatalf("positive sample count = %d, want 0", got)
|
|
}
|
|
}
|
|
|
|
func TestTrainingEnsureDetectorValidationSampleIncludesPositiveExample(t *testing.T) {
|
|
root := t.TempDir()
|
|
trainImages := filepath.Join(root, "detector", "dataset", "images", "train")
|
|
trainLabels := filepath.Join(root, "detector", "dataset", "labels", "train")
|
|
|
|
if err := os.MkdirAll(trainImages, 0755); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if err := os.MkdirAll(trainLabels, 0755); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
for i := 0; i < minDetectorTrainCount; i++ {
|
|
id := fmt.Sprintf("sample-%02d", i)
|
|
if err := os.WriteFile(filepath.Join(trainImages, id+".jpg"), []byte("image"), 0644); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
label := []byte{}
|
|
if i == minDetectorTrainCount-1 {
|
|
label = []byte("0 0.5 0.5 1 1\n")
|
|
}
|
|
if err := os.WriteFile(filepath.Join(trainLabels, id+".txt"), label, 0644); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
}
|
|
|
|
if err := trainingEnsureDetectorValidationSample(root); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
valImages := filepath.Join(root, "detector", "dataset", "images", "val")
|
|
valLabels := filepath.Join(root, "detector", "dataset", "labels", "val")
|
|
if got := trainingCountDetectorSamples(valImages, valLabels); got < minDetectorValCount {
|
|
t.Fatalf("validation sample count = %d, want at least %d", got, minDetectorValCount)
|
|
}
|
|
if got := trainingCountPositiveDetectorSamples(valImages, valLabels); got < 1 {
|
|
t.Fatalf("positive validation sample count = %d, want at least 1", got)
|
|
}
|
|
}
|
|
|
|
func TestTrainingEffectiveCorrectionClearsNegativeAnnotation(t *testing.T) {
|
|
effective := trainingEffectiveCorrection(TrainingAnnotation{
|
|
Negative: true,
|
|
Prediction: TrainingPrediction{
|
|
SexPosition: "doggy",
|
|
Boxes: []TrainingBox{
|
|
{Label: "person", X: 0, Y: 0, W: 1, H: 1},
|
|
},
|
|
},
|
|
})
|
|
|
|
if effective.SexPosition != trainingNoSexPositionLabel {
|
|
t.Fatalf("sex position = %q, want %s", effective.SexPosition, trainingNoSexPositionLabel)
|
|
}
|
|
if len(effective.Boxes) != 0 {
|
|
t.Fatalf("negative annotation has %d boxes, want 0", len(effective.Boxes))
|
|
}
|
|
}
|
|
|
|
func TestNoSexPositionAliasesTreatUnknownAsNoPosition(t *testing.T) {
|
|
for _, value := range []string{"Unknown", "unknown", "unbekannt", "none", "no_position"} {
|
|
if !isNoSexPositionLabel(value) {
|
|
t.Fatalf("%q should be treated as no sex position", value)
|
|
}
|
|
if got := normalizeSexPositionLabel(value); got != trainingNoSexPositionLabel {
|
|
t.Fatalf("normalizeSexPositionLabel(%q) = %q, want %q", value, got, trainingNoSexPositionLabel)
|
|
}
|
|
}
|
|
}
|
|
|
|
func trainingTestPoseKeypoints(overrides map[string]TrainingKeypoint) []TrainingKeypoint {
|
|
names := []string{
|
|
"nose",
|
|
"left_eye", "right_eye",
|
|
"left_ear", "right_ear",
|
|
"left_shoulder", "right_shoulder",
|
|
"left_elbow", "right_elbow",
|
|
"left_wrist", "right_wrist",
|
|
"left_hip", "right_hip",
|
|
"left_knee", "right_knee",
|
|
"left_ankle", "right_ankle",
|
|
}
|
|
|
|
out := make([]TrainingKeypoint, 0, len(names))
|
|
for i, name := range names {
|
|
point := TrainingKeypoint{
|
|
Name: name,
|
|
X: 0.20 + float64(i)*0.01,
|
|
Y: 0.25 + float64(i)*0.01,
|
|
Conf: 0.85,
|
|
}
|
|
|
|
if override, ok := overrides[name]; ok {
|
|
override.Name = name
|
|
point = override
|
|
}
|
|
|
|
out = append(out, point)
|
|
}
|
|
|
|
return out
|
|
}
|
|
|
|
func TestTrainingApplyPoseToPredictionAggregatesMultiplePersons(t *testing.T) {
|
|
pred := TrainingPrediction{
|
|
SexPosition: trainingNoSexPositionLabel,
|
|
Boxes: []TrainingBox{
|
|
{Label: "person_female", X: 0.10, Y: 0.10, W: 0.35, H: 0.75},
|
|
{Label: "person_male", X: 0.42, Y: 0.12, W: 0.34, H: 0.72},
|
|
},
|
|
}
|
|
|
|
pose := TrainingPosePrediction{
|
|
Available: true,
|
|
Persons: []TrainingPosePerson{
|
|
{
|
|
Label: "missionary",
|
|
Score: 0.58,
|
|
Box: TrainingBox{X: 0.10, Y: 0.10, W: 0.35, H: 0.75},
|
|
Keypoints: trainingTestPoseKeypoints(nil),
|
|
},
|
|
{
|
|
Label: "doggy",
|
|
Score: 0.55,
|
|
Box: TrainingBox{X: 0.42, Y: 0.12, W: 0.34, H: 0.72},
|
|
Keypoints: trainingTestPoseKeypoints(nil),
|
|
},
|
|
{
|
|
Label: "doggy",
|
|
Score: 0.54,
|
|
Box: TrainingBox{X: 0.52, Y: 0.14, W: 0.30, H: 0.70},
|
|
Keypoints: trainingTestPoseKeypoints(nil),
|
|
},
|
|
},
|
|
}
|
|
|
|
got := trainingApplyPoseToPrediction(pred, pose)
|
|
if got.SexPosition != "doggy" {
|
|
t.Fatalf("sex position = %q, want doggy", got.SexPosition)
|
|
}
|
|
}
|
|
|
|
func TestTrainingApplyPoseToPredictionUsesObjectPositionContext(t *testing.T) {
|
|
pred := TrainingPrediction{
|
|
SexPosition: trainingNoSexPositionLabel,
|
|
Boxes: []TrainingBox{
|
|
{Label: "person_female", X: 0.08, Y: 0.08, W: 0.42, H: 0.78},
|
|
{Label: "person_male", X: 0.42, Y: 0.12, W: 0.35, H: 0.76},
|
|
{Label: "penis", X: 0.50, Y: 0.48, W: 0.07, H: 0.08},
|
|
},
|
|
}
|
|
|
|
pose := TrainingPosePrediction{
|
|
Available: true,
|
|
Persons: []TrainingPosePerson{
|
|
{
|
|
Label: "missionary",
|
|
Score: 0.58,
|
|
Box: TrainingBox{X: 0.08, Y: 0.08, W: 0.42, H: 0.78},
|
|
Keypoints: trainingTestPoseKeypoints(map[string]TrainingKeypoint{
|
|
"nose": {X: 0.18, Y: 0.18, Conf: 0.90},
|
|
}),
|
|
},
|
|
{
|
|
Label: "blowjob",
|
|
Score: 0.55,
|
|
Box: TrainingBox{X: 0.42, Y: 0.12, W: 0.35, H: 0.76},
|
|
Keypoints: trainingTestPoseKeypoints(map[string]TrainingKeypoint{
|
|
"nose": {X: 0.53, Y: 0.52, Conf: 0.95},
|
|
}),
|
|
},
|
|
},
|
|
}
|
|
|
|
got := trainingApplyPoseToPrediction(pred, pose)
|
|
if got.SexPosition != "blowjob" {
|
|
t.Fatalf("sex position = %q, want blowjob", got.SexPosition)
|
|
}
|
|
}
|
|
|
|
func TestTrainingApplyPoseToPredictionUsesOcclusionTolerantCowgirlGeometry(t *testing.T) {
|
|
pred := TrainingPrediction{
|
|
SexPosition: trainingNoSexPositionLabel,
|
|
Boxes: []TrainingBox{
|
|
{Label: "person_female", X: 0.30, Y: 0.10, W: 0.28, H: 0.62},
|
|
{Label: "person_male", X: 0.20, Y: 0.42, W: 0.55, H: 0.24},
|
|
},
|
|
}
|
|
|
|
pose := TrainingPosePrediction{
|
|
Available: true,
|
|
Persons: []TrainingPosePerson{
|
|
{
|
|
Label: "person",
|
|
Score: 0.72,
|
|
Box: TrainingBox{X: 0.30, Y: 0.10, W: 0.28, H: 0.62},
|
|
Keypoints: trainingTestPoseKeypoints(map[string]TrainingKeypoint{
|
|
"left_shoulder": {X: 0.40, Y: 0.22, Conf: 0.92},
|
|
"right_shoulder": {X: 0.48, Y: 0.22, Conf: 0.92},
|
|
"left_hip": {X: 0.38, Y: 0.38, Conf: 0.92},
|
|
"right_hip": {X: 0.50, Y: 0.38, Conf: 0.92},
|
|
"left_knee": {X: 0.31, Y: 0.58, Conf: 0.90},
|
|
"right_knee": {X: 0.57, Y: 0.58, Conf: 0.90},
|
|
}),
|
|
},
|
|
{
|
|
Label: "person",
|
|
Score: 0.70,
|
|
Box: TrainingBox{X: 0.20, Y: 0.42, W: 0.55, H: 0.24},
|
|
Keypoints: trainingTestPoseKeypoints(map[string]TrainingKeypoint{
|
|
"left_shoulder": {X: 0.25, Y: 0.50, Conf: 0.90},
|
|
"right_shoulder": {X: 0.38, Y: 0.50, Conf: 0.90},
|
|
"left_hip": {X: 0.52, Y: 0.53, Conf: 0.90},
|
|
"right_hip": {X: 0.66, Y: 0.53, Conf: 0.90},
|
|
"left_knee": {X: 0.58, Y: 0.56, Conf: 0.88},
|
|
"right_knee": {X: 0.70, Y: 0.56, Conf: 0.88},
|
|
}),
|
|
},
|
|
},
|
|
}
|
|
|
|
got := trainingApplyPoseToPrediction(pred, pose)
|
|
if got.SexPosition != "cowgirl" {
|
|
t.Fatalf("sex position = %q, want cowgirl", got.SexPosition)
|
|
}
|
|
if got.SexPositionScore < 0.30 {
|
|
t.Fatalf("sex position score = %.3f, want >= 0.30", got.SexPositionScore)
|
|
}
|
|
}
|
|
|
|
func TestTrainingApplyPoseToPredictionKeepsUnreliablePoseOutOfContext(t *testing.T) {
|
|
pred := TrainingPrediction{
|
|
SexPosition: trainingNoSexPositionLabel,
|
|
Boxes: []TrainingBox{
|
|
{Label: "person_female", X: 0.08, Y: 0.08, W: 0.42, H: 0.78},
|
|
{Label: "penis", X: 0.50, Y: 0.48, W: 0.07, H: 0.08},
|
|
},
|
|
}
|
|
|
|
pose := TrainingPosePrediction{
|
|
Available: true,
|
|
Persons: []TrainingPosePerson{
|
|
{
|
|
Label: "blowjob",
|
|
Score: 0.24,
|
|
Box: TrainingBox{X: 0.08, Y: 0.08, W: 0.42, H: 0.78},
|
|
Keypoints: trainingTestPoseKeypoints(map[string]TrainingKeypoint{
|
|
"nose": {X: 0.53, Y: 0.52, Conf: 0.95},
|
|
}),
|
|
},
|
|
},
|
|
}
|
|
|
|
got := trainingApplyPoseToPrediction(pred, pose)
|
|
if got.SexPosition != trainingNoSexPositionLabel {
|
|
t.Fatalf("sex position = %q, want %s", got.SexPosition, trainingNoSexPositionLabel)
|
|
}
|
|
if len(got.Persons) != 1 {
|
|
t.Fatalf("persons = %d, want 1", len(got.Persons))
|
|
}
|
|
if got.Persons[0].Reliable {
|
|
t.Fatalf("low-score pose should be marked unreliable: %+v", got.Persons[0])
|
|
}
|
|
if got.Persons[0].VisibleKeypoints == 0 || got.Persons[0].Quality == 0 {
|
|
t.Fatalf("pose quality should be annotated: %+v", got.Persons[0])
|
|
}
|
|
}
|
|
|
|
func TestTrainingApplyPoseToPredictionUsesBoxContextWithoutPose(t *testing.T) {
|
|
pred := TrainingPrediction{
|
|
ModelAvailable: true,
|
|
Source: "yolo26_detector",
|
|
SexPosition: trainingNoSexPositionLabel,
|
|
Boxes: []TrainingBox{
|
|
{Label: "person_female", X: 0.18, Y: 0.12, W: 0.42, H: 0.72},
|
|
{Label: "person_male", X: 0.42, Y: 0.18, W: 0.38, H: 0.68},
|
|
{Label: "penis", X: 0.48, Y: 0.50, W: 0.08, H: 0.08},
|
|
{Label: "pussy", X: 0.54, Y: 0.51, W: 0.08, H: 0.08},
|
|
},
|
|
}
|
|
|
|
got := trainingApplyPoseToPrediction(pred, TrainingPosePrediction{Available: false})
|
|
if got.SexPosition == trainingNoSexPositionLabel {
|
|
t.Fatalf("sex position = %q, want uncertain context prediction", got.SexPosition)
|
|
}
|
|
if got.SexPositionScore < trainingPositionContextMinScore {
|
|
t.Fatalf("sex position score = %.3f, want >= %.3f", got.SexPositionScore, trainingPositionContextMinScore)
|
|
}
|
|
if got.SexPositionScore > trainingPositionContextMaxScore {
|
|
t.Fatalf("sex position score = %.3f, want <= %.3f", got.SexPositionScore, trainingPositionContextMaxScore)
|
|
}
|
|
if !strings.Contains(got.Source, "box_context") {
|
|
t.Fatalf("source = %q, want box_context", got.Source)
|
|
}
|
|
}
|
|
|
|
func TestTrainingApplyPoseToPredictionBoostsPoseWithBoxContext(t *testing.T) {
|
|
pred := TrainingPrediction{
|
|
ModelAvailable: true,
|
|
Source: "yolo26_detector",
|
|
SexPosition: trainingNoSexPositionLabel,
|
|
Boxes: []TrainingBox{
|
|
{Label: "penis", X: 0.48, Y: 0.50, W: 0.08, H: 0.08},
|
|
{Label: "pussy", X: 0.54, Y: 0.51, W: 0.08, H: 0.08},
|
|
},
|
|
}
|
|
|
|
pose := TrainingPosePrediction{
|
|
Available: true,
|
|
Persons: []TrainingPosePerson{
|
|
{
|
|
Label: "doggy",
|
|
Score: 0.31,
|
|
Box: TrainingBox{X: 0.20, Y: 0.12, W: 0.60, H: 0.76},
|
|
Keypoints: trainingTestPoseKeypoints(nil),
|
|
},
|
|
},
|
|
}
|
|
|
|
got := trainingApplyPoseToPrediction(pred, pose)
|
|
if got.SexPosition != "doggy" {
|
|
t.Fatalf("sex position = %q, want doggy", got.SexPosition)
|
|
}
|
|
if got.SexPositionScore <= 0.31 {
|
|
t.Fatalf("sex position score = %.3f, want context boost over 0.31", got.SexPositionScore)
|
|
}
|
|
if !strings.Contains(got.Source, "yolo_pose") || !strings.Contains(got.Source, "box_context") {
|
|
t.Fatalf("source = %q, want yolo_pose and box_context", got.Source)
|
|
}
|
|
}
|
|
|
|
func TestBuildClipPositionHitsFromEvidenceRequiresTemporalSupport(t *testing.T) {
|
|
hits := buildClipPositionHitsFromEvidence([]analyzePositionEvidence{
|
|
{Time: 1, Label: "doggy", Score: 0.55, HasPose: true, PersonCount: 2},
|
|
{Time: 2, Label: "cowgirl", Score: 0.60, HasPose: true, PersonCount: 2},
|
|
{Time: 6, Label: "doggy", Score: 0.58, HasPose: true, PersonCount: 2},
|
|
}, 10)
|
|
|
|
if len(hits) != 0 {
|
|
t.Fatalf("hits = %+v, want no temporally supported position", hits)
|
|
}
|
|
}
|
|
|
|
func TestBuildClipPositionHitsFromEvidenceBuildsStablePosition(t *testing.T) {
|
|
hits := buildClipPositionHitsFromEvidence([]analyzePositionEvidence{
|
|
{Time: 1, Label: "doggy", Score: 0.45, HasPose: true, HasContext: true, PersonCount: 2},
|
|
{Time: 2, Label: "doggy", Score: 0.47, HasPose: true, HasContext: true, PersonCount: 2},
|
|
{Time: 3, Label: "doggy", Score: 0.44, HasPose: true, HasContext: true, PersonCount: 2},
|
|
{Time: 6, Label: "cowgirl", Score: 0.62, HasPose: true, PersonCount: 2},
|
|
}, 10)
|
|
|
|
if len(hits) == 0 {
|
|
t.Fatal("expected stable doggy position hit")
|
|
}
|
|
if hits[0].Label != "position:doggy" {
|
|
t.Fatalf("label = %q, want position:doggy", hits[0].Label)
|
|
}
|
|
if hits[0].End <= hits[0].Start {
|
|
t.Fatalf("invalid hit span: %+v", hits[0])
|
|
}
|
|
}
|
|
|
|
func TestBuildClipPositionHitsFromEvidenceCapsContextOnlyScore(t *testing.T) {
|
|
hits := buildClipPositionHitsFromEvidence([]analyzePositionEvidence{
|
|
{Time: 1, Label: "missionary", Score: 0.90, HasContext: true, PersonCount: 2},
|
|
{Time: 2, Label: "missionary", Score: 0.88, HasContext: true, PersonCount: 2},
|
|
{Time: 3, Label: "missionary", Score: 0.92, HasContext: true, PersonCount: 2},
|
|
}, 8)
|
|
|
|
if len(hits) == 0 {
|
|
t.Fatal("expected context-only position hit")
|
|
}
|
|
if hits[0].Score > trainingPositionContextMaxScore {
|
|
t.Fatalf("score = %.3f, want <= %.3f", hits[0].Score, trainingPositionContextMaxScore)
|
|
}
|
|
}
|
|
|
|
func TestTrainingFilterPosePersonsByContextDropsUnannotatedPeople(t *testing.T) {
|
|
persons := []TrainingPosePerson{
|
|
{
|
|
Label: "doggy",
|
|
Score: 0.80,
|
|
Box: TrainingBox{X: 0.10, Y: 0.10, W: 0.30, H: 0.70},
|
|
Keypoints: trainingTestPoseKeypoints(nil),
|
|
},
|
|
{
|
|
Label: "doggy",
|
|
Score: 0.70,
|
|
Box: TrainingBox{X: 0.70, Y: 0.10, W: 0.20, H: 0.70},
|
|
Keypoints: trainingTestPoseKeypoints(nil),
|
|
},
|
|
}
|
|
|
|
filtered := trainingFilterPosePersonsByContext(persons, []TrainingBox{
|
|
{Label: "person_female", X: 0.08, Y: 0.08, W: 0.34, H: 0.74},
|
|
})
|
|
|
|
if len(filtered) != 1 {
|
|
t.Fatalf("filtered persons = %d, want 1", len(filtered))
|
|
}
|
|
if filtered[0].Box.X != persons[0].Box.X {
|
|
t.Fatalf("kept wrong person: %+v", filtered[0].Box)
|
|
}
|
|
}
|