134 lines
3.9 KiB
Go
134 lines
3.9 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 != "unknown" {
|
|
t.Fatalf("sex position = %q, want unknown", effective.SexPosition)
|
|
}
|
|
if len(effective.Boxes) != 0 {
|
|
t.Fatalf("negative annotation has %d boxes, want 0", len(effective.Boxes))
|
|
}
|
|
}
|