diff --git a/backend/main.go b/backend/main.go index ee26bbc..f18a5bf 100644 --- a/backend/main.go +++ b/backend/main.go @@ -11,6 +11,7 @@ import ( "errors" "fmt" "io" + "math" "net/http" "net/url" "os" @@ -1211,6 +1212,9 @@ func uniqueDestPath(dstDir, file string) (string, error) { } func clamp01(v float64) float64 { + if math.IsNaN(v) || math.IsInf(v, 0) { + return 0 + } if v < 0 { return 0 } diff --git a/backend/ml/detection_labels.json b/backend/ml/detection_labels.json new file mode 100644 index 0000000..73abb11 --- /dev/null +++ b/backend/ml/detection_labels.json @@ -0,0 +1,61 @@ +{ + "sexPositions": [ + "unknown", + "solo", + "missionary", + "doggy", + "cowgirl", + "reverse_cowgirl", + "cunnilingus", + "standing", + "standing_doggy", + "spooning", + "sideways", + "sitting", + "facesitting", + "handjob", + "blowjob", + "toy_play", + "fingering", + "69", + "other" + ], + "bodyParts": [ + "anus", + "back", + "breasts", + "buttocks", + "chest", + "penis", + "tongue", + "vagina" + ], + "objects": [ + "bath", + "blindfold", + "buttplug", + "collar", + "dildo", + "handcuffs", + "rope", + "shower", + "strapon", + "towel", + "vibrator" + ], + "clothing": [ + "bikini", + "bra", + "dress", + "fishnet", + "heels", + "hotpants", + "lingerie", + "nude", + "panties", + "shirt", + "skirt", + "stockings", + "top" + ] +} \ No newline at end of file diff --git a/backend/ml/predict_detector_model.py b/backend/ml/predict_detector_model.py new file mode 100644 index 0000000..a23cfaa --- /dev/null +++ b/backend/ml/predict_detector_model.py @@ -0,0 +1,86 @@ +# backend\ml\predict_detector_model.py + +import argparse +import json +from pathlib import Path + +from PIL import Image +from ultralytics import YOLO + + +def clamp01(v): + return max(0.0, min(1.0, float(v))) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--root", required=True) + parser.add_argument("--image", required=True) + parser.add_argument("--conf", type=float, default=0.30) + parser.add_argument("--imgsz", type=int, default=640) + args = parser.parse_args() + + root = Path(args.root) + image_path = Path(args.image) + + model_path = root / "detector" / "best.pt" + if not model_path.exists(): + print(json.dumps({ + "available": False, + "source": "detector_missing", + "boxes": [], + })) + return + + img = Image.open(image_path).convert("RGB") + img_w, img_h = img.size + + model = YOLO(str(model_path)) + results = model.predict( + source=str(image_path), + conf=args.conf, + imgsz=args.imgsz, + verbose=False, + device=0 if __import__("torch").cuda.is_available() else "cpu", + ) + + boxes = [] + + if results: + r = results[0] + names = r.names or {} + + if r.boxes is not None: + for b in r.boxes: + cls_id = int(b.cls[0].item()) + score = float(b.conf[0].item()) + label = str(names.get(cls_id, cls_id)) + + x1, y1, x2, y2 = [float(v) for v in b.xyxy[0].tolist()] + + x = clamp01(x1 / img_w) + y = clamp01(y1 / img_h) + w = clamp01((x2 - x1) / img_w) + h = clamp01((y2 - y1) / img_h) + + if w <= 0 or h <= 0: + continue + + boxes.append({ + "label": label, + "score": score, + "x": x, + "y": y, + "w": w, + "h": h, + }) + + print(json.dumps({ + "available": True, + "source": "yolo_detector", + "boxes": boxes, + }, ensure_ascii=False)) + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/backend/ml/predict_scene_model.py b/backend/ml/predict_scene_model.py new file mode 100644 index 0000000..ed0091e --- /dev/null +++ b/backend/ml/predict_scene_model.py @@ -0,0 +1,197 @@ +# backend\ml\predict_scene_model.py + +import argparse +import json +from pathlib import Path + +import numpy as np +import torch +from PIL import Image +from torchvision import models + +import subprocess +import sys + + +def empty_prediction(source="no_model"): + return { + "modelAvailable": False, + "source": source, + "peopleCount": 0, + "maleCount": 0, + "femaleCount": 0, + "unknownCount": 0, + "sexPosition": "unknown", + "sexPositionScore": 0, + "bodyPartsPresent": [], + "objectsPresent": [], + "clothingPresent": [], + "boxes": [], + } + + +def load_encoder(): + weights = models.ResNet18_Weights.DEFAULT + model = models.resnet18(weights=weights) + model.fc = torch.nn.Identity() + model.eval() + + preprocess = weights.transforms() + return model, preprocess + + +def embed_image(model, preprocess, image_path: Path): + img = Image.open(image_path).convert("RGB") + x = preprocess(img).unsqueeze(0) + + with torch.no_grad(): + emb = model(x).cpu().numpy()[0].astype("float32") + + norm = np.linalg.norm(emb) + if norm > 0: + emb = emb / norm + + return emb + + +def weighted_mode_string(items, weights, fallback="unknown"): + scores = {} + + for item, weight in zip(items, weights): + key = str(item or fallback) + scores[key] = scores.get(key, 0.0) + float(weight) + + if not scores: + return fallback, 0.0 + + label, score = max(scores.items(), key=lambda x: x[1]) + total = sum(scores.values()) or 1.0 + return label, float(score / total) + + +def weighted_int(items, weights): + if not items: + return 0 + + total_w = float(np.sum(weights)) or 1.0 + value = sum(float(v or 0) * float(w) for v, w in zip(items, weights)) / total_w + return int(round(value)) + + +def weighted_multilabel(targets, key, weights, threshold=0.35): + scores = {} + total = float(np.sum(weights)) or 1.0 + + for target, weight in zip(targets, weights): + for label in target.get(key) or []: + label = str(label).strip() + if not label: + continue + scores[label] = scores.get(label, 0.0) + float(weight) + + out = [] + for label, score in scores.items(): + conf = float(score / total) + if conf >= threshold: + out.append({ + "label": label, + "score": conf, + }) + + out.sort(key=lambda x: x["score"], reverse=True) + return out + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--root", required=True) + parser.add_argument("--image", required=True) + args = parser.parse_args() + + root = Path(args.root) + image_path = Path(args.image) + + model_path = root / "model" / "scene_knn.npz" + targets_path = root / "model" / "targets.json" + + if not model_path.exists() or not targets_path.exists(): + print(json.dumps(empty_prediction("model_missing"))) + return + + data = np.load(model_path) + embeddings = data["embeddings"].astype("float32") + + with targets_path.open("r", encoding="utf-8") as f: + targets = json.load(f) + + if len(embeddings) == 0 or len(targets) == 0: + print(json.dumps(empty_prediction("model_empty"))) + return + + encoder, preprocess = load_encoder() + emb = embed_image(encoder, preprocess, image_path) + + sims = embeddings @ emb + + k = min(7, len(sims)) + idx = np.argsort(-sims)[:k] + top_targets = [targets[int(i)] for i in idx] + + raw_weights = np.maximum(sims[idx], 0.0) + if float(np.sum(raw_weights)) <= 0: + raw_weights = np.ones_like(raw_weights, dtype="float32") + + weights = raw_weights / (float(np.sum(raw_weights)) or 1.0) + + sex_labels = [t.get("sexPosition") or "unknown" for t in top_targets] + sex_position, sex_score = weighted_mode_string(sex_labels, weights, "unknown") + + pred = { + "modelAvailable": True, + "source": "resnet18_knn", + "peopleCount": weighted_int([t.get("peopleCount") for t in top_targets], weights), + "maleCount": weighted_int([t.get("maleCount") for t in top_targets], weights), + "femaleCount": weighted_int([t.get("femaleCount") for t in top_targets], weights), + "unknownCount": weighted_int([t.get("unknownCount") for t in top_targets], weights), + "sexPosition": sex_position, + "sexPositionScore": sex_score, + "bodyPartsPresent": weighted_multilabel(top_targets, "bodyPartsPresent", weights), + "objectsPresent": weighted_multilabel(top_targets, "objectsPresent", weights), + "clothingPresent": weighted_multilabel(top_targets, "clothingPresent", weights), + "boxes": predict_boxes(root, image_path), + } + + print(json.dumps(pred, ensure_ascii=False)) + +def predict_boxes(root: Path, image_path: Path): + script = Path(__file__).parent / "predict_detector_model.py" + + if not script.exists(): + return [] + + try: + proc = subprocess.run( + [ + sys.executable, + str(script), + "--root", + str(root), + "--image", + str(image_path), + ], + capture_output=True, + text=True, + timeout=20, + ) + + if proc.returncode != 0: + return [] + + data = json.loads(proc.stdout or "{}") + return data.get("boxes") or [] + except Exception: + return [] + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/backend/ml/requirements-ml.txt b/backend/ml/requirements-ml.txt new file mode 100644 index 0000000..f73288e --- /dev/null +++ b/backend/ml/requirements-ml.txt @@ -0,0 +1,4 @@ +torch +torchvision +pillow +numpy \ No newline at end of file diff --git a/backend/ml/train_detector_model.py b/backend/ml/train_detector_model.py new file mode 100644 index 0000000..0059807 --- /dev/null +++ b/backend/ml/train_detector_model.py @@ -0,0 +1,56 @@ +# backend\ml\train_detector_model.py + +import argparse +import json +from pathlib import Path +from ultralytics import YOLO + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--root", required=True) + parser.add_argument("--base", default="yolo11n.pt") + parser.add_argument("--epochs", default="80") + parser.add_argument("--imgsz", default="640") + args = parser.parse_args() + + root = Path(args.root).resolve() + yaml_path = root / "detector" / "dataset" / "dataset.yaml" + runs_dir = root / "detector" / "runs" + + if not yaml_path.exists(): + raise SystemExit(f"dataset.yaml not found: {yaml_path}") + + model = YOLO(args.base) + + result = model.train( + data=str(yaml_path), + epochs=int(args.epochs), + imgsz=int(args.imgsz), + project=str(runs_dir), + name="detect", + exist_ok=True, + device="cpu", + workers=2, + patience=20, + ) + + best = runs_dir / "detect" / "weights" / "best.pt" + if not best.exists(): + raise SystemExit(f"best.pt not found after training: {best}") + + out_dir = root / "detector" / "model" + out_dir.mkdir(parents=True, exist_ok=True) + + final_model = out_dir / "best.pt" + final_model.write_bytes(best.read_bytes()) + + print(json.dumps({ + "ok": True, + "model": str(final_model), + "runs": str(runs_dir / "detect"), + })) + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/backend/ml/train_scene_model.py b/backend/ml/train_scene_model.py new file mode 100644 index 0000000..3a66462 --- /dev/null +++ b/backend/ml/train_scene_model.py @@ -0,0 +1,167 @@ +# backend\ml\train_scene_model.py + +import argparse +import json +from pathlib import Path + +import numpy as np +import torch +from PIL import Image +from torchvision import models, transforms + + +def load_encoder(): + weights = models.ResNet18_Weights.DEFAULT + model = models.resnet18(weights=weights) + model.fc = torch.nn.Identity() + model.eval() + + preprocess = weights.transforms() + return model, preprocess + + +def read_jsonl(path: Path): + if not path.exists(): + return [] + + rows = [] + with path.open("r", encoding="utf-8") as f: + for line in f: + line = line.strip() + if not line: + continue + try: + rows.append(json.loads(line)) + except Exception: + pass + return rows + + +def prediction_target(annotation): + pred = annotation.get("prediction") or {} + + return { + "peopleCount": int(pred.get("peopleCount") or 0), + "maleCount": int(pred.get("maleCount") or 0), + "femaleCount": int(pred.get("femaleCount") or 0), + "unknownCount": int(pred.get("unknownCount") or 0), + "sexPosition": str(pred.get("sexPosition") or "unknown"), + "bodyPartsPresent": [x.get("label") for x in pred.get("bodyPartsPresent") or [] if x.get("label")], + "objectsPresent": [x.get("label") for x in pred.get("objectsPresent") or [] if x.get("label")], + "clothingPresent": [x.get("label") for x in pred.get("clothingPresent") or [] if x.get("label")], + } + + +def correction_target(annotation): + corr = annotation.get("correction") or {} + + return { + "peopleCount": int(corr.get("peopleCount") or 0), + "maleCount": int(corr.get("maleCount") or 0), + "femaleCount": int(corr.get("femaleCount") or 0), + "unknownCount": int(corr.get("unknownCount") or 0), + "sexPosition": str(corr.get("sexPosition") or "unknown"), + "bodyPartsPresent": list(corr.get("bodyPartsPresent") or []), + "objectsPresent": list(corr.get("objectsPresent") or []), + "clothingPresent": list(corr.get("clothingPresent") or []), + } + + +def target_from_annotation(annotation): + if annotation.get("accepted") is True: + return prediction_target(annotation) + + return correction_target(annotation) + + +def embed_image(model, preprocess, image_path: Path): + img = Image.open(image_path).convert("RGB") + x = preprocess(img).unsqueeze(0) + + with torch.no_grad(): + emb = model(x).cpu().numpy()[0].astype("float32") + + norm = np.linalg.norm(emb) + if norm > 0: + emb = emb / norm + + return emb + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--root", required=True) + args = parser.parse_args() + + root = Path(args.root) + feedback_path = root / "feedback.jsonl" + frames_dir = root / "frames" + model_dir = root / "model" + model_dir.mkdir(parents=True, exist_ok=True) + + rows = read_jsonl(feedback_path) + + model, preprocess = load_encoder() + + embeddings = [] + targets = [] + used = 0 + skipped = 0 + + for row in rows: + sample_id = str(row.get("sampleId") or "").strip() + if not sample_id: + skipped += 1 + continue + + image_path = frames_dir / f"{sample_id}.jpg" + if not image_path.exists(): + skipped += 1 + continue + + try: + emb = embed_image(model, preprocess, image_path) + target = target_from_annotation(row) + + embeddings.append(emb) + targets.append(target) + used += 1 + except Exception: + skipped += 1 + + if used < 5: + raise SystemExit(f"too few usable samples: {used}") + + embeddings_np = np.stack(embeddings).astype("float32") + + np.savez_compressed( + model_dir / "scene_knn.npz", + embeddings=embeddings_np, + ) + + with (model_dir / "targets.json").open("w", encoding="utf-8") as f: + json.dump(targets, f, ensure_ascii=False, indent=2) + + with (model_dir / "status.json").open("w", encoding="utf-8") as f: + json.dump( + { + "ok": True, + "usedSamples": used, + "skippedSamples": skipped, + "model": "resnet18_knn", + }, + f, + ensure_ascii=False, + indent=2, + ) + + print(json.dumps({ + "ok": True, + "usedSamples": used, + "skippedSamples": skipped, + "modelPath": str(model_dir / "scene_knn.npz"), + })) + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/backend/ml_embed.go b/backend/ml_embed.go new file mode 100644 index 0000000..3f83dfa --- /dev/null +++ b/backend/ml_embed.go @@ -0,0 +1,42 @@ +package main + +import ( + "embed" + "os" + "path/filepath" +) + +//go:embed ml/*.py ml/*.json +var embeddedMLFiles embed.FS + +func trainingEmbeddedMLDir() (string, error) { + dir := filepath.Join(os.TempDir(), "nsfwapp-ml") + + if err := os.MkdirAll(dir, 0755); err != nil { + return "", err + } + + files := []string{ + "predict_scene_model.py", + "train_scene_model.py", + "predict_detector_model.py", + "train_detector_model.py", + "detection_labels.json", + } + + for _, name := range files { + srcPath := filepath.Join("ml", name) + + b, err := embeddedMLFiles.ReadFile(filepath.ToSlash(srcPath)) + if err != nil { + return "", err + } + + dstPath := filepath.Join(dir, name) + if err := os.WriteFile(dstPath, b, 0644); err != nil { + return "", err + } + } + + return dir, nil +} diff --git a/backend/routes.go b/backend/routes.go index db35d58..dbd3d33 100644 --- a/backend/routes.go +++ b/backend/routes.go @@ -69,7 +69,16 @@ func registerRoutes(mux *http.ServeMux, auth *AuthManager) *ModelStore { api.HandleFunc("/api/record/delete-many", handleDeleteMany) api.HandleFunc("/api/record/keep-many", handleKeepMany) api.HandleFunc("/api/record/concurrency", recordConcurrentDownloadsStatus) - mux.HandleFunc("/api/record/release-file", handleReleaseFileTasks) + api.HandleFunc("/api/record/release-file", handleReleaseFileTasks) + + // Training / ML Annotation + api.HandleFunc("/api/training/labels", trainingLabelsHandler) + api.HandleFunc("/api/training/next", trainingNextHandler) + api.HandleFunc("/api/training/frame", trainingFrameHandler) + api.HandleFunc("/api/training/feedback", trainingFeedbackHandler) + api.HandleFunc("/api/training/train", trainingTrainHandler) + api.HandleFunc("/api/training/status", trainingStatusHandler) + api.HandleFunc("/api/training/delete-all", trainingDeleteAllHandler) api.HandleFunc("/api/chaturbate/online", chaturbateOnlineHandler) api.HandleFunc("/api/chaturbate/biocontext", chaturbateBioContextHandler) diff --git a/backend/training.go b/backend/training.go new file mode 100644 index 0000000..62a1a44 --- /dev/null +++ b/backend/training.go @@ -0,0 +1,1180 @@ +// backend\training.go + +package main + +import ( + "crypto/sha1" + "encoding/hex" + "encoding/json" + "errors" + "fmt" + "math" + "math/rand" + "net/http" + "os" + "os/exec" + "path/filepath" + "sort" + "strconv" + "strings" + "syscall" + "time" +) + +type TrainingLabels struct { + SexPositions []string `json:"sexPositions"` + BodyParts []string `json:"bodyParts"` + Objects []string `json:"objects"` + Clothing []string `json:"clothing"` +} + +type TrainingBox struct { + Label string `json:"label"` + Score float64 `json:"score,omitempty"` + X float64 `json:"x"` + Y float64 `json:"y"` + W float64 `json:"w"` + H float64 `json:"h"` +} + +type TrainingScoredLabel struct { + Label string `json:"label"` + Score float64 `json:"score"` +} + +type TrainingPrediction struct { + ModelAvailable bool `json:"modelAvailable"` + Source string `json:"source,omitempty"` + + PeopleCount int `json:"peopleCount"` + MaleCount int `json:"maleCount"` + FemaleCount int `json:"femaleCount"` + UnknownCount int `json:"unknownCount"` + SexPosition string `json:"sexPosition"` + SexPositionScore float64 `json:"sexPositionScore"` + BodyPartsPresent []TrainingScoredLabel `json:"bodyPartsPresent"` + ObjectsPresent []TrainingScoredLabel `json:"objectsPresent"` + ClothingPresent []TrainingScoredLabel `json:"clothingPresent"` + Boxes []TrainingBox `json:"boxes,omitempty"` +} + +type TrainingCorrection struct { + PeopleCount int `json:"peopleCount"` + MaleCount int `json:"maleCount"` + FemaleCount int `json:"femaleCount"` + UnknownCount int `json:"unknownCount"` + SexPosition string `json:"sexPosition"` + BodyPartsPresent []string `json:"bodyPartsPresent"` + ObjectsPresent []string `json:"objectsPresent"` + ClothingPresent []string `json:"clothingPresent"` + Boxes []TrainingBox `json:"boxes,omitempty"` +} + +type TrainingSample struct { + SampleID string `json:"sampleId"` + FrameURL string `json:"frameUrl"` + SourceFile string `json:"sourceFile"` + SourcePath string `json:"sourcePath,omitempty"` + SourceSizeBytes int64 `json:"sourceSizeBytes,omitempty"` + Second float64 `json:"second"` + CreatedAt string `json:"createdAt"` + Prediction TrainingPrediction `json:"prediction"` +} + +type TrainingFeedbackRequest struct { + SampleID string `json:"sampleId"` + Accepted bool `json:"accepted"` + Correction *TrainingCorrection `json:"correction,omitempty"` + Notes string `json:"notes,omitempty"` +} + +type TrainingAnnotation struct { + SampleID string `json:"sampleId"` + FrameURL string `json:"frameUrl"` + SourceFile string `json:"sourceFile"` + SourcePath string `json:"sourcePath,omitempty"` + SourceSizeBytes int64 `json:"sourceSizeBytes,omitempty"` + Second float64 `json:"second"` + CreatedAt string `json:"createdAt"` + AnsweredAt string `json:"answeredAt"` + Prediction TrainingPrediction `json:"prediction"` + Accepted bool `json:"accepted"` + Correction *TrainingCorrection `json:"correction,omitempty"` + Notes string `json:"notes,omitempty"` +} + +const minTrainingFeedbackCount = 5 + +func trainingLabelsHandler(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodGet { + trainingWriteError(w, http.StatusMethodNotAllowed, "method not allowed") + return + } + + trainingWriteJSON(w, http.StatusOK, defaultTrainingLabelsFromJSON()) +} + +func trainingNextHandler(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodGet { + trainingWriteError(w, http.StatusMethodNotAllowed, "method not allowed") + return + } + + forceNew := r.URL.Query().Get("force") == "1" || + strings.EqualFold(r.URL.Query().Get("force"), "true") + + root, err := trainingRootDir() + if err != nil { + trainingWriteError(w, http.StatusInternalServerError, err.Error()) + return + } + + if !forceNew { + if sample, ok, err := trainingLatestOpenSample(root); err != nil { + trainingWriteError(w, http.StatusInternalServerError, err.Error()) + return + } else if ok { + trainingWriteJSON(w, http.StatusOK, sample) + return + } + } + + sample, err := trainingCreateNextSample() + if err != nil { + trainingWriteError(w, http.StatusInternalServerError, err.Error()) + return + } + + trainingWriteJSON(w, http.StatusOK, sample) +} + +func trainingLatestOpenSample(root string) (*TrainingSample, bool, error) { + answered, err := trainingAnsweredSampleIDs(root) + if err != nil { + return nil, false, err + } + + samplesDir := filepath.Join(root, "samples") + + entries, err := os.ReadDir(samplesDir) + if err != nil { + if os.IsNotExist(err) { + return nil, false, nil + } + return nil, false, err + } + + type sampleFile struct { + id string + path string + modTime time.Time + } + + files := []sampleFile{} + + for _, entry := range entries { + if entry.IsDir() { + continue + } + + name := entry.Name() + if strings.ToLower(filepath.Ext(name)) != ".json" { + continue + } + + id := strings.TrimSuffix(name, filepath.Ext(name)) + if id == "" || answered[id] { + continue + } + + info, err := entry.Info() + if err != nil { + continue + } + + files = append(files, sampleFile{ + id: id, + path: filepath.Join(samplesDir, name), + modTime: info.ModTime(), + }) + } + + sort.Slice(files, func(i, j int) bool { + return files[i].modTime.After(files[j].modTime) + }) + + for _, file := range files { + sample, err := trainingReadSample(root, file.id) + if err != nil { + continue + } + + framePath := filepath.Join(root, "frames", sample.SampleID+".jpg") + if !fileExistsNonEmpty(framePath) { + continue + } + + return sample, true, nil + } + + return nil, false, nil +} + +func trainingAnsweredSampleIDs(root string) (map[string]bool, error) { + out := map[string]bool{} + + path := filepath.Join(root, "feedback.jsonl") + + b, err := os.ReadFile(path) + if err != nil { + if os.IsNotExist(err) { + return out, nil + } + return nil, err + } + + for _, line := range strings.Split(string(b), "\n") { + line = strings.TrimSpace(line) + if line == "" { + continue + } + + var row TrainingAnnotation + if err := json.Unmarshal([]byte(line), &row); err != nil { + continue + } + + id := strings.TrimSpace(row.SampleID) + if id != "" { + out[id] = true + } + } + + return out, nil +} + +func trainingFrameHandler(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodGet { + trainingWriteError(w, http.StatusMethodNotAllowed, "method not allowed") + return + } + + id := strings.TrimSpace(r.URL.Query().Get("id")) + if id == "" || strings.Contains(id, "/") || strings.Contains(id, "\\") { + trainingWriteError(w, http.StatusBadRequest, "invalid frame id") + return + } + + root, err := trainingRootDir() + if err != nil { + trainingWriteError(w, http.StatusInternalServerError, err.Error()) + return + } + + path := filepath.Join(root, "frames", id+".jpg") + + if _, err := os.Stat(path); err != nil { + trainingWriteError(w, http.StatusNotFound, "frame not found") + return + } + + w.Header().Set("Cache-Control", "no-store") + http.ServeFile(w, r, path) +} + +func trainingFeedbackHandler(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodPost { + trainingWriteError(w, http.StatusMethodNotAllowed, "method not allowed") + return + } + + var req TrainingFeedbackRequest + if err := json.NewDecoder(r.Body).Decode(&req); err != nil { + trainingWriteError(w, http.StatusBadRequest, "invalid json") + return + } + + req.SampleID = strings.TrimSpace(req.SampleID) + if req.SampleID == "" { + trainingWriteError(w, http.StatusBadRequest, "sampleId missing") + return + } + + root, err := trainingRootDir() + if err != nil { + trainingWriteError(w, http.StatusInternalServerError, err.Error()) + return + } + + sample, err := trainingReadSample(root, req.SampleID) + if err != nil { + trainingWriteError(w, http.StatusNotFound, "sample not found") + return + } + + annotation := TrainingAnnotation{ + SampleID: sample.SampleID, + FrameURL: sample.FrameURL, + SourceFile: sample.SourceFile, + SourcePath: sample.SourcePath, + SourceSizeBytes: sample.SourceSizeBytes, + Second: sample.Second, + CreatedAt: sample.CreatedAt, + AnsweredAt: time.Now().UTC().Format(time.RFC3339), + Prediction: sample.Prediction, + Accepted: req.Accepted, + Correction: req.Correction, + Notes: strings.TrimSpace(req.Notes), + } + + if !annotation.Accepted && annotation.Correction == nil { + trainingWriteError(w, http.StatusBadRequest, "correction missing") + return + } + + if err := trainingAppendAnnotation(root, annotation); err != nil { + trainingWriteError(w, http.StatusInternalServerError, err.Error()) + return + } + + if req.Correction != nil && len(req.Correction.Boxes) > 0 { + if err := trainingWriteDetectorSample(root, sample, req.Correction.Boxes); err != nil { + fmt.Println("⚠️ detector sample write failed:", err) + } + } + + trainingWriteJSON(w, http.StatusOK, map[string]any{ + "ok": true, + }) +} + +func trainingHasDetectorTrainingData(imagesDir string, labelsDir string) bool { + imageExts := map[string]bool{ + ".jpg": true, + ".jpeg": true, + ".png": true, + ".webp": true, + } + + entries, err := os.ReadDir(imagesDir) + if err != nil { + return false + } + + count := 0 + + for _, e := range entries { + if e.IsDir() { + continue + } + + ext := strings.ToLower(filepath.Ext(e.Name())) + if !imageExts[ext] { + continue + } + + id := strings.TrimSuffix(e.Name(), filepath.Ext(e.Name())) + labelPath := filepath.Join(labelsDir, id+".txt") + + if fileExistsNonEmpty(labelPath) { + count++ + } + } + + return count >= 5 +} + +func trainingTrainHandler(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodPost { + trainingWriteError(w, http.StatusMethodNotAllowed, "method not allowed") + return + } + + root, err := trainingRootDir() + if err != nil { + trainingWriteError(w, http.StatusInternalServerError, err.Error()) + return + } + + if err := trainingEnsureDetectorDirs(root); err != nil { + trainingWriteError(w, http.StatusInternalServerError, err.Error()) + return + } + + feedbackPath := filepath.Join(root, "feedback.jsonl") + count, _ := trainingCountAnnotations(feedbackPath) + + if count < minTrainingFeedbackCount { + trainingWriteError( + w, + http.StatusBadRequest, + fmt.Sprintf( + "Zu wenige Bewertungen. Mindestens %d, aktuell %d.", + minTrainingFeedbackCount, + count, + ), + ) + return + } + + python := trainingPythonExe() + + // -------------------------------------------------- + // 1) Scene-Modell trainieren: ResNet18-KNN + // -------------------------------------------------- + sceneScript := trainingScriptPath("train_scene_model.py") + + sceneCmd := exec.Command( + python, + sceneScript, + "--root", root, + ) + + sceneCmd.SysProcAttr = &syscall.SysProcAttr{ + HideWindow: true, + CreationFlags: 0x08000000, + } + + sceneOut, err := sceneCmd.CombinedOutput() + if err != nil { + trainingWriteError( + w, + http.StatusInternalServerError, + fmt.Sprintf("scene training failed: %v: %s", err, strings.TrimSpace(string(sceneOut))), + ) + return + } + + // -------------------------------------------------- + // 2) Detector trainieren: YOLO + // Nur starten, wenn dataset.yaml existiert. + // -------------------------------------------------- + detectorOutput := "" + detectorStatus := "skipped" + + detectorDatasetYAML := filepath.Join(root, "detector", "dataset", "dataset.yaml") + detectorTrainImages := filepath.Join(root, "detector", "dataset", "images", "train") + detectorTrainLabels := filepath.Join(root, "detector", "dataset", "labels", "train") + detectorValImages := filepath.Join(root, "detector", "dataset", "images", "val") + detectorValLabels := filepath.Join(root, "detector", "dataset", "labels", "val") + + trainCount := trainingCountDetectorSamples(detectorTrainImages, detectorTrainLabels) + valCount := trainingCountDetectorSamples(detectorValImages, detectorValLabels) + + if fileExistsNonEmpty(detectorDatasetYAML) && trainCount >= 5 && valCount >= 1 { + detectorScript := trainingScriptPath("train_detector_model.py") + + detectorCmd := exec.Command( + python, + detectorScript, + "--root", root, + "--base", "yolo11n.pt", + "--epochs", "80", + "--imgsz", "640", + ) + + detectorCmd.SysProcAttr = &syscall.SysProcAttr{ + HideWindow: true, + CreationFlags: 0x08000000, + } + + detectorOut, detectorErr := detectorCmd.CombinedOutput() + detectorOutput = strings.TrimSpace(string(detectorOut)) + + if detectorErr != nil { + detectorStatus = "failed" + fmt.Println("⚠️ detector training failed:", detectorErr, detectorOutput) + } else { + detectorStatus = "trained" + } + } else { + detectorStatus = "skipped_no_detector_data" + detectorOutput = fmt.Sprintf( + "Detector übersprungen: zu wenige YOLO-Box-Labels. Train=%d, Val=%d. Benötigt: mindestens 5 Train und 1 Val.", + trainCount, + valCount, + ) + } + + message := "Training abgeschlossen. Neue Bilder werden jetzt mit dem Scene-Modell analysiert." + if detectorStatus == "trained" { + message = "Training abgeschlossen. Scene-Modell und Detector wurden trainiert." + } + if detectorStatus == "failed" { + message = "Scene-Modell wurde trainiert, Detector-Training ist fehlgeschlagen." + } + + trainingWriteJSON(w, http.StatusOK, map[string]any{ + "ok": true, + "message": message, + "annotationCount": count, + "sceneStatus": "trained", + "sceneOutput": strings.TrimSpace(string(sceneOut)), + "detectorStatus": detectorStatus, + "detectorOutput": detectorOutput, + }) +} + +func trainingStatusHandler(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodGet { + trainingWriteError(w, http.StatusMethodNotAllowed, "method not allowed") + return + } + + root, err := trainingRootDir() + if err != nil { + trainingWriteError(w, http.StatusInternalServerError, err.Error()) + return + } + + feedbackPath := filepath.Join(root, "feedback.jsonl") + count, _ := trainingCountAnnotations(feedbackPath) + + trainingWriteJSON(w, http.StatusOK, map[string]any{ + "ok": true, + "feedbackCount": count, + "requiredCount": minTrainingFeedbackCount, + "canTrain": count >= minTrainingFeedbackCount, + }) +} + +func trainingDeleteAllHandler(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodDelete && r.Method != http.MethodPost { + trainingWriteError(w, http.StatusMethodNotAllowed, "method not allowed") + return + } + + root, err := trainingRootDir() + if err != nil { + trainingWriteError(w, http.StatusInternalServerError, err.Error()) + return + } + + if err := os.RemoveAll(root); err != nil { + trainingWriteError(w, http.StatusInternalServerError, fmt.Sprintf("Trainingsdaten konnten nicht gelöscht werden: %v", err)) + return + } + + // Basisordner direkt wieder anlegen, damit die nächsten API-Aufrufe sauber funktionieren. + if _, err := trainingRootDir(); err != nil { + trainingWriteError(w, http.StatusInternalServerError, err.Error()) + return + } + + trainingWriteJSON(w, http.StatusOK, map[string]any{ + "ok": true, + "message": "Alle Trainingsdaten wurden gelöscht.", + }) +} + +func trainingCountDetectorSamples(imagesDir string, labelsDir string) int { + imageExts := map[string]bool{ + ".jpg": true, + ".jpeg": true, + ".png": true, + ".webp": true, + } + + entries, err := os.ReadDir(imagesDir) + if err != nil { + return 0 + } + + count := 0 + + for _, e := range entries { + if e.IsDir() { + continue + } + + ext := strings.ToLower(filepath.Ext(e.Name())) + if !imageExts[ext] { + continue + } + + id := strings.TrimSuffix(e.Name(), filepath.Ext(e.Name())) + labelPath := filepath.Join(labelsDir, id+".txt") + + if fileExistsNonEmpty(labelPath) { + count++ + } + } + + return count +} + +func trainingWriteDetectorDatasetYAML(root string) error { + datasetDir := filepath.Join(root, "detector", "dataset") + + absDatasetDir, err := filepath.Abs(datasetDir) + if err != nil { + return err + } + + namesYAML, err := trainingDetectorDatasetNamesYAML() + if err != nil { + return err + } + + path := filepath.Join(datasetDir, "dataset.yaml") + yoloPath := filepath.ToSlash(absDatasetDir) + + body := fmt.Sprintf(`path: %s +train: images/train +val: images/val + +names: +%s`, yoloPath, namesYAML) + + return os.WriteFile(path, []byte(body), 0644) +} + +func trainingEnsureDetectorDirs(root string) error { + dirs := []string{ + filepath.Join(root, "detector"), + filepath.Join(root, "detector", "dataset"), + filepath.Join(root, "detector", "dataset", "images"), + filepath.Join(root, "detector", "dataset", "images", "train"), + filepath.Join(root, "detector", "dataset", "images", "val"), + filepath.Join(root, "detector", "dataset", "labels"), + filepath.Join(root, "detector", "dataset", "labels", "train"), + filepath.Join(root, "detector", "dataset", "labels", "val"), + filepath.Join(root, "detector", "runs"), + } + + for _, dir := range dirs { + if err := os.MkdirAll(dir, 0755); err != nil { + return err + } + } + + if err := trainingWriteDetectorDatasetYAML(root); err != nil { + return err + } + + return nil +} + +func trainingCreateNextSample() (*TrainingSample, error) { + settings := getSettings() + + doneDir := strings.TrimSpace(settings.DoneDir) + if doneDir == "" { + return nil, errors.New("doneDir ist leer") + } + + videoPath, err := trainingPickRandomVideo(doneDir) + if err != nil { + return nil, err + } + + duration := trainingProbeDurationSeconds(videoPath) + second := trainingRandomSecond(duration) + + root, err := trainingRootDir() + if err != nil { + return nil, err + } + + if err := trainingEnsureDetectorDirs(root); err != nil { + return nil, err + } + + if err := os.MkdirAll(filepath.Join(root, "frames"), 0755); err != nil { + return nil, err + } + if err := os.MkdirAll(filepath.Join(root, "samples"), 0755); err != nil { + return nil, err + } + + id := trainingMakeSampleID(videoPath, second) + framePath := filepath.Join(root, "frames", id+".jpg") + + if err := trainingExtractFrame(videoPath, framePath, second); err != nil { + // Fallback auf Sekunde 0, falls random second nicht funktioniert. + second = 0 + id = trainingMakeSampleID(videoPath, second) + framePath = filepath.Join(root, "frames", id+".jpg") + + if err2 := trainingExtractFrame(videoPath, framePath, second); err2 != nil { + return nil, fmt.Errorf("frame extraction failed: %v / fallback: %w", err, err2) + } + } + + prediction := trainingPredictFrame(framePath) + + var sourceSizeBytes int64 + if st, err := os.Stat(videoPath); err == nil && st != nil && !st.IsDir() { + sourceSizeBytes = st.Size() + } + + sample := &TrainingSample{ + SampleID: id, + FrameURL: "/api/training/frame?id=" + id, + SourceFile: filepath.Base(videoPath), + SourcePath: videoPath, + SourceSizeBytes: sourceSizeBytes, + Second: second, + CreatedAt: time.Now().UTC().Format(time.RFC3339), + Prediction: prediction, + } + + if err := trainingWriteSample(root, sample); err != nil { + return nil, err + } + + return sample, nil +} + +func trainingPickRandomVideo(doneDir string) (string, error) { + extOK := map[string]bool{ + ".mp4": true, + ".mkv": true, + ".webm": true, + ".mov": true, + ".avi": true, + } + + var files []string + + err := filepath.WalkDir(doneDir, func(path string, d os.DirEntry, err error) error { + if err != nil { + return nil + } + if d.IsDir() { + name := strings.ToLower(d.Name()) + if name == ".trash" || name == "training" || name == "generated" { + return filepath.SkipDir + } + return nil + } + + ext := strings.ToLower(filepath.Ext(path)) + if extOK[ext] { + files = append(files, path) + } + return nil + }) + + if err != nil { + return "", err + } + + if len(files) == 0 { + return "", errors.New("keine Videos im doneDir gefunden") + } + + return files[rand.Intn(len(files))], nil +} + +func trainingExtractFrame(videoPath string, framePath string, second float64) error { + settings := getSettings() + + ffmpeg := strings.TrimSpace(settings.FFmpegPath) + if ffmpeg == "" { + ffmpeg = "ffmpeg" + } + + _ = os.Remove(framePath) + + ss := strconv.FormatFloat(math.Max(0, second), 'f', 3, 64) + + cmd := exec.Command( + ffmpeg, + "-hide_banner", + "-loglevel", "error", + "-ss", ss, + "-i", videoPath, + "-frames:v", "1", + "-q:v", "2", + "-y", + framePath, + ) + + cmd.SysProcAttr = &syscall.SysProcAttr{ + HideWindow: true, + CreationFlags: 0x08000000, + } + + out, err := cmd.CombinedOutput() + if err != nil { + return fmt.Errorf("%w: %s", err, strings.TrimSpace(string(out))) + } + + if st, err := os.Stat(framePath); err != nil || st.Size() == 0 { + return errors.New("ffmpeg erzeugte kein Frame") + } + + return nil +} + +func trainingPredictFrame(framePath string) TrainingPrediction { + root, err := trainingRootDir() + if err != nil { + return trainingEmptyPrediction("root_error") + } + + python := trainingPythonExe() + script := trainingScriptPath("predict_scene_model.py") + + cmd := exec.Command( + python, + script, + "--root", root, + "--image", framePath, + ) + + cmd.SysProcAttr = &syscall.SysProcAttr{ + HideWindow: true, + CreationFlags: 0x08000000, + } + + out, err := cmd.CombinedOutput() + if err != nil { + fmt.Println("⚠️ training predict failed:", err, strings.TrimSpace(string(out))) + return trainingEmptyPrediction("predict_failed") + } + + var pred TrainingPrediction + if err := json.Unmarshal(out, &pred); err != nil { + fmt.Println("⚠️ training predict json failed:", err, strings.TrimSpace(string(out))) + return trainingEmptyPrediction("predict_json_failed") + } + + if pred.SexPosition == "" { + pred.SexPosition = "unknown" + } + if pred.BodyPartsPresent == nil { + pred.BodyPartsPresent = []TrainingScoredLabel{} + } + if pred.ObjectsPresent == nil { + pred.ObjectsPresent = []TrainingScoredLabel{} + } + if pred.ClothingPresent == nil { + pred.ClothingPresent = []TrainingScoredLabel{} + } + if pred.Boxes == nil { + pred.Boxes = []TrainingBox{} + } + + return pred +} + +func trainingWriteDetectorSample(root string, sample *TrainingSample, boxes []TrainingBox) error { + if sample == nil { + return errors.New("sample missing") + } + + classMap, err := trainingDetectorClassMap() + if err != nil { + return err + } + + srcFrame := filepath.Join(root, "frames", sample.SampleID+".jpg") + if _, err := os.Stat(srcFrame); err != nil { + return fmt.Errorf("frame missing: %w", err) + } + + // Erstmal alles in train schreiben. + // Später kannst du 80/20 train/val splitten. + split := trainingStableSplit(sample.SampleID) + + imgDir := filepath.Join(root, "detector", "dataset", "images", split) + lblDir := filepath.Join(root, "detector", "dataset", "labels", split) + + if err := os.MkdirAll(imgDir, 0755); err != nil { + return err + } + if err := os.MkdirAll(lblDir, 0755); err != nil { + return err + } + + dstFrame := filepath.Join(imgDir, sample.SampleID+".jpg") + if err := copyFile(srcFrame, dstFrame); err != nil { + return err + } + + var lines []string + + for _, box := range boxes { + label := strings.TrimSpace(box.Label) + if label == "" { + continue + } + + classID, ok := classMap[label] + if !ok { + continue + } + + x := clamp01(box.X) + y := clamp01(box.Y) + w := clamp01(box.W) + h := clamp01(box.H) + + if w <= 0 || h <= 0 { + continue + } + + // Erwartung: box.X/Y/W/H sind bereits YOLO-normalisiert: + // x_center, y_center, width, height. + lines = append(lines, fmt.Sprintf( + "%d %.6f %.6f %.6f %.6f", + classID, + x, + y, + w, + h, + )) + } + + if len(lines) == 0 { + return errors.New("no valid detector boxes") + } + + labelPath := filepath.Join(lblDir, sample.SampleID+".txt") + return os.WriteFile(labelPath, []byte(strings.Join(lines, "\n")+"\n"), 0644) +} + +func trainingStableSplit(sampleID string) string { + sum := sha1.Sum([]byte(sampleID)) + if int(sum[0])%5 == 0 { + return "val" + } + return "train" +} + +func copyFile(src string, dst string) error { + b, err := os.ReadFile(src) + if err != nil { + return err + } + return os.WriteFile(dst, b, 0644) +} + +func trainingEmptyPrediction(source string) TrainingPrediction { + return TrainingPrediction{ + ModelAvailable: false, + Source: source, + PeopleCount: 0, + MaleCount: 0, + FemaleCount: 0, + UnknownCount: 0, + SexPosition: "unknown", + SexPositionScore: 0, + BodyPartsPresent: []TrainingScoredLabel{}, + ObjectsPresent: []TrainingScoredLabel{}, + ClothingPresent: []TrainingScoredLabel{}, + Boxes: []TrainingBox{}, + } +} + +func trainingPythonExe() string { + v := strings.TrimSpace(os.Getenv("TRAINING_PYTHON")) + if v != "" { + return v + } + return "python" +} + +func trainingProjectRoot() string { + wd, err := os.Getwd() + if err != nil { + return "." + } + + // Fall A: Prozess läuft im Projekt-Root: + // ./backend/ml/predict_scene_model.py existiert + if _, err := os.Stat(filepath.Join(wd, "backend", "ml", "predict_scene_model.py")); err == nil { + return wd + } + + // Fall B: Prozess läuft direkt in /backend: + // ./ml/predict_scene_model.py existiert + if _, err := os.Stat(filepath.Join(wd, "ml", "predict_scene_model.py")); err == nil { + return filepath.Dir(wd) + } + + // Fall C: Fallback: Parent testen + parent := filepath.Dir(wd) + if _, err := os.Stat(filepath.Join(parent, "backend", "ml", "predict_scene_model.py")); err == nil { + return parent + } + + return wd +} + +func trainingScriptPath(name string) string { + // 1) Eingebettete Scripts bevorzugen. + if dir, err := trainingEmbeddedMLDir(); err == nil { + p := filepath.Join(dir, name) + if _, err := os.Stat(p); err == nil { + return p + } + } + + // 2) Dev-Fallback: Projektstruktur. + root := trainingProjectRoot() + + p := filepath.Join(root, "backend", "ml", name) + if _, err := os.Stat(p); err == nil { + return p + } + + p = filepath.Join("ml", name) + if _, err := os.Stat(p); err == nil { + return p + } + + return filepath.Join(root, "backend", "ml", name) +} + +func trainingRootDir() (string, error) { + root := filepath.Join("generated", "training") + + if err := os.MkdirAll(root, 0755); err != nil { + return "", err + } + + return root, nil +} + +func trainingWriteSample(root string, sample *TrainingSample) error { + path := filepath.Join(root, "samples", sample.SampleID+".json") + b, err := json.MarshalIndent(sample, "", " ") + if err != nil { + return err + } + return os.WriteFile(path, b, 0644) +} + +func trainingReadSample(root string, sampleID string) (*TrainingSample, error) { + if strings.Contains(sampleID, "/") || strings.Contains(sampleID, "\\") { + return nil, errors.New("invalid sample id") + } + + path := filepath.Join(root, "samples", sampleID+".json") + + b, err := os.ReadFile(path) + if err != nil { + return nil, err + } + + var sample TrainingSample + if err := json.Unmarshal(b, &sample); err != nil { + return nil, err + } + + return &sample, nil +} + +func trainingAppendAnnotation(root string, annotation TrainingAnnotation) error { + path := filepath.Join(root, "feedback.jsonl") + + f, err := os.OpenFile(path, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0644) + if err != nil { + return err + } + defer f.Close() + + b, err := json.Marshal(annotation) + if err != nil { + return err + } + + if _, err := f.Write(append(b, '\n')); err != nil { + return err + } + + return nil +} + +func trainingCountAnnotations(path string) (int, error) { + b, err := os.ReadFile(path) + if err != nil { + return 0, err + } + + text := strings.TrimSpace(string(b)) + if text == "" { + return 0, nil + } + + return len(strings.Split(text, "\n")), nil +} + +func trainingProbeDurationSeconds(videoPath string) float64 { + settings := getSettings() + + ffmpeg := strings.TrimSpace(settings.FFmpegPath) + ffprobe := "ffprobe" + + if ffmpeg != "" { + dir := filepath.Dir(ffmpeg) + base := filepath.Base(ffmpeg) + if strings.Contains(strings.ToLower(base), "ffmpeg") { + ffprobeBase := strings.Replace(base, "ffmpeg", "ffprobe", 1) + ffprobe = filepath.Join(dir, ffprobeBase) + } + } + + cmd := exec.Command( + ffprobe, + "-v", "error", + "-show_entries", "format=duration", + "-of", "default=noprint_wrappers=1:nokey=1", + videoPath, + ) + + cmd.SysProcAttr = &syscall.SysProcAttr{ + HideWindow: true, + CreationFlags: 0x08000000, + } + + out, err := cmd.Output() + if err != nil { + return 0 + } + + v, err := strconv.ParseFloat(strings.TrimSpace(string(out)), 64) + if err != nil || math.IsNaN(v) || math.IsInf(v, 0) || v <= 0 { + return 0 + } + + return v +} + +func trainingRandomSecond(duration float64) float64 { + if duration <= 2 { + return 0 + } + + minSec := 1.0 + maxSec := math.Max(minSec, duration-1) + + return minSec + rand.Float64()*(maxSec-minSec) +} + +func trainingMakeSampleID(videoPath string, second float64) string { + h := sha1.New() + _, _ = h.Write([]byte(videoPath)) + _, _ = h.Write([]byte("|")) + _, _ = h.Write([]byte(strconv.FormatFloat(second, 'f', 3, 64))) + _, _ = h.Write([]byte("|")) + _, _ = h.Write([]byte(strconv.FormatInt(time.Now().UnixNano(), 10))) + return hex.EncodeToString(h.Sum(nil))[:20] +} + +func trainingWriteJSON(w http.ResponseWriter, status int, v any) { + w.Header().Set("Content-Type", "application/json") + w.Header().Set("Cache-Control", "no-store") + w.WriteHeader(status) + _ = json.NewEncoder(w).Encode(v) +} + +func trainingWriteError(w http.ResponseWriter, status int, msg string) { + trainingWriteJSON(w, status, map[string]any{ + "ok": false, + "error": msg, + }) +} diff --git a/backend/training_label_loader.go b/backend/training_label_loader.go new file mode 100644 index 0000000..afa0b6f --- /dev/null +++ b/backend/training_label_loader.go @@ -0,0 +1,149 @@ +// backend\training_label_loader.go + +package main + +import ( + "encoding/json" + "fmt" + "os" + "path/filepath" + "strings" +) + +type TrainingGroupedLabels struct { + 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.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") + } + + return grouped, nil +} + +func trainingDetectorLabels() ([]string, error) { + grouped, err := trainingGroupedLabels() + if err != nil { + return nil, err + } + + labels := []string{} + 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{ + SexPositions: grouped.SexPositions, + BodyParts: grouped.BodyParts, + Objects: grouped.Objects, + Clothing: grouped.Clothing, + } +} diff --git a/backend/yolo11n.pt b/backend/yolo11n.pt new file mode 100644 index 0000000..45b273b Binary files /dev/null and b/backend/yolo11n.pt differ diff --git a/frontend/src/App.tsx b/frontend/src/App.tsx index 33ce2b7..ca31ab4 100644 --- a/frontend/src/App.tsx +++ b/frontend/src/App.tsx @@ -23,6 +23,7 @@ import { UserGroupIcon, Squares2X2Icon, Cog6ToothIcon, + BeakerIcon, } from '@heroicons/react/24/solid' import PerformanceMonitor from './components/ui/PerformanceMonitor' import { useNotify } from './components/ui/notify' @@ -31,6 +32,7 @@ import LoginPage from './components/ui/LoginPage' import VideoSplitModal from './components/ui/VideoSplitModal' import TextInput from './components/ui/TextInput' import LastUpdatedText from './components/ui/LastUpdatedText' +import TrainingTab from './components/ui/TrainingTab' const COOKIE_STORAGE_KEY = 'record_cookies' @@ -2577,6 +2579,11 @@ export default function App() { label: 'Kategorien', icon: Squares2X2Icon, }, + { + id: 'training', + label: 'Training', + icon: BeakerIcon, + }, { id: 'settings', label: 'Einstellungen', @@ -2954,6 +2961,7 @@ export default function App() { if (d.tab === 'finished') setSelectedTab('finished') if (d.tab === 'categories') setSelectedTab('categories') if (d.tab === 'models') setSelectedTab('models') + if (d.tab === 'training') setSelectedTab('training') if (d.tab === 'running') setSelectedTab('running') if (d.tab === 'settings') setSelectedTab('settings') } @@ -3970,7 +3978,13 @@ export default function App() { {selectedTab === 'models' ? ( ) : null} + + {selectedTab === 'training' ? ( + + ) : null} + {selectedTab === 'categories' ? : null} + {selectedTab === 'settings' ? ( { return out } -function formatDuration(ms: number): string { - if (!Number.isFinite(ms) || ms <= 0) return '—' - const totalSec = Math.floor(ms / 1000) - const h = Math.floor(totalSec / 3600) - const m = Math.floor((totalSec % 3600) / 60) - const s = totalSec % 60 - if (h > 0) return `${h}h ${m}m` - if (m > 0) return `${m}m ${s}s` - return `${s}s` -} - -function formatBytes(bytes?: number | null): string { - if (typeof bytes !== 'number' || !Number.isFinite(bytes) || bytes <= 0) return '—' - const units = ['B', 'KB', 'MB', 'GB', 'TB'] - let v = bytes - let i = 0 - while (v >= 1024 && i < units.length - 1) { - v /= 1024 - i++ - } - const digits = i === 0 ? 0 : v >= 100 ? 0 : v >= 10 ? 1 : 2 - return `${v.toFixed(digits)} ${units[i]}` -} - -function formatDateTime(v?: string | number | Date | null): string { - if (!v) return '—' - const d = v instanceof Date ? v : new Date(v) - const t = d.getTime() - if (!Number.isFinite(t)) return '—' - return d.toLocaleString(undefined, { - year: 'numeric', - month: '2-digit', - day: '2-digit', - hour: '2-digit', - minute: '2-digit', - }) -} - const pickNum = (...vals: any[]): number | null => { for (const v of vals) { const n = typeof v === 'string' ? Number(v) : v diff --git a/frontend/src/components/ui/TrainingTab.tsx b/frontend/src/components/ui/TrainingTab.tsx new file mode 100644 index 0000000..b3faf03 --- /dev/null +++ b/frontend/src/components/ui/TrainingTab.tsx @@ -0,0 +1,805 @@ +// frontend/src/components/ui/TrainingTab.tsx +'use client' + +import { useCallback, useEffect, useMemo, useState } from 'react' +import Button from './Button' +import LoadingSpinner from './LoadingSpinner' +import { formatBytes, formatDuration } from './formatters' + +type ScoredLabel = { + label: string + score: number +} + +type TrainingStatus = { + feedbackCount: number + requiredCount: number + canTrain: boolean +} + +type TrainingPrediction = { + modelAvailable: boolean + source?: string + + peopleCount: number + maleCount: number + femaleCount: number + unknownCount: number + sexPosition: string + sexPositionScore: number + bodyPartsPresent: ScoredLabel[] + objectsPresent: ScoredLabel[] + clothingPresent: ScoredLabel[] + boxes?: TrainingBox[] +} + +type TrainingSample = { + sampleId: string + frameUrl: string + sourceFile: string + sourcePath?: string + sourceSizeBytes?: number + second: number + createdAt: string + prediction: TrainingPrediction +} + +type TrainingLabels = { + sexPositions: string[] + bodyParts: string[] + objects: string[] + clothing: string[] +} + +type CorrectionState = { + peopleCount: number + maleCount: number + femaleCount: number + unknownCount: number + sexPosition: string + bodyPartsPresent: string[] + objectsPresent: string[] + clothingPresent: string[] +} + +type TrainingBox = { + label: string + score?: number + x: number + y: number + w: number + h: number +} + +const emptyLabels: TrainingLabels = { + sexPositions: ['unknown'], + bodyParts: [], + objects: [], + clothing: [], +} + +function percent(v: number) { + if (!Number.isFinite(v)) return '—' + return `${Math.round(v * 100)}%` +} + +function scoredLabelsText(items?: ScoredLabel[] | null) { + const list = Array.isArray(items) ? items : [] + + if (list.length === 0) return '—' + + return list + .map((x) => { + const label = String(x?.label ?? '').trim() + if (!label) return null + + return `${label} (${percent(Number(x?.score ?? 0))})` + }) + .filter(Boolean) + .join(', ') +} + +function detectionBoxes(sample: TrainingSample | null) { + const boxes = sample?.prediction.boxes ?? [] + + return boxes.filter((box) => { + const label = String(box.label || '').trim() + if (!label) return false + + const isBodyPart = sample?.prediction.bodyPartsPresent?.some((x) => x.label === label) + const isObject = sample?.prediction.objectsPresent?.some((x) => x.label === label) + + return isBodyPart || isObject + }) +} + +function clampPercent(v: number) { + if (!Number.isFinite(v)) return 0 + return Math.max(0, Math.min(100, v)) +} + +function toggleArrayValue(arr: string[], value: string) { + return arr.includes(value) + ? arr.filter((x) => x !== value) + : [...arr, value] +} + +function safeCount(value: unknown) { + const n = Number(value) + return Number.isFinite(n) && n > 0 ? Math.floor(n) : 0 +} + +function totalPeopleFromGenderCounts(state: Pick) { + return safeCount(state.maleCount) + safeCount(state.femaleCount) +} + +function predictionToCorrection(sample: TrainingSample | null): CorrectionState { + const p = sample?.prediction + + const maleCount = safeCount(p?.maleCount) + const femaleCount = safeCount(p?.femaleCount) + + return { + peopleCount: maleCount + femaleCount, + maleCount, + femaleCount, + unknownCount: 0, + sexPosition: p?.sexPosition || 'unknown', + bodyPartsPresent: (p?.bodyPartsPresent ?? []).map((x) => x.label), + objectsPresent: (p?.objectsPresent ?? []).map((x) => x.label), + clothingPresent: (p?.clothingPresent ?? []).map((x) => x.label), + } +} + +function sortLabelList(values?: string[], opts?: { keepUnknownFirst?: boolean }) { + const list = [...(values ?? [])].sort((a, b) => + a.localeCompare(b, undefined, { sensitivity: 'base' }) + ) + + if (!opts?.keepUnknownFirst) return list + + return [ + ...list.filter((x) => x === 'unknown'), + ...list.filter((x) => x !== 'unknown'), + ] +} + +function sortTrainingLabels(input: Partial | null | undefined): TrainingLabels { + return { + sexPositions: sortLabelList(input?.sexPositions, { keepUnknownFirst: true }), + bodyParts: sortLabelList(input?.bodyParts), + objects: sortLabelList(input?.objects), + clothing: sortLabelList(input?.clothing), + } +} + +export default function TrainingTab() { + const [labels, setLabels] = useState(emptyLabels) + const [sample, setSample] = useState(null) + const [correction, setCorrection] = useState(() => predictionToCorrection(null)) + const [loading, setLoading] = useState(false) + const [saving, setSaving] = useState(false) + const [training, setTraining] = useState(false) + const [trainingStatus, setTrainingStatus] = useState(null) + const [deletingTrainingData, setDeletingTrainingData] = useState(false) + const [error, setError] = useState(null) + const [message, setMessage] = useState(null) + + const computedPeopleCount = useMemo( + () => totalPeopleFromGenderCounts(correction), + [correction.maleCount, correction.femaleCount] + ) + + const imageSrc = useMemo(() => { + if (!sample?.frameUrl) return '' + return `${sample.frameUrl}&t=${encodeURIComponent(sample.sampleId)}` + }, [sample]) + + const boxes = useMemo(() => detectionBoxes(sample), [sample]) + + const canStartTraining = Boolean(trainingStatus?.canTrain) + const feedbackCount = trainingStatus?.feedbackCount ?? 0 + const requiredCount = trainingStatus?.requiredCount ?? 5 + + const loadLabels = useCallback(async () => { + const res = await fetch('/api/training/labels', { cache: 'no-store' }) + if (!res.ok) return + + const data = await res.json().catch(() => null) + if (data) setLabels(sortTrainingLabels(data)) + }, []) + + const loadNext = useCallback(async (opts?: { forceNew?: boolean }) => { + setLoading(true) + setError(null) + setMessage(null) + + try { + const url = opts?.forceNew + ? '/api/training/next?force=1' + : '/api/training/next' + + const res = await fetch(url, { cache: 'no-store' }) + const data = await res.json().catch(() => null) + + if (!res.ok) { + throw new Error(data?.error || `HTTP ${res.status}`) + } + + setSample(data) + setCorrection(predictionToCorrection(data)) + } catch (e) { + setError(e instanceof Error ? e.message : String(e)) + } finally { + setLoading(false) + } + }, []) + + const loadTrainingStatus = useCallback(async () => { + const res = await fetch('/api/training/status', { cache: 'no-store' }) + const data = await res.json().catch(() => null) + + if (!res.ok || !data) return + + setTrainingStatus({ + feedbackCount: Number(data.feedbackCount ?? 0), + requiredCount: Number(data.requiredCount ?? 5), + canTrain: Boolean(data.canTrain), + }) + }, []) + + useEffect(() => { + void loadLabels() + void loadNext() + void loadTrainingStatus() + }, [loadLabels, loadNext, loadTrainingStatus]) + + const saveFeedback = useCallback( + async (accepted: boolean) => { + if (!sample) return + + setSaving(true) + setError(null) + setMessage(null) + + try { + const correctionPayload: CorrectionState = { + ...correction, + peopleCount: totalPeopleFromGenderCounts(correction), + unknownCount: 0, + } + + const payload = { + sampleId: sample.sampleId, + accepted, + correction: accepted ? undefined : correctionPayload, + } + + const res = await fetch('/api/training/feedback', { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify(payload), + }) + + const data = await res.json().catch(() => null) + + if (!res.ok) { + throw new Error(data?.error || `HTTP ${res.status}`) + } + + await loadTrainingStatus() + await loadNext() + } catch (e) { + setError(e instanceof Error ? e.message : String(e)) + } finally { + setSaving(false) + } + }, + [sample, correction, loadNext, loadTrainingStatus] + ) + + const startTraining = useCallback(async () => { + setTraining(true) + setError(null) + setMessage(null) + + try { + const res = await fetch('/api/training/train', { + method: 'POST', + }) + + const data = await res.json().catch(() => null) + + if (!res.ok) { + throw new Error(data?.error || `HTTP ${res.status}`) + } + + setMessage(data?.message || 'Training gestartet.') + } catch (e) { + setError(e instanceof Error ? e.message : String(e)) + } finally { + setTraining(false) + } + }, []) + + const deleteAllTrainingData = useCallback(async () => { + const confirmed = window.confirm( + 'Wirklich alle Trainingsdaten löschen? Das entfernt Feedback, Frames, Samples und Detector-Daten. Diese Aktion kann nicht rückgängig gemacht werden.' + ) + + if (!confirmed) return + + setDeletingTrainingData(true) + setError(null) + setMessage(null) + + try { + const res = await fetch('/api/training/delete-all', { + method: 'DELETE', + }) + + const data = await res.json().catch(() => null) + + if (!res.ok) { + throw new Error(data?.error || `HTTP ${res.status}`) + } + + setSample(null) + setCorrection(predictionToCorrection(null)) + setTrainingStatus({ + feedbackCount: 0, + requiredCount, + canTrain: false, + }) + + setMessage(data?.message || 'Alle Trainingsdaten wurden gelöscht.') + + await loadTrainingStatus() + await loadNext({ forceNew: true }) + } catch (e) { + setError(e instanceof Error ? e.message : String(e)) + } finally { + setDeletingTrainingData(false) + } + }, [loadNext, loadTrainingStatus, requiredCount]) + + return ( +
+ {/* Sidebar links */} + + + {/* Mitte */} +
+
+ {loading ? ( + + ) : imageSrc ? ( +
+ Training Frame + + {boxes.length > 0 ? ( +
+ {boxes.map((box, index) => { + const left = clampPercent(box.x * 100) + const top = clampPercent(box.y * 100) + const width = clampPercent(box.w * 100) + const height = clampPercent(box.h * 100) + + return ( +
+
+ {box.label} {typeof box.score === 'number' ? percent(box.score) : ''} +
+
+ ) + })} +
+ ) : null} +
+ ) : ( +
Kein Frame geladen
+ )} +
+ +
+ + + + + +
+ +
+ Prüfe das Bild. Wenn die Erkennung stimmt: „Passt so“. Wenn nicht: rechts korrigieren und speichern. +
+
+ + {/* Rechte Details/Korrektur */} + +
+ ) +} \ No newline at end of file diff --git a/frontend/src/components/ui/VideoSplitModal.tsx b/frontend/src/components/ui/VideoSplitModal.tsx index 22436e5..4a9d720 100644 --- a/frontend/src/components/ui/VideoSplitModal.tsx +++ b/frontend/src/components/ui/VideoSplitModal.tsx @@ -12,6 +12,7 @@ import { getSegmentLabelText, getHighestPrioritySegmentLabelItem, } from './Icons' +import { formatDateTime } from './formatters' type Segment = { index: number @@ -364,20 +365,6 @@ function formatBytes(bytes?: number | null): string { return `${v.toFixed(digits)} ${units[i]}` } -function formatDateTime(v?: string | number | Date | null): string { - if (!v) return '—' - const d = v instanceof Date ? v : new Date(v) - const t = d.getTime() - if (!Number.isFinite(t)) return '—' - return d.toLocaleString(undefined, { - year: 'numeric', - month: '2-digit', - day: '2-digit', - hour: '2-digit', - minute: '2-digit', - }) -} - function pickNum(...vals: any[]): number | null { for (const v of vals) { const n = typeof v === 'string' ? Number(v) : v diff --git a/frontend/src/components/ui/formatters.ts b/frontend/src/components/ui/formatters.ts index 304b42f..834b5d4 100644 --- a/frontend/src/components/ui/formatters.ts +++ b/frontend/src/components/ui/formatters.ts @@ -35,7 +35,6 @@ export function formatResolution(r?: { w: number; h: number } | null): string { return `${p}p` } - export function formatDuration(ms: number): string { if (!Number.isFinite(ms) || ms <= 0) return '—' const totalSec = Math.floor(ms / 1000) @@ -58,4 +57,19 @@ export function formatBytes(bytes?: number | null): string { } const digits = i === 0 ? 0 : v >= 100 ? 0 : v >= 10 ? 1 : 2 return `${v.toFixed(digits)} ${units[i]}` +} + + +export function formatDateTime(v?: string | number | Date | null): string { + if (!v) return '—' + const d = v instanceof Date ? v : new Date(v) + const t = d.getTime() + if (!Number.isFinite(t)) return '—' + return d.toLocaleString(undefined, { + year: 'numeric', + month: '2-digit', + day: '2-digit', + hour: '2-digit', + minute: '2-digit', + }) } \ No newline at end of file