Add explicit-content caption router

scene_search/routing.py

@@ -0,0 +1,109 @@
+from __future__ import annotations
+
+from bisect import bisect_left
+from pathlib import Path
+
+import torch
+from PIL import Image
+from transformers import AutoImageProcessor, AutoModelForImageClassification
+
+
+class ExplicitContentRouter:
+    """Lightweight image classifier used only to choose a captioning model."""
+
+    def __init__(
+        self,
+        model_id: str,
+        device: str = "cuda",
+        batch_size: int = 32,
+    ) -> None:
+        self.model_id = model_id
+        self.batch_size = batch_size
+        if device.startswith("cuda") and not torch.cuda.is_available():
+            device = "cpu"
+        self.device = torch.device(device)
+
+        self.processor = AutoImageProcessor.from_pretrained(model_id)
+        self.model = AutoModelForImageClassification.from_pretrained(model_id)
+        self.model.eval()
+        self.model.to(self.device)
+
+        labels = {
+            int(index): str(label).strip().lower()
+            for index, label in self.model.config.id2label.items()
+        }
+        matches = [index for index, label in labels.items() if label == "nsfw"]
+        if not matches:
+            raise RuntimeError(
+                f"Router model {model_id!r} has no 'nsfw' output label: {labels}"
+            )
+        self.explicit_index = matches[0]
+
+    def score_paths(self, paths: list[Path]) -> list[float]:
+        scores: list[float] = []
+        for start in range(0, len(paths), self.batch_size):
+            batch_paths = paths[start:start + self.batch_size]
+            images: list[Image.Image] = []
+            try:
+                for path in batch_paths:
+                    images.append(Image.open(path).convert("RGB"))
+                inputs = self.processor(images=images, return_tensors="pt")
+                inputs = {
+                    key: value.to(self.device)
+                    for key, value in inputs.items()
+                }
+                with torch.inference_mode():
+                    logits = self.model(**inputs).logits
+                    probabilities = torch.softmax(logits, dim=-1)
+                scores.extend(
+                    float(value)
+                    for value in probabilities[:, self.explicit_index]
+                    .detach()
+                    .cpu()
+                    .tolist()
+                )
+            finally:
+                for image in images:
+                    image.close()
+        return scores
+
+
+def apply_context_routing(
+    rows: list[dict],
+    scores: dict[str, float],
+    threshold: float,
+    context_ms: int,
+) -> dict[str, tuple[str, str]]:
+    """Return frame id -> (route, reason), expanding positive hits to nearby frames."""
+
+    positive_times: dict[str, list[int]] = {}
+    for row in rows:
+        frame_id = str(row["id"])
+        score = scores[frame_id]
+        if score >= threshold:
+            positive_times.setdefault(str(row["media_id"]), []).append(
+                int(row["timestamp_ms"])
+            )
+
+    for times in positive_times.values():
+        times.sort()
+
+    decisions: dict[str, tuple[str, str]] = {}
+    for row in rows:
+        frame_id = str(row["id"])
+        media_id = str(row["media_id"])
+        timestamp_ms = int(row["timestamp_ms"])
+        score = scores[frame_id]
+
+        if score >= threshold:
+            decisions[frame_id] = ("specialist", "score")
+            continue
+
+        times = positive_times.get(media_id, [])
+        left = bisect_left(times, timestamp_ms - context_ms)
+        if left < len(times) and times[left] <= timestamp_ms + context_ms:
+            decisions[frame_id] = ("specialist", "context")
+        else:
+            decisions[frame_id] = ("general", "score")
+
+    return decisions