nsfwapp/backend/ml/predict_pose_model.py

# backend/ml/predict_pose_model.py

import argparse
import json
from pathlib import Path

from PIL import Image
from ultralytics import YOLO
import torch

KEYPOINT_NAMES = [
    "nose",
    "left_eye", "right_eye",
    "left_ear", "right_ear",
    "left_shoulder", "right_shoulder",
    "left_elbow", "right_elbow",
    "left_wrist", "right_wrist",
    "left_hip", "right_hip",
    "left_knee", "right_knee",
    "left_ankle", "right_ankle",
]

BASE_MODEL_NAME = "yolo26n-pose.pt"


def existing_file(path):
    try:
        p = Path(path).expanduser().resolve()
        if p.exists() and p.is_file() and p.stat().st_size > 0:
            return p
    except Exception:
        pass
    return None


def base_model_candidates(root):
    script_dir = Path(__file__).resolve().parent

    return [
        script_dir / BASE_MODEL_NAME,
        Path.cwd() / BASE_MODEL_NAME,
        root / BASE_MODEL_NAME,
        root.parent / BASE_MODEL_NAME,
        root.parent.parent / BASE_MODEL_NAME,
    ]


def resolve_model_path(root, requested):
    if requested:
        p = existing_file(requested)
        if p:
            return p, "yolo_pose_model"
        return Path(requested).expanduser(), "pose_missing"

    trained = root / "pose" / "model" / "best.pt"
    p = existing_file(trained)
    if p:
        return p, "yolo_pose"

    for candidate in base_model_candidates(root):
        p = existing_file(candidate)
        if p:
            return p, "yolo26_pose_base"

    return trained, "pose_missing"


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--root", required=True)
    parser.add_argument("--image", required=True)
    parser.add_argument("--model", default="")
    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, model_source = resolve_model_path(root, args.model)
    if not model_path.exists():
        print(json.dumps({
            "available": False,
            "source": model_source,
            "modelPath": str(model_path),
            "persons": [],
        }, ensure_ascii=False))
        return

    img = Image.open(image_path).convert("RGB")
    img_w, img_h = img.size

    try:
        model = YOLO(str(model_path))
    except Exception as e:
        print(json.dumps({
            "available": False,
            "source": "pose_load_failed",
            "modelPath": str(model_path),
            "error": repr(e),
            "persons": [],
        }, ensure_ascii=False))
        return

    device = 0 if torch.cuda.is_available() else "cpu"

    try:
        results = model.predict(
            source=str(image_path),
            conf=float(args.conf),
            imgsz=int(args.imgsz),
            verbose=False,
            device=device,
        )
    except Exception as e:
        print(json.dumps({
            "available": False,
            "source": "pose_predict_failed",
            "modelPath": str(model_path),
            "image": str(image_path),
            "error": repr(e),
            "persons": [],
        }, ensure_ascii=False))
        return

    persons = []

    if results:
        r = results[0]
        names = r.names or {}

        kpts_xyn = None
        kpts_conf = None

        if r.keypoints is not None:
            try:
                kpts_xyn = r.keypoints.xyn.cpu().numpy()
            except Exception:
                kpts_xyn = None
            try:
                kpts_conf = r.keypoints.conf.cpu().numpy()
            except Exception:
                kpts_conf = None

        if r.boxes is not None:
            for i, b in enumerate(r.boxes):
                score = float(b.conf[0].item())
                label = ""
                try:
                    cls_id = int(b.cls[0].item())
                    label = str(names.get(cls_id, cls_id)).strip().lower()
                except Exception:
                    label = ""

                x1, y1, x2, y2 = [float(v) for v in b.xyxy[0].tolist()]
                x1 = max(0.0, min(float(img_w), x1))
                y1 = max(0.0, min(float(img_h), y1))
                x2 = max(0.0, min(float(img_w), x2))
                y2 = max(0.0, min(float(img_h), y2))

                if x2 <= x1 or y2 <= y1:
                    continue

                box = {
                    "x": x1 / img_w,
                    "y": y1 / img_h,
                    "w": (x2 - x1) / img_w,
                    "h": (y2 - y1) / img_h,
                }

                keypoints = []
                if kpts_xyn is not None and i < len(kpts_xyn):
                    person_kpts = kpts_xyn[i]
                    for ki, (kx, ky) in enumerate(person_kpts):
                        kconf = 0.0
                        if kpts_conf is not None and i < len(kpts_conf) and ki < len(kpts_conf[i]):
                            kconf = float(kpts_conf[i][ki])

                        name = KEYPOINT_NAMES[ki] if ki < len(KEYPOINT_NAMES) else str(ki)
                        keypoints.append({
                            "name": name,
                            "x": float(kx),
                            "y": float(ky),
                            "conf": kconf,
                        })

                persons.append({
                    "label": label,
                    "score": score,
                    "box": box,
                    "keypoints": keypoints,
                })

    print(json.dumps({
        "available": True,
        "source": model_source,
        "modelPath": str(model_path),
        "image": str(image_path),
        "conf": float(args.conf),
        "imgsz": int(args.imgsz),
        "device": str(device),
        "imageWidth": img_w,
        "imageHeight": img_h,
        "personCount": len(persons),
        "persons": persons,
    }, ensure_ascii=False))


if __name__ == "__main__":
    main()