nsfwapp/backend/training_label_loader.go

// backend\training_label_loader.go

package main

import (
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"strings"
)

type TrainingGroupedLabels struct {
	People       []string `json:"people"`
	SexPositions []string `json:"sexPositions"`
	BodyParts    []string `json:"bodyParts"`
	Objects      []string `json:"objects"`
	Clothing     []string `json:"clothing"`
}

func trainingDetectionLabelsPath() string {
	if dir, err := trainingEmbeddedMLDir(); err == nil {
		p := filepath.Join(dir, "detection_labels.json")
		if _, err := os.Stat(p); err == nil {
			return p
		}
	}

	projectRoot := trainingProjectRoot()

	candidates := []string{
		filepath.Join(projectRoot, "backend", "ml", "detection_labels.json"),
		filepath.Join(projectRoot, "backend", "detection_labels.json"),
		filepath.Join(projectRoot, "detection_labels.json"),
		filepath.Join("ml", "detection_labels.json"),
		filepath.Join("detection_labels.json"),
	}

	for _, p := range candidates {
		if _, err := os.Stat(p); err == nil {
			return p
		}
	}

	return filepath.Join(projectRoot, "backend", "ml", "detection_labels.json")
}

func trainingGroupedLabels() (TrainingGroupedLabels, error) {
	path := trainingDetectionLabelsPath()

	b, err := os.ReadFile(path)
	if err != nil {
		return TrainingGroupedLabels{}, fmt.Errorf("detection_labels.json konnte nicht gelesen werden: %w", err)
	}

	var grouped TrainingGroupedLabels
	if err := json.Unmarshal(b, &grouped); err != nil {
		return TrainingGroupedLabels{}, fmt.Errorf("detection_labels.json ist ungültig: %w", err)
	}

	grouped.People = uniqueNonEmptyLabels(grouped.People)
	grouped.SexPositions = uniqueNonEmptyLabels(grouped.SexPositions)
	grouped.BodyParts = uniqueNonEmptyLabels(grouped.BodyParts)
	grouped.Objects = uniqueNonEmptyLabels(grouped.Objects)
	grouped.Clothing = uniqueNonEmptyLabels(grouped.Clothing)

	if len(grouped.SexPositions) == 0 {
		grouped.SexPositions = []string{"unknown"}
	}

	if len(grouped.BodyParts)+len(grouped.Objects)+len(grouped.Clothing) == 0 {
		return TrainingGroupedLabels{}, fmt.Errorf("detection_labels.json enthält keine Detection-Labels")
	}

	if len(grouped.People)+len(grouped.BodyParts)+len(grouped.Objects)+len(grouped.Clothing) == 0 {
		return TrainingGroupedLabels{}, fmt.Errorf("detection_labels.json enthält keine Detection-Labels")
	}

	return grouped, nil
}

func trainingDetectorLabels() ([]string, error) {
	grouped, err := trainingGroupedLabels()
	if err != nil {
		return nil, err
	}

	labels := []string{}

	// Wichtig:
	// People zuerst oder zuletzt ist egal, aber die Reihenfolge bestimmt YOLO-Class-IDs.
	// Wenn du schon ein bestehendes Detector-Modell hast, musst du danach neu trainieren.
	labels = append(labels, grouped.People...)
	labels = append(labels, grouped.BodyParts...)
	labels = append(labels, grouped.Objects...)
	labels = append(labels, grouped.Clothing...)

	return uniqueNonEmptyLabels(labels), nil
}

func uniqueNonEmptyLabels(values []string) []string {
	seen := map[string]bool{}
	out := []string{}

	for _, value := range values {
		label := strings.TrimSpace(value)
		if label == "" || seen[label] {
			continue
		}

		seen[label] = true
		out = append(out, label)
	}

	return out
}

func trainingDetectorClassMap() (map[string]int, error) {
	labels, err := trainingDetectorLabels()
	if err != nil {
		return nil, err
	}

	out := map[string]int{}
	for i, label := range labels {
		out[label] = i
	}

	return out, nil
}

func trainingDetectorDatasetNamesYAML() (string, error) {
	labels, err := trainingDetectorLabels()
	if err != nil {
		return "", err
	}

	var b strings.Builder
	for i, label := range labels {
		b.WriteString(fmt.Sprintf("  %d: %s\n", i, label))
	}

	return b.String(), nil
}

func defaultTrainingLabelsFromJSON() TrainingLabels {
	grouped, err := trainingGroupedLabels()
	if err != nil {
		grouped = TrainingGroupedLabels{
			SexPositions: []string{"unknown"},
		}
	}

	return TrainingLabels{
		People:       grouped.People,
		SexPositions: grouped.SexPositions,
		BodyParts:    grouped.BodyParts,
		Objects:      grouped.Objects,
		Clothing:     grouped.Clothing,
	}
}