nsfwapp/backend/ai_server.py

# backend/ai_server.py

import json
import os
from pathlib import Path
from typing import List, Optional

from fastapi import FastAPI
from pydantic import BaseModel
from ultralytics import YOLO


BASE_DIR = Path(__file__).resolve().parent


def existing_file(path: Path) -> Optional[Path]:
    try:
        if path.exists() and path.is_file() and path.stat().st_size > 0:
            return path
    except OSError:
        pass
    return None


def resolve_training_root() -> Path:
    env_root = os.environ.get("TRAINING_ROOT", "").strip()
    if env_root:
        root = Path(env_root).expanduser().resolve()
        root.mkdir(parents=True, exist_ok=True)
        return root

    candidates = [
        # Wenn ai_server.py aus backend/ läuft:
        BASE_DIR / "generated" / "training",

        # Wenn ai_server.py aus backend/ml/ laufen würde:
        BASE_DIR.parent / "generated" / "training",

        # Wenn ai_server.py embedded aus Temp läuft, aber backendRoot als cwd gesetzt wurde:
        Path.cwd() / "generated" / "training",

        # Wenn Working Directory Projektroot ist:
        Path.cwd() / "backend" / "generated" / "training",
    ]

    for root in candidates:
        if (
            existing_file(root / "detection_labels.json")
            or existing_file(root / "detector" / "model" / "best.pt")
        ):
            root.mkdir(parents=True, exist_ok=True)
            return root.resolve()

    # Fallback: Server soll trotzdem starten.
    root = (Path.cwd() / "generated" / "training").resolve()
    root.mkdir(parents=True, exist_ok=True)
    return root


TRAINING_ROOT = resolve_training_root()
DEFAULT_MODEL_PATH = TRAINING_ROOT / "detector" / "model" / "best.pt"


def resolve_detection_labels_path() -> Path:
    env_path = os.environ.get("DETECTION_LABELS_PATH", "").strip()
    if env_path:
        p = Path(env_path).expanduser().resolve()
        if existing_file(p):
            return p
        raise RuntimeError(f"DETECTION_LABELS_PATH not found: {p}")

    candidates = [
        TRAINING_ROOT / "detection_labels.json",

        # Wenn ai_server.py direkt neben detection_labels.json embedded liegt:
        BASE_DIR / "detection_labels.json",

        # Dev-Fallbacks:
        BASE_DIR / "ml" / "detection_labels.json",
        BASE_DIR.parent / "ml" / "detection_labels.json",
        Path.cwd() / "ml" / "detection_labels.json",
        Path.cwd() / "backend" / "ml" / "detection_labels.json",
    ]

    for p in candidates:
        if existing_file(p):
            return p.resolve()

    raise RuntimeError(
        "detection_labels.json not found. Checked: "
        + ", ".join(str(p) for p in candidates)
    )


def resolve_model_path() -> str:
    env_path = os.environ.get("YOLO_MODEL", "").strip()
    if env_path:
        p = Path(env_path).expanduser().resolve()
        if existing_file(p):
            return str(p)
        raise RuntimeError(f"YOLO_MODEL not found: {p}")

    if existing_file(DEFAULT_MODEL_PATH):
        return str(DEFAULT_MODEL_PATH)

    raise RuntimeError(f"YOLO model not found: {DEFAULT_MODEL_PATH}")


# Server darf auch ohne Labels/Model starten.
DETECTION_LABELS_PATH: Optional[Path] = None

LABEL_GROUPS = {
    "people": set(),
    "sexPositions": {"unknown"},
    "bodyParts": set(),
    "objects": set(),
    "clothing": set(),
}

BODY_LABELS = LABEL_GROUPS["bodyParts"]
OBJECT_LABELS = LABEL_GROUPS["objects"]
CLOTHING_LABELS = LABEL_GROUPS["clothing"]
POSITION_LABELS = set()

PERSON_LABELS = {
    "person_male",
    "person_female",
}

_MODEL_PATH = ""
_MODEL_ERROR = ""
_LABEL_ERROR = ""

_DEVICE = os.environ.get("YOLO_DEVICE", "")
_CONF = float(os.environ.get("YOLO_CONF", "0.25"))
_BATCH = int(os.environ.get("YOLO_BATCH", "16"))
_IMGSZ = int(os.environ.get("YOLO_IMGSZ", "640"))
_HALF = os.environ.get("YOLO_HALF", "0").lower() in {"1", "true", "yes", "on"}

model = None

app = FastAPI()


class PredictBatchRequest(BaseModel):
    paths: List[str]
    detectorOnly: bool = False
    imageSize: int = 640
    model: Optional[str] = None


def empty_prediction(source: str = "model_missing") -> dict:
    return {
        "modelAvailable": False,
        "source": source,
        "sexPosition": "unknown",
        "sexPositionScore": 0.0,
        "peoplePresent": [],
        "bodyPartsPresent": [],
        "objectsPresent": [],
        "clothingPresent": [],
        "boxes": [],
    }


def load_label_groups_safe() -> None:
    global DETECTION_LABELS_PATH
    global LABEL_GROUPS
    global BODY_LABELS
    global OBJECT_LABELS
    global CLOTHING_LABELS
    global POSITION_LABELS
    global PERSON_LABELS
    global _LABEL_ERROR

    try:
        path = resolve_detection_labels_path()
        DETECTION_LABELS_PATH = path

        with path.open("r", encoding="utf-8") as f:
            data = json.load(f)

        LABEL_GROUPS = {
            "people": set(
                str(x).strip().lower()
                for x in data.get("people", [])
                if str(x).strip()
            ),
            "sexPositions": set(
                str(x).strip().lower()
                for x in data.get("sexPositions", [])
                if str(x).strip()
            ),
            "bodyParts": set(
                str(x).strip().lower()
                for x in data.get("bodyParts", [])
                if str(x).strip()
            ),
            "objects": set(
                str(x).strip().lower()
                for x in data.get("objects", [])
                if str(x).strip()
            ),
            "clothing": set(
                str(x).strip().lower()
                for x in data.get("clothing", [])
                if str(x).strip()
            ),
        }

        if not LABEL_GROUPS["sexPositions"]:
            LABEL_GROUPS["sexPositions"] = {"unknown"}

        _LABEL_ERROR = ""

    except Exception as exc:
        DETECTION_LABELS_PATH = None
        _LABEL_ERROR = str(exc)

        LABEL_GROUPS = {
            "people": set(),
            "sexPositions": {"unknown"},
            "bodyParts": set(),
            "objects": set(),
            "clothing": set(),
        }

    BODY_LABELS = LABEL_GROUPS["bodyParts"]
    OBJECT_LABELS = LABEL_GROUPS["objects"]
    CLOTHING_LABELS = LABEL_GROUPS["clothing"]

    POSITION_LABELS = {
        label for label in LABEL_GROUPS["sexPositions"]
        if label and label != "unknown"
    }

    PERSON_LABELS = {
        label for label in LABEL_GROUPS["people"]
        if label
    }


def get_model():
    global model
    global _MODEL_PATH
    global _MODEL_ERROR

    if model is not None:
        return model

    try:
        path = resolve_model_path()
        loaded = YOLO(path)

        model = loaded
        _MODEL_PATH = path
        _MODEL_ERROR = ""

        # Labels erst laden, wenn Inference wirklich gebraucht wird.
        load_label_groups_safe()

        return model

    except Exception as exc:
        model = None
        _MODEL_PATH = ""
        _MODEL_ERROR = str(exc)
        return None


def scored(label: str, score: float) -> dict:
    return {
        "label": label,
        "score": float(score),
    }


def best_score(items: list[dict], label: str, score: float) -> None:
    for item in items:
        if item["label"] == label:
            if score > item["score"]:
                item["score"] = float(score)
            return

    items.append(scored(label, score))


def prediction_from_result(result) -> dict:
    names = result.names or {}

    boxes_out = []
    people_present = []
    body_parts = []
    objects = []
    clothing = []

    sex_position = "unknown"
    sex_position_score = 0.0

    if result.boxes is not None:
        xywhn = result.boxes.xywhn.cpu().tolist()
        cls_values = result.boxes.cls.cpu().tolist()
        conf_values = result.boxes.conf.cpu().tolist()

        for box_xywhn, cls_id, conf in zip(xywhn, cls_values, conf_values):
            label = str(names.get(int(cls_id), int(cls_id))).strip().lower()
            score = float(conf)

            if not label:
                continue

            cx, cy, w, h = [float(v) for v in box_xywhn]

            x = max(0.0, min(1.0, cx - w / 2.0))
            y = max(0.0, min(1.0, cy - h / 2.0))
            w = max(0.0, min(1.0 - x, w))
            h = max(0.0, min(1.0 - y, h))

            is_person = label in PERSON_LABELS
            is_body = label in BODY_LABELS
            is_object = label in OBJECT_LABELS
            is_clothing = label in CLOTHING_LABELS
            is_position = label in POSITION_LABELS

            if is_position:
                if score > sex_position_score:
                    sex_position = label
                    sex_position_score = score
                continue

            if not (is_person or is_body or is_object or is_clothing):
                continue

            boxes_out.append({
                "label": label,
                "score": score,
                "x": x,
                "y": y,
                "w": w,
                "h": h,
            })

            if is_person:
                best_score(people_present, label, score)

            if is_body:
                best_score(body_parts, label, score)

            if is_object:
                best_score(objects, label, score)

            if is_clothing:
                best_score(clothing, label, score)

    return {
        "modelAvailable": True,
        "source": f"yolo-server:{Path(_MODEL_PATH).name}",
        "sexPosition": sex_position,
        "sexPositionScore": sex_position_score,
        "peoplePresent": people_present,
        "bodyPartsPresent": body_parts,
        "objectsPresent": objects,
        "clothingPresent": clothing,
        "boxes": boxes_out,
    }


@app.post("/predict-batch")
def predict_batch(req: PredictBatchRequest):
    paths = [str(path).strip() for path in req.paths if str(path).strip()]
    if not paths:
        return {
            "ok": False,
            "predictions": [],
            "error": "no paths supplied",
        }

    current_model = get_model()
    if current_model is None:
        return {
            "ok": True,
            "predictions": [empty_prediction("model_missing") for _ in paths],
            "error": _MODEL_ERROR or f"YOLO model not found: {DEFAULT_MODEL_PATH}",
        }

    if DETECTION_LABELS_PATH is None or _LABEL_ERROR:
        return {
            "ok": False,
            "predictions": [],
            "error": f"detection labels missing: {_LABEL_ERROR}",
        }

    imgsz = int(req.imageSize or _IMGSZ or 640)

    try:
        results = current_model.predict(
            source=paths,
            imgsz=imgsz,
            conf=_CONF,
            batch=_BATCH,
            device=_DEVICE or None,
            half=_HALF,
            verbose=False,
        )

        predictions = [prediction_from_result(result) for result in results]

        return {
            "ok": True,
            "predictions": predictions,
        }

    except Exception as exc:
        return {
            "ok": False,
            "predictions": [],
            "error": str(exc),
        }


@app.get("/health")
def health():
    current_model = get_model()
    names = getattr(current_model, "names", {}) or {} if current_model is not None else {}

    return {
        "ok": True,
        "ready": current_model is not None,
        "modelAvailable": current_model is not None,
        "model": _MODEL_PATH,
        "modelError": _MODEL_ERROR,
        "expectedModel": str(DEFAULT_MODEL_PATH),
        "trainingRoot": str(TRAINING_ROOT),
        "classCount": len(names),
        "classes": list(names.values())[:80] if isinstance(names, dict) else names,
        "labelConfig": str(DETECTION_LABELS_PATH) if DETECTION_LABELS_PATH else "",
        "labelError": _LABEL_ERROR,
    }