US02-01: Extract Safety and Existing Web UI Donors (#54)

photo_pipeline/integrations/nsfw_model.py

@@ -0,0 +1,84 @@
+"""Local NSFW model adapter.
+
+Wraps the on-device ViT classifier behind a small interface so the safety service
+depends on ``NsfwModel.score`` rather than on torch/transformers or the archived
+CLI. Heavy imports are lazy: importing this module never loads the model, so the
+rest of the app (and its tests) stay light. Scoring itself is non-deterministic
+across hardware and is exercised in dedicated model tests, not the unit suite.
+
+Extracted from nsfwtag.scoring.score_images (donor_ledger.yaml: nt-score-model).
+The donor's CSV cache is dropped here — scores are persisted in the database by
+asset ID (safety repository), not a path-keyed CSV.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+MODEL_ID = "AdamCodd/vit-base-nsfw-detector"
+BATCH = 16
+
+
+class NsfwModel:
+    def __init__(self, model_id: str = MODEL_ID, batch: int = BATCH) -> None:
+        self.model_id = model_id
+        self.batch = batch
+        self._model = None
+        self._device = None
+        self._size = None
+        self._nsfw_idx = None
+
+    def _ensure_loaded(self) -> None:
+        if self._model is not None:
+            return
+        import torch
+        from transformers import AutoModelForImageClassification
+
+        self._device = "mps" if torch.backends.mps.is_available() else "cpu"
+        try:
+            model = AutoModelForImageClassification.from_pretrained(
+                self.model_id, local_files_only=True
+            )
+        except Exception:
+            model = AutoModelForImageClassification.from_pretrained(self.model_id)
+        self._model = model.to(self._device).eval()
+        self._nsfw_idx = next(
+            i for i, label in model.config.id2label.items() if label.lower() == "nsfw"
+        )
+        self._size = getattr(model.config, "image_size", 224)
+
+    def score(self, paths: list[Path | str]) -> list[tuple[str, float]]:
+        """Return ``[(path, nsfw_probability)]`` for each readable image."""
+        if not paths:
+            return []
+        self._ensure_loaded()
+        import numpy as np
+        import torch
+        from PIL import Image, ImageFile
+
+        ImageFile.LOAD_TRUNCATED_IMAGES = True
+
+        def preprocess(image):
+            image = image.convert("RGB").resize((self._size, self._size), Image.BILINEAR)
+            array = (np.asarray(image, dtype="float32") / 255.0 - 0.5) / 0.5
+            return torch.from_numpy(array).permute(2, 0, 1)
+
+        results: list[tuple[str, float]] = []
+        items = [str(p) for p in paths]
+        for start in range(0, len(items), self.batch):
+            tensors, batch_paths = [], []
+            for path in items[start : start + self.batch]:
+                try:
+                    tensors.append(preprocess(Image.open(path)))
+                    batch_paths.append(path)
+                except Exception:
+                    continue
+            if not tensors:
+                continue
+            with torch.no_grad():
+                probs = self._model(
+                    pixel_values=torch.stack(tensors).to(self._device)
+                ).logits.softmax(-1)
+            for path, prob in zip(batch_paths, probs):
+                results.append((path, float(prob[self._nsfw_idx].item())))
+        return results