dtoro
574d71371f
Drops face recognition, OCR, object detection, and semantic embeddings. The sole remaining vision task is a CLIP-based binary classifier (photography vs other); photos in "other" get needs_review=true so screenshots, documents, memes and scans can be triaged from a new filter pill in the UI. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
107 lines
3.7 KiB
Python
107 lines
3.7 KiB
Python
"""
|
|
Binary content classifier: 'photography' vs 'other'.
|
|
|
|
Uses OpenCLIP ViT-B/32 image features (ONNX) and two pre-computed text
|
|
prompt centroids. Text centroids are computed once with the native
|
|
open_clip text encoder and cached to {models_dir}/classifier/vectors.npz
|
|
so steady-state worker startup doesn't pay the PyTorch cost.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
|
|
from app.config import VisionSettings
|
|
from app.services.vision.base import ClassificationResult, ContentClassifier
|
|
from app.services.vision.embed import CLIPVisualEncoder
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
PROMPTS = {
|
|
"photography": [
|
|
"a photograph taken with a camera",
|
|
"a real photo of a real scene or person",
|
|
"a candid photograph",
|
|
"a portrait photograph",
|
|
"a landscape photograph",
|
|
],
|
|
"other": [
|
|
"a screenshot of a computer screen",
|
|
"a screenshot of a phone screen",
|
|
"a screen capture of a user interface",
|
|
"a scanned document",
|
|
"a photo of a document with printed text",
|
|
"a photo of a receipt",
|
|
"a photo of a bill or invoice",
|
|
"an internet meme with text overlay",
|
|
"a funny image with caption text",
|
|
"a digital illustration or graphic design",
|
|
],
|
|
}
|
|
|
|
|
|
def _compute_text_centroids() -> dict[str, np.ndarray]:
|
|
"""Compute the 'photography' and 'other' centroid vectors using the
|
|
open_clip text encoder. Only called on the cache-miss path."""
|
|
import open_clip
|
|
import torch
|
|
|
|
logger.info("Computing CLIP text centroids for binary classifier")
|
|
model, _, _ = open_clip.create_model_and_transforms(
|
|
"ViT-B-32", pretrained="laion2b_s34b_b79k"
|
|
)
|
|
model.eval()
|
|
tokenizer = open_clip.get_tokenizer("ViT-B-32")
|
|
|
|
centroids: dict[str, np.ndarray] = {}
|
|
for label, prompts in PROMPTS.items():
|
|
tokens = tokenizer(prompts)
|
|
with torch.no_grad():
|
|
feats = model.encode_text(tokens)
|
|
feats = feats / feats.norm(dim=-1, keepdim=True)
|
|
avg = feats.mean(dim=0)
|
|
avg = avg / avg.norm()
|
|
centroids[label] = avg.numpy().astype(np.float32)
|
|
return centroids
|
|
|
|
|
|
class CLIPContentClassifier(ContentClassifier):
|
|
def __init__(self, settings: VisionSettings):
|
|
self._min_confidence = settings.classifier.min_confidence
|
|
self._encoder = CLIPVisualEncoder(settings)
|
|
|
|
cache_dir = Path(settings.models_dir) / "classifier"
|
|
cache_dir.mkdir(parents=True, exist_ok=True)
|
|
cache_path = cache_dir / "vectors.npz"
|
|
|
|
if cache_path.exists():
|
|
logger.info("Loading cached text centroids from %s", cache_path)
|
|
data = np.load(cache_path)
|
|
self._photo = data["photography"].astype(np.float32)
|
|
self._other = data["other"].astype(np.float32)
|
|
else:
|
|
centroids = _compute_text_centroids()
|
|
self._photo = centroids["photography"]
|
|
self._other = centroids["other"]
|
|
np.savez(cache_path, photography=self._photo, other=self._other)
|
|
logger.info("Cached text centroids to %s", cache_path)
|
|
|
|
def classify(self, image: np.ndarray) -> ClassificationResult:
|
|
vec = self._encoder.encode(image)
|
|
s_photo = float(np.dot(vec, self._photo))
|
|
s_other = float(np.dot(vec, self._other))
|
|
|
|
if s_photo >= s_other:
|
|
label = "photography"
|
|
margin = s_photo - s_other
|
|
else:
|
|
label = "other"
|
|
margin = s_other - s_photo
|
|
|
|
# 0.01 margin → ~0.3 conf, 0.03+ → ~1.0
|
|
confidence = min(1.0, margin * 30)
|
|
return ClassificationResult(label=label, confidence=confidence)
|