app/classifier.py

"""NSFW classifier — wraps an ONNX-exported Vision Transformer (or compatible).

Default model: LukeJacob2023/nsfw-image-detector (ViT-base, 5 classes).
The companion `<model>.json` file (written by scripts/convert_model.py) supplies
the exact image size, channel mean/std, and label order — so swapping models
needs only a re-run of convert_model.py with --model <other-id>.
"""
from __future__ import annotations

import io
import json
import logging
import threading
from dataclasses import dataclass
from pathlib import Path
from typing import Dict, Sequence

import numpy as np
import onnxruntime as ort
from PIL import Image

from app.config import CATEGORIES

LOGGER = logging.getLogger(__name__)


@dataclass(frozen=True)
class Prediction:
    """Per-image classification result."""
    scores: Dict[str, float]  # category -> probability in [0, 1]

    @property
    def top_category(self) -> str:
        return max(self.scores, key=self.scores.get)

    @property
    def top_score(self) -> float:
        return self.scores[self.top_category]

    def is_blocked(
        self,
        threshold_percent: int,
        blocked_categories: Sequence[str],
    ) -> tuple[bool, str | None, float]:
        thr = threshold_percent / 100.0
        best_cat: str | None = None
        best_score = 0.0
        for cat in blocked_categories:
            score = self.scores.get(cat, 0.0)
            if score >= thr and score > best_score:
                best_cat = cat
                best_score = score
        return (best_cat is not None, best_cat, best_score)


def _softmax(x: np.ndarray) -> np.ndarray:
    x = x - np.max(x, axis=-1, keepdims=True)
    e = np.exp(x)
    return e / np.sum(e, axis=-1, keepdims=True)


class NSFWClassifier:
    """Loads ONNX model + metadata once. Thread-safe inference."""

    # Fallback values used only if the metadata file is missing
    _FALLBACK_SIZE = (224, 224)
    _FALLBACK_MEAN = (0.5, 0.5, 0.5)
    _FALLBACK_STD = (0.5, 0.5, 0.5)

    def __init__(self, model_path: str):
        self._model_path = Path(model_path)
        self._session: ort.InferenceSession | None = None
        self._input_name: str | None = None
        self._lock = threading.Lock()

        self._labels: list[str] = list(CATEGORIES)
        self._size: tuple[int, int] = self._FALLBACK_SIZE
        self._mean = np.array(self._FALLBACK_MEAN, dtype=np.float32)
        self._std = np.array(self._FALLBACK_STD, dtype=np.float32)
        self._needs_softmax = True  # ViT exports raw logits

    def load(self) -> None:
        LOGGER.info("Loading NSFW model from %s", self._model_path)

        meta_path = self._model_path.with_suffix(".json")
        if meta_path.exists():
            meta = json.loads(meta_path.read_text())
            self._labels = list(meta["labels"])
            h, w = meta.get("image_size", [224, 224])
            self._size = (int(h), int(w))
            self._mean = np.array(meta["image_mean"], dtype=np.float32)
            self._std = np.array(meta["image_std"], dtype=np.float32)
            LOGGER.info(
                "Model metadata: id=%s labels=%s size=%s",
                meta.get("model_id"),
                self._labels,
                self._size,
            )
        else:
            LOGGER.warning(
                "Metadata file %s not found — falling back to defaults", meta_path
            )

        sess_options = ort.SessionOptions()
        sess_options.intra_op_num_threads = 1
        sess_options.inter_op_num_threads = 1
        self._session = ort.InferenceSession(
            str(self._model_path),
            sess_options=sess_options,
            providers=["CPUExecutionProvider"],
        )
        self._input_name = self._session.get_inputs()[0].name

        in_shape = self._session.get_inputs()[0].shape
        out_shape = self._session.get_outputs()[0].shape
        LOGGER.info("Model ready. Input=%s, output=%s", in_shape, out_shape)

    def _preprocess(self, image_bytes: bytes) -> np.ndarray:
        """Resize → normalize with ImageNet (or model-supplied) stats → NCHW float32."""
        h, w = self._size
        with Image.open(io.BytesIO(image_bytes)) as img:
            img = img.convert("RGB").resize((w, h), Image.BILINEAR)
            arr = np.asarray(img, dtype=np.float32) / 255.0
        # HWC → normalize per-channel
        arr = (arr - self._mean) / self._std
        # HWC → CHW → batch
        arr = np.transpose(arr, (2, 0, 1))
        return np.expand_dims(arr, axis=0).astype(np.float32)

    def classify(self, image_bytes: bytes) -> Prediction:
        if self._session is None or self._input_name is None:
            raise RuntimeError("Model not loaded — call load() first")

        x = self._preprocess(image_bytes)

        with self._lock:
            outputs = self._session.run(None, {self._input_name: x})

        logits = np.asarray(outputs[0]).reshape(-1)
        if logits.shape[0] != len(self._labels):
            raise RuntimeError(
                f"Unexpected output shape {logits.shape}, "
                f"expected {len(self._labels)} probabilities"
            )

        probs = _softmax(logits) if self._needs_softmax else logits
        return Prediction(
            scores={label: float(p) for label, p in zip(self._labels, probs)}
        )

    def classify_many(self, images: Sequence[bytes]) -> list[Prediction]:
        return [self.classify(img) for img in images]