nsfwapp/backend/ai_server.py

# backend\ai_server.py

import os
from typing import List, Optional

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


BODY_LABELS = {
    "anus",
    "ass",
    "breasts",
    "penis",
    "tongue",
    "pussy",
}

OBJECT_LABELS = {
    "blindfold",
    "buttplug",
    "collar",
    "dildo",
    "handcuffs",
    "shower",
    "strapon",
    "towel",
    "vibrator",
}

CLOTHING_LABELS = {
    "bikini",
    "bra",
    "dress",
    "heels",
    "hotpants",
    "lingerie",
    "panties",
    "skirt",
    "stockings",
    "croptop",
}

POSITION_LABELS = {
    "missionary",
    "doggy",
    "cowgirl",
    "reverse_cowgirl",
    "cunnilingus",
    "prone_bone",
    "standing",
    "standing_doggy",
    "spooning",
    "sitting",
    "facesitting",
    "handjob",
    "blowjob",
    "toy_play",
    "fingering",
    "69",
    "other",
}


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


app = FastAPI()

from pathlib import Path

BASE_DIR = Path(__file__).resolve().parent
DEFAULT_MODEL_PATH = BASE_DIR / "generated" / "training" / "detector" / "model" / "best.pt"

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

    default_path = DEFAULT_MODEL_PATH
    if default_path.exists():
        return str(default_path)

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


_MODEL_PATH = resolve_model_path()
_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 = YOLO(_MODEL_PATH)


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 = []
    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))

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

            if label in BODY_LABELS:
                best_score(body_parts, label, score)

            if label in OBJECT_LABELS:
                best_score(objects, label, score)

            if label in CLOTHING_LABELS:
                best_score(clothing, label, score)

            if label in POSITION_LABELS and score > sex_position_score:
                sex_position = label
                sex_position_score = score

    people_count = sum(
        1 for box in boxes_out
        if box["label"] in {"person", "person_male", "person_female", "person_unknown"}
    )
    male_count = sum(1 for box in boxes_out if box["label"] in {"person_male", "male_person"})
    female_count = sum(1 for box in boxes_out if box["label"] in {"person_female", "female_person"})
    unknown_count = max(0, people_count - male_count - female_count)

    return {
        "modelAvailable": True,
        "source": f"yolo-server:{Path(_MODEL_PATH).name}",
        "peopleCount": people_count,
        "maleCount": male_count,
        "femaleCount": female_count,
        "unknownCount": unknown_count,
        "sexPosition": sex_position,
        "sexPositionScore": sex_position_score,
        "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",
        }

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

    try:
        results = 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():
    names = getattr(model, "names", {}) or {}

    return {
        "ok": True,
        "model": _MODEL_PATH,
        "classCount": len(names),
        "classes": list(names.values())[:80] if isinstance(names, dict) else names,
    }