607 lines
14 KiB
Go
607 lines
14 KiB
Go
package main
|
|
|
|
import (
|
|
"bufio"
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"math"
|
|
"os"
|
|
"os/exec"
|
|
"path/filepath"
|
|
"sort"
|
|
"strconv"
|
|
"strings"
|
|
)
|
|
|
|
const (
|
|
trainingVideoMAEClipSeconds = 4.0
|
|
trainingVideoMAENumFrames = 16
|
|
trainingVideoMAEFrameSize = 224
|
|
minVideoMAETrainCount = 5
|
|
minVideoMAEValCount = 1
|
|
)
|
|
|
|
func trainingVideoMAEEpochs() int {
|
|
raw := strings.TrimSpace(os.Getenv("TRAINING_VIDEOMAE_EPOCHS"))
|
|
if raw == "" {
|
|
return 8
|
|
}
|
|
|
|
n, err := strconv.Atoi(raw)
|
|
if err != nil || n < 1 {
|
|
return 8
|
|
}
|
|
if n > 200 {
|
|
return 200
|
|
}
|
|
return n
|
|
}
|
|
|
|
func trainingVideoMAEBatchSize() int {
|
|
raw := strings.TrimSpace(os.Getenv("TRAINING_VIDEOMAE_BATCH"))
|
|
if raw == "" {
|
|
return 2
|
|
}
|
|
|
|
n, err := strconv.Atoi(raw)
|
|
if err != nil || n < 1 {
|
|
return 2
|
|
}
|
|
if n > 32 {
|
|
return 32
|
|
}
|
|
return n
|
|
}
|
|
|
|
type trainingVideoMAEManifestEntry struct {
|
|
SampleID string `json:"sampleId"`
|
|
Split string `json:"split"`
|
|
Label string `json:"label"`
|
|
ClipDir string `json:"clipDir"`
|
|
SourcePath string `json:"sourcePath,omitempty"`
|
|
Second float64 `json:"second,omitempty"`
|
|
FrameCount int `json:"frameCount"`
|
|
}
|
|
|
|
func trainingResolveVideoMAEModel(root string) trainingModelResolution {
|
|
modelDir := filepath.Join(root, "videomae", "model")
|
|
configPath := filepath.Join(modelDir, "config.json")
|
|
if fileExistsNonEmpty(configPath) {
|
|
return trainingModelResolution{
|
|
BestPath: modelDir,
|
|
EffectivePath: modelDir,
|
|
Source: "videomae_clip",
|
|
TrainedExists: true,
|
|
EffectiveExists: true,
|
|
}
|
|
}
|
|
|
|
return trainingModelResolution{
|
|
BestPath: modelDir,
|
|
EffectivePath: modelDir,
|
|
Source: "videomae_missing",
|
|
TrainedExists: false,
|
|
EffectiveExists: false,
|
|
}
|
|
}
|
|
|
|
func trainingEnsureVideoMAEDirs(root string) error {
|
|
dirs := []string{
|
|
filepath.Join(root, "videomae"),
|
|
filepath.Join(root, "videomae", "dataset"),
|
|
filepath.Join(root, "videomae", "dataset", "clips"),
|
|
filepath.Join(root, "videomae", "dataset", "clips", "train"),
|
|
filepath.Join(root, "videomae", "dataset", "clips", "val"),
|
|
filepath.Join(root, "videomae", "model"),
|
|
filepath.Join(root, "videomae", "runs"),
|
|
}
|
|
|
|
for _, dir := range dirs {
|
|
if err := os.MkdirAll(dir, 0755); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
|
|
return trainingWriteVideoMAELabelsFile(root)
|
|
}
|
|
|
|
func trainingVideoMAELabels() ([]string, error) {
|
|
grouped, err := trainingGroupedLabels()
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
out := []string{trainingNoSexPositionLabel}
|
|
seen := map[string]bool{
|
|
trainingNoSexPositionLabel: true,
|
|
}
|
|
|
|
for _, value := range grouped.SexPositions {
|
|
label := normalizeSexPositionLabel(value)
|
|
if seen[label] {
|
|
continue
|
|
}
|
|
|
|
seen[label] = true
|
|
out = append(out, label)
|
|
}
|
|
|
|
return out, nil
|
|
}
|
|
|
|
func trainingVideoMAELabelSet() (map[string]bool, error) {
|
|
labels, err := trainingVideoMAELabels()
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
out := map[string]bool{}
|
|
for _, label := range labels {
|
|
out[label] = true
|
|
}
|
|
return out, nil
|
|
}
|
|
|
|
func trainingWriteVideoMAELabelsFile(root string) error {
|
|
labels, err := trainingVideoMAELabels()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
body, err := json.MarshalIndent(map[string]any{
|
|
"labels": labels,
|
|
}, "", " ")
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
return os.WriteFile(filepath.Join(root, "videomae", "dataset", "labels.json"), body, 0644)
|
|
}
|
|
|
|
func trainingVideoMAEManifestPath(root string) string {
|
|
return filepath.Join(root, "videomae", "dataset", "manifest.jsonl")
|
|
}
|
|
|
|
func trainingVideoMAELabelForAnnotation(item TrainingAnnotation) string {
|
|
if item.Negative {
|
|
return trainingNoSexPositionLabel
|
|
}
|
|
|
|
effective := trainingEffectiveCorrection(item)
|
|
label := normalizeSexPositionLabel(effective.SexPosition)
|
|
if isNoSexPositionLabel(label) {
|
|
return trainingNoSexPositionLabel
|
|
}
|
|
|
|
return label
|
|
}
|
|
|
|
func trainingVideoMAEFrameFallbackPath(root string, sampleID string) string {
|
|
return filepath.Join(root, "frames", sampleID+".jpg")
|
|
}
|
|
|
|
func trainingVideoMAEAnnotationHasSource(root string, item TrainingAnnotation) bool {
|
|
sourcePath := strings.TrimSpace(item.SourcePath)
|
|
if sourcePath != "" && trainingSupportedImportVideo(sourcePath) && fileExistsNonEmpty(sourcePath) {
|
|
return true
|
|
}
|
|
|
|
return fileExistsNonEmpty(trainingVideoMAEFrameFallbackPath(root, item.SampleID))
|
|
}
|
|
|
|
func trainingCountVideoMAEEligibleAnnotations(root string) (int, error) {
|
|
items, err := trainingReadAnnotations(root)
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
|
|
labelSet, err := trainingVideoMAELabelSet()
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
|
|
count := 0
|
|
for _, item := range items {
|
|
sampleID := strings.TrimSpace(item.SampleID)
|
|
if sampleID == "" || strings.Contains(sampleID, "/") || strings.Contains(sampleID, "\\") {
|
|
continue
|
|
}
|
|
|
|
if !labelSet[trainingVideoMAELabelForAnnotation(item)] {
|
|
continue
|
|
}
|
|
|
|
if trainingVideoMAEAnnotationHasSource(root, item) {
|
|
count++
|
|
}
|
|
}
|
|
|
|
return count, nil
|
|
}
|
|
|
|
func trainingReadVideoMAEManifest(root string) ([]trainingVideoMAEManifestEntry, error) {
|
|
path := trainingVideoMAEManifestPath(root)
|
|
f, err := os.Open(path)
|
|
if err != nil {
|
|
if os.IsNotExist(err) {
|
|
return []trainingVideoMAEManifestEntry{}, nil
|
|
}
|
|
return nil, err
|
|
}
|
|
defer f.Close()
|
|
|
|
out := []trainingVideoMAEManifestEntry{}
|
|
scanner := bufio.NewScanner(f)
|
|
for scanner.Scan() {
|
|
line := strings.TrimSpace(scanner.Text())
|
|
if line == "" {
|
|
continue
|
|
}
|
|
|
|
var entry trainingVideoMAEManifestEntry
|
|
if err := json.Unmarshal([]byte(line), &entry); err != nil {
|
|
continue
|
|
}
|
|
|
|
out = append(out, entry)
|
|
}
|
|
|
|
if err := scanner.Err(); err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
return out, nil
|
|
}
|
|
|
|
func trainingWriteVideoMAEManifest(root string, entries []trainingVideoMAEManifestEntry) error {
|
|
path := trainingVideoMAEManifestPath(root)
|
|
if err := os.MkdirAll(filepath.Dir(path), 0755); err != nil {
|
|
return err
|
|
}
|
|
|
|
var b strings.Builder
|
|
for _, entry := range entries {
|
|
line, err := json.Marshal(entry)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
b.Write(line)
|
|
b.WriteByte('\n')
|
|
}
|
|
|
|
return os.WriteFile(path, []byte(b.String()), 0644)
|
|
}
|
|
|
|
func trainingCountVideoMAEManifestSamples(root string) (train int, val int) {
|
|
entries, err := trainingReadVideoMAEManifest(root)
|
|
if err != nil {
|
|
return 0, 0
|
|
}
|
|
|
|
for _, entry := range entries {
|
|
if entry.FrameCount <= 0 || !fileExistsNonEmpty(filepath.Join(entry.ClipDir, "frame_001.jpg")) {
|
|
continue
|
|
}
|
|
|
|
switch strings.ToLower(strings.TrimSpace(entry.Split)) {
|
|
case "train":
|
|
train++
|
|
case "val":
|
|
val++
|
|
}
|
|
}
|
|
|
|
return train, val
|
|
}
|
|
|
|
func trainingRemoveVideoMAEGeneratedClips(root string) error {
|
|
clipsDir := filepath.Join(root, "videomae", "dataset", "clips")
|
|
if err := os.RemoveAll(clipsDir); err != nil {
|
|
return err
|
|
}
|
|
|
|
for _, split := range []string{"train", "val"} {
|
|
if err := os.MkdirAll(filepath.Join(clipsDir, split), 0755); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
func trainingVideoMAEClipFrameCount(clipDir string) int {
|
|
entries, err := os.ReadDir(clipDir)
|
|
if err != nil {
|
|
return 0
|
|
}
|
|
|
|
count := 0
|
|
for _, entry := range entries {
|
|
if entry.IsDir() {
|
|
continue
|
|
}
|
|
|
|
ext := strings.ToLower(filepath.Ext(entry.Name()))
|
|
if ext == ".jpg" || ext == ".jpeg" || ext == ".png" || ext == ".webp" {
|
|
count++
|
|
}
|
|
}
|
|
|
|
return count
|
|
}
|
|
|
|
func trainingExtractVideoMAEClipFrames(
|
|
ctx context.Context,
|
|
videoPath string,
|
|
centerSecond float64,
|
|
clipDir string,
|
|
) error {
|
|
if strings.TrimSpace(videoPath) == "" {
|
|
return errors.New("video path missing")
|
|
}
|
|
|
|
if err := os.MkdirAll(clipDir, 0755); err != nil {
|
|
return err
|
|
}
|
|
|
|
ffmpegPath := strings.TrimSpace(getSettings().FFmpegPath)
|
|
if ffmpegPath == "" {
|
|
ffmpegPath = "ffmpeg"
|
|
}
|
|
|
|
start := math.Max(0, centerSecond-trainingVideoMAEClipSeconds/2)
|
|
fps := float64(trainingVideoMAENumFrames) / trainingVideoMAEClipSeconds
|
|
vf := fmt.Sprintf(
|
|
"fps=%.6f,scale=%d:%d:force_original_aspect_ratio=decrease,pad=%d:%d:(ow-iw)/2:(oh-ih)/2",
|
|
fps,
|
|
trainingVideoMAEFrameSize,
|
|
trainingVideoMAEFrameSize,
|
|
trainingVideoMAEFrameSize,
|
|
trainingVideoMAEFrameSize,
|
|
)
|
|
|
|
cmd := exec.CommandContext(
|
|
ctx,
|
|
ffmpegPath,
|
|
"-hide_banner",
|
|
"-loglevel", "error",
|
|
"-ss", fmt.Sprintf("%.3f", start),
|
|
"-i", videoPath,
|
|
"-t", fmt.Sprintf("%.3f", trainingVideoMAEClipSeconds),
|
|
"-vf", vf,
|
|
"-frames:v", strconv.Itoa(trainingVideoMAENumFrames),
|
|
"-q:v", "3",
|
|
filepath.Join(clipDir, "frame_%03d.jpg"),
|
|
)
|
|
|
|
hideCommandWindow(cmd)
|
|
|
|
out, err := cmd.CombinedOutput()
|
|
if err != nil {
|
|
return fmt.Errorf("ffmpeg clip extract failed: %w: %s", err, strings.TrimSpace(string(out)))
|
|
}
|
|
|
|
if trainingVideoMAEClipFrameCount(clipDir) <= 0 {
|
|
return errors.New("ffmpeg erzeugte keine Clip-Frames")
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
func trainingWriteStillVideoMAEClip(framePath string, clipDir string) error {
|
|
if !fileExistsNonEmpty(framePath) {
|
|
return errors.New("fallback frame missing")
|
|
}
|
|
|
|
if err := os.MkdirAll(clipDir, 0755); err != nil {
|
|
return err
|
|
}
|
|
|
|
for i := 1; i <= trainingVideoMAENumFrames; i++ {
|
|
dst := filepath.Join(clipDir, fmt.Sprintf("frame_%03d.jpg", i))
|
|
if err := copyFile(framePath, dst); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
func trainingWriteVideoMAEClipForAnnotation(
|
|
ctx context.Context,
|
|
root string,
|
|
item TrainingAnnotation,
|
|
clipDir string,
|
|
) (int, error) {
|
|
if err := os.RemoveAll(clipDir); err != nil {
|
|
return 0, err
|
|
}
|
|
if err := os.MkdirAll(clipDir, 0755); err != nil {
|
|
return 0, err
|
|
}
|
|
|
|
sourcePath := strings.TrimSpace(item.SourcePath)
|
|
if sourcePath != "" && trainingSupportedImportVideo(sourcePath) && fileExistsNonEmpty(sourcePath) {
|
|
if err := trainingExtractVideoMAEClipFrames(ctx, sourcePath, item.Second, clipDir); err == nil {
|
|
return trainingVideoMAEClipFrameCount(clipDir), nil
|
|
} else {
|
|
appLogln("videomae clip extract fallback:", item.SampleID, err)
|
|
}
|
|
}
|
|
|
|
fallbackFrame := trainingVideoMAEFrameFallbackPath(root, item.SampleID)
|
|
if err := trainingWriteStillVideoMAEClip(fallbackFrame, clipDir); err != nil {
|
|
return 0, err
|
|
}
|
|
|
|
return trainingVideoMAEClipFrameCount(clipDir), nil
|
|
}
|
|
|
|
func trainingCopyVideoMAEClip(srcDir string, dstDir string) (int, error) {
|
|
if err := os.RemoveAll(dstDir); err != nil {
|
|
return 0, err
|
|
}
|
|
if err := os.MkdirAll(dstDir, 0755); err != nil {
|
|
return 0, err
|
|
}
|
|
|
|
entries, err := os.ReadDir(srcDir)
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
|
|
copied := 0
|
|
for _, entry := range entries {
|
|
if entry.IsDir() {
|
|
continue
|
|
}
|
|
|
|
ext := strings.ToLower(filepath.Ext(entry.Name()))
|
|
if ext != ".jpg" && ext != ".jpeg" && ext != ".png" && ext != ".webp" {
|
|
continue
|
|
}
|
|
|
|
if err := copyFile(filepath.Join(srcDir, entry.Name()), filepath.Join(dstDir, entry.Name())); err != nil {
|
|
return copied, err
|
|
}
|
|
copied++
|
|
}
|
|
|
|
return copied, nil
|
|
}
|
|
|
|
func trainingEnsureVideoMAEValidationEntries(
|
|
root string,
|
|
entries []trainingVideoMAEManifestEntry,
|
|
) []trainingVideoMAEManifestEntry {
|
|
trainCount := 0
|
|
valCount := 0
|
|
trainEntries := []trainingVideoMAEManifestEntry{}
|
|
|
|
for _, entry := range entries {
|
|
switch strings.ToLower(strings.TrimSpace(entry.Split)) {
|
|
case "train":
|
|
trainCount++
|
|
trainEntries = append(trainEntries, entry)
|
|
case "val":
|
|
valCount++
|
|
}
|
|
}
|
|
|
|
if valCount >= minVideoMAEValCount || trainCount < minVideoMAETrainCount {
|
|
return entries
|
|
}
|
|
|
|
sort.SliceStable(trainEntries, func(i, j int) bool {
|
|
if trainEntries[i].Label == trainEntries[j].Label {
|
|
return trainEntries[i].SampleID < trainEntries[j].SampleID
|
|
}
|
|
return trainEntries[i].Label < trainEntries[j].Label
|
|
})
|
|
|
|
for _, entry := range trainEntries {
|
|
if valCount >= minVideoMAEValCount {
|
|
break
|
|
}
|
|
|
|
copyID := entry.SampleID + "_valcopy"
|
|
dstDir := filepath.Join(root, "videomae", "dataset", "clips", "val", copyID)
|
|
frameCount, err := trainingCopyVideoMAEClip(entry.ClipDir, dstDir)
|
|
if err != nil || frameCount <= 0 {
|
|
if err != nil {
|
|
appLogln("videomae val copy failed:", entry.SampleID, err)
|
|
}
|
|
continue
|
|
}
|
|
|
|
copyEntry := entry
|
|
copyEntry.SampleID = copyID
|
|
copyEntry.Split = "val"
|
|
copyEntry.ClipDir = dstDir
|
|
copyEntry.FrameCount = frameCount
|
|
entries = append(entries, copyEntry)
|
|
valCount++
|
|
}
|
|
|
|
return entries
|
|
}
|
|
|
|
func trainingSyncVideoMAEDataset(
|
|
ctx context.Context,
|
|
root string,
|
|
) (trainCount int, valCount int, written int, err error) {
|
|
if err := trainingEnsureVideoMAEDirs(root); err != nil {
|
|
return 0, 0, 0, err
|
|
}
|
|
if err := trainingRemoveVideoMAEGeneratedClips(root); err != nil {
|
|
return 0, 0, 0, err
|
|
}
|
|
|
|
items, err := trainingReadAnnotations(root)
|
|
if err != nil {
|
|
return 0, 0, 0, err
|
|
}
|
|
|
|
labelSet, err := trainingVideoMAELabelSet()
|
|
if err != nil {
|
|
return 0, 0, 0, err
|
|
}
|
|
|
|
entries := []trainingVideoMAEManifestEntry{}
|
|
|
|
for _, item := range items {
|
|
if ctx.Err() != nil {
|
|
return 0, 0, written, ctx.Err()
|
|
}
|
|
|
|
sampleID := strings.TrimSpace(item.SampleID)
|
|
if sampleID == "" || strings.Contains(sampleID, "/") || strings.Contains(sampleID, "\\") {
|
|
continue
|
|
}
|
|
|
|
label := trainingVideoMAELabelForAnnotation(item)
|
|
if !labelSet[label] {
|
|
continue
|
|
}
|
|
|
|
split := trainingStableSplit(sampleID)
|
|
clipDir := filepath.Join(root, "videomae", "dataset", "clips", split, sampleID)
|
|
frameCount, clipErr := trainingWriteVideoMAEClipForAnnotation(ctx, root, item, clipDir)
|
|
if clipErr != nil {
|
|
appLogln("videomae sample sync skipped:", sampleID, clipErr)
|
|
continue
|
|
}
|
|
|
|
entry := trainingVideoMAEManifestEntry{
|
|
SampleID: sampleID,
|
|
Split: split,
|
|
Label: label,
|
|
ClipDir: clipDir,
|
|
SourcePath: strings.TrimSpace(item.SourcePath),
|
|
Second: item.Second,
|
|
FrameCount: frameCount,
|
|
}
|
|
entries = append(entries, entry)
|
|
written++
|
|
}
|
|
|
|
entries = trainingEnsureVideoMAEValidationEntries(root, entries)
|
|
|
|
for _, entry := range entries {
|
|
switch strings.ToLower(strings.TrimSpace(entry.Split)) {
|
|
case "train":
|
|
trainCount++
|
|
case "val":
|
|
valCount++
|
|
}
|
|
}
|
|
|
|
if err := trainingWriteVideoMAEManifest(root, entries); err != nil {
|
|
return trainCount, valCount, written, err
|
|
}
|
|
|
|
return trainCount, valCount, written, nil
|
|
}
|