85 lines
3.2 KiB
Python
85 lines
3.2 KiB
Python
"""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
|