nsfwapp/backend/ml/predict_scene_model.py

# 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()