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"`
}

const trainingNoSexPositionLabel = "keine"

func isNoSexPositionLabel(label string) bool {
	switch strings.ToLower(strings.TrimSpace(label)) {
	case "", trainingNoSexPositionLabel:
		return true
	default:
		return false
	}
}

func normalizeSexPositionLabel(label string) string {
	clean := strings.TrimSpace(label)
	if isNoSexPositionLabel(clean) {
		return trainingNoSexPositionLabel
	}
	return clean
}

func normalizeSexPositionLabels(values []string) []string {
	seen := map[string]bool{}
	out := []string{}
	hasNoPosition := false

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

		if isNoSexPositionLabel(label) {
			hasNoPosition = true
			continue
		}

		if seen[label] {
			continue
		}

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

	if hasNoPosition {
		out = append([]string{trainingNoSexPositionLabel}, out...)
	}

	return out
}

func trainingDetectionLabelsPath() string {
	if p, err := ensureTrainingDetectionLabelsFile(); err == nil {
		return p
	}

	// Fallback: Temp-embedded ML.
	if dir, err := trainingEmbeddedMLDir(); err == nil {
		p := filepath.Join(dir, "detection_labels.json")
		if st, err := os.Stat(p); err == nil && st != nil && !st.IsDir() && st.Size() > 0 {
			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 st, err := os.Stat(p); err == nil && st != nil && !st.IsDir() && st.Size() > 0 {
			return p
		}
	}

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

func trainingGeneratedRootDirNoLabels() (string, error) {
	backendRoot, err := trainingBackendRootDir()
	if err != nil {
		return "", err
	}

	root, err := filepath.Abs(filepath.Join(backendRoot, "generated", "training"))
	if err != nil {
		return "", err
	}

	if err := os.MkdirAll(root, 0755); err != nil {
		return "", err
	}

	return root, nil
}

func ensureTrainingDetectionLabelsFile() (string, error) {
	root, err := trainingGeneratedRootDirNoLabels()
	if err != nil {
		return "", err
	}

	if err := os.MkdirAll(root, 0755); err != nil {
		return "", appErrorf("generated/training konnte nicht erstellt werden: %w", err)
	}

	dstPath := filepath.Join(root, "detection_labels.json")

	if st, err := os.Stat(dstPath); err == nil && st != nil && !st.IsDir() && st.Size() > 0 {
		return dstPath, nil
	}

	b, err := embeddedMLFiles.ReadFile("ml/detection_labels.json")
	if err != nil {
		return "", appErrorf("embedded detection_labels.json fehlt: %w", err)
	}

	if err := os.MkdirAll(filepath.Dir(dstPath), 0755); err != nil {
		return "", appErrorf("generated/training konnte nicht erstellt werden: %w", err)
	}

	if err := os.WriteFile(dstPath, b, 0644); err != nil {
		return "", appErrorf("detection_labels.json konnte nicht nach generated/training kopiert werden: %w", err)
	}

	return dstPath, nil
}

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

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

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

	grouped.People = uniqueNonEmptyLabels(grouped.People)
	grouped.SexPositions = normalizeSexPositionLabels(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{trainingNoSexPositionLabel}
	}

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

	if len(grouped.People)+len(grouped.BodyParts)+len(grouped.Objects)+len(grouped.Clothing) == 0 {
		return TrainingGroupedLabels{}, appErrorf("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{}

	labels = append(labels, grouped.People...)

	for _, label := range grouped.SexPositions {
		clean := strings.TrimSpace(label)
		if isNoSexPositionLabel(clean) {
			continue
		}

		labels = append(labels, clean)
	}

	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 trainingPoseLabels() ([]string, error) {
	grouped, err := trainingGroupedLabels()
	if err != nil {
		return nil, err
	}

	labels := []string{}
	seen := map[string]bool{}

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

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

	if len(labels) == 0 {
		return nil, fmt.Errorf("no pose sex position labels configured")
	}

	return labels, nil
}

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

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

	return out, nil
}

func trainingPoseDatasetNamesYAML() (string, error) {
	labels, err := trainingPoseLabels()
	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{trainingNoSexPositionLabel},
		}
	}

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