refactor: strip AI pipeline to binary photo/other classifier
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>
This commit is contained in:
@@ -1,117 +1,106 @@
|
||||
"""
|
||||
CLIP zero-shot content-type classifier.
|
||||
Binary content classifier: 'photography' vs 'other'.
|
||||
|
||||
Uses the native OpenCLIP PyTorch text encoder for high-quality text
|
||||
embeddings (the ONNX text encoder has degraded quality due to the
|
||||
eot_indices workaround). Image embeddings use the ONNX visual encoder
|
||||
which works well.
|
||||
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
|
||||
import torch
|
||||
|
||||
from app.config import VisionSettings
|
||||
from app.services.vision.base import ContentClassifier, ClassificationResult
|
||||
from app.services.vision.base import ClassificationResult, ContentClassifier
|
||||
from app.services.vision.embed import CLIPVisualEncoder
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
CATEGORY_PROMPTS = {
|
||||
"screenshot": [
|
||||
"a screenshot of a computer screen",
|
||||
"a screenshot of a phone screen",
|
||||
"a screen capture of a user interface",
|
||||
],
|
||||
"document": [
|
||||
"a scanned document",
|
||||
"a photo of a document with printed text",
|
||||
"a photo of a page of text on paper",
|
||||
],
|
||||
"receipt": [
|
||||
"a photo of a receipt",
|
||||
"a photo of a bill or invoice",
|
||||
],
|
||||
"meme": [
|
||||
"an internet meme with text overlay",
|
||||
"a funny image with caption text",
|
||||
],
|
||||
"artwork": [
|
||||
"a painting or drawing",
|
||||
"a sketch or illustration",
|
||||
"digital art or graphic design",
|
||||
],
|
||||
"photograph": [
|
||||
|
||||
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):
|
||||
"""Zero-shot content classifier using CLIP text-image similarity.
|
||||
Uses native PyTorch for text encoding, ONNX for image encoding."""
|
||||
|
||||
def __init__(self, settings: VisionSettings):
|
||||
import open_clip
|
||||
|
||||
self._min_confidence = settings.classifier.min_confidence
|
||||
self._encoder = CLIPVisualEncoder(settings)
|
||||
|
||||
# Load native model for text encoding only.
|
||||
# Use whichever model family the embedder is configured for so
|
||||
# the classification text vectors live in the same space as the
|
||||
# image embeddings.
|
||||
embedder_name = settings.embedder.name
|
||||
if embedder_name.startswith("siglip2"):
|
||||
model_arch = "ViT-B-16-SigLIP-384"
|
||||
pretrained = "webli"
|
||||
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:
|
||||
model_arch = "ViT-B-32"
|
||||
pretrained = "laion2b_s34b_b79k"
|
||||
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)
|
||||
|
||||
logger.info("Loading %s text encoder for content classification", model_arch)
|
||||
model, _, _ = open_clip.create_model_and_transforms(
|
||||
model_arch, pretrained=pretrained
|
||||
)
|
||||
model.eval()
|
||||
self._model = model
|
||||
self._tokenizer = open_clip.get_tokenizer(model_arch)
|
||||
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))
|
||||
|
||||
# Get the ONNX image embedder from the registry
|
||||
from app.services.vision.registry import registry
|
||||
self._embedder = registry.get_embedder()
|
||||
if s_photo >= s_other:
|
||||
label = "photography"
|
||||
margin = s_photo - s_other
|
||||
else:
|
||||
label = "other"
|
||||
margin = s_other - s_photo
|
||||
|
||||
# Pre-compute text embeddings for each category
|
||||
self._category_embeddings: dict[str, np.ndarray] = {}
|
||||
for category, prompts in CATEGORY_PROMPTS.items():
|
||||
tokens = self._tokenizer(prompts)
|
||||
with torch.no_grad():
|
||||
text_features = model.encode_text(tokens)
|
||||
text_features /= text_features.norm(dim=-1, keepdim=True)
|
||||
avg = text_features.mean(dim=0)
|
||||
avg /= avg.norm()
|
||||
self._category_embeddings[category] = avg.numpy().astype(np.float32)
|
||||
|
||||
logger.info("Content classifier ready with %d categories", len(self._category_embeddings))
|
||||
|
||||
def classify(self, image: np.ndarray) -> list[ClassificationResult]:
|
||||
img_vec = self._embedder.embed_image(image)
|
||||
|
||||
# Cosine similarity against each category
|
||||
scores = {}
|
||||
for category, cat_vec in self._category_embeddings.items():
|
||||
scores[category] = float(np.dot(img_vec, cat_vec))
|
||||
|
||||
# Sort by score descending
|
||||
ranked = sorted(scores.items(), key=lambda x: -x[1])
|
||||
best_cat, best_score = ranked[0]
|
||||
second_score = ranked[1][1]
|
||||
|
||||
margin = best_score - second_score
|
||||
# Normalize: 0.01 margin → ~0.5 confidence, 0.03+ → ~1.0
|
||||
# 0.01 margin → ~0.3 conf, 0.03+ → ~1.0
|
||||
confidence = min(1.0, margin * 30)
|
||||
|
||||
if confidence >= self._min_confidence:
|
||||
return [ClassificationResult(label=best_cat, confidence=confidence)]
|
||||
|
||||
return []
|
||||
return ClassificationResult(label=label, confidence=confidence)
|
||||
|
||||
Reference in New Issue
Block a user