mirror of
https://github.com/MrSleeps/jellyfin-scene-search.git
synced 2026-11-15 22:41:28 +00:00
Sleeps
110 lines
3.6 KiB
Python
110 lines
3.6 KiB
Python
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
|