91 lines
2.4 KiB
Go
91 lines
2.4 KiB
Go
package main
|
|
|
|
import (
|
|
"os"
|
|
"path/filepath"
|
|
"testing"
|
|
)
|
|
|
|
func TestTrainingVideoMAELabelForAnnotationUsesNegativeAsKeine(t *testing.T) {
|
|
got := trainingVideoMAELabelForAnnotation(TrainingAnnotation{
|
|
Negative: true,
|
|
Correction: &TrainingCorrection{
|
|
SexPosition: "doggy",
|
|
},
|
|
})
|
|
|
|
if got != trainingNoSexPositionLabel {
|
|
t.Fatalf("label = %q, want %q", got, trainingNoSexPositionLabel)
|
|
}
|
|
}
|
|
|
|
func TestTrainingVideoMAELabelForAnnotationUsesCorrection(t *testing.T) {
|
|
got := trainingVideoMAELabelForAnnotation(TrainingAnnotation{
|
|
Correction: &TrainingCorrection{
|
|
SexPosition: "cowgirl",
|
|
},
|
|
})
|
|
|
|
if got != "cowgirl" {
|
|
t.Fatalf("label = %q, want cowgirl", got)
|
|
}
|
|
}
|
|
|
|
func TestBuildAnalyzeVideoMAEClipsUsesWindowAndStride(t *testing.T) {
|
|
samples := []videoFrameSample{
|
|
{Time: 0, Path: "f0.jpg"},
|
|
{Time: 1, Path: "f1.jpg"},
|
|
{Time: 2, Path: "f2.jpg"},
|
|
{Time: 3, Path: "f3.jpg"},
|
|
{Time: 4, Path: "f4.jpg"},
|
|
{Time: 5, Path: "f5.jpg"},
|
|
}
|
|
|
|
clips := buildAnalyzeVideoMAEClips(samples, 6)
|
|
if len(clips) != 3 {
|
|
t.Fatalf("clip count = %d, want 3", len(clips))
|
|
}
|
|
if clips[0].Time != 0 || clips[1].Time != 2 || clips[2].Time != 4 {
|
|
t.Fatalf("clip times = %.1f, %.1f, %.1f; want 0, 2, 4", clips[0].Time, clips[1].Time, clips[2].Time)
|
|
}
|
|
if len(clips[1].Paths) < 4 {
|
|
t.Fatalf("middle clip paths = %d, want at least 4", len(clips[1].Paths))
|
|
}
|
|
}
|
|
|
|
func TestLimitAnalyzeVideoMAEClipsKeepsEvenCoverage(t *testing.T) {
|
|
clips := make([]analyzeVideoMAEClipReqItem, 10)
|
|
for i := range clips {
|
|
clips[i] = analyzeVideoMAEClipReqItem{Time: float64(i)}
|
|
}
|
|
|
|
limited := limitAnalyzeVideoMAEClips(clips, 4)
|
|
if len(limited) != 4 {
|
|
t.Fatalf("clip count = %d, want 4", len(limited))
|
|
}
|
|
|
|
got := []float64{limited[0].Time, limited[1].Time, limited[2].Time, limited[3].Time}
|
|
want := []float64{0, 3, 6, 9}
|
|
for i := range want {
|
|
if got[i] != want[i] {
|
|
t.Fatalf("clip times = %v, want %v", got, want)
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestTrainingResolveVideoMAEModelUsesConfigDir(t *testing.T) {
|
|
root := t.TempDir()
|
|
modelDir := filepath.Join(root, "videomae", "model")
|
|
if err := os.MkdirAll(modelDir, 0755); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if err := os.WriteFile(filepath.Join(modelDir, "config.json"), []byte("{}"), 0644); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
got := trainingResolveVideoMAEModel(root)
|
|
if !got.EffectiveExists || got.EffectivePath != modelDir || got.Source != "videomae_clip" {
|
|
t.Fatalf("unexpected model resolution: %+v", got)
|
|
}
|
|
}
|