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,146 +1,58 @@
|
||||
"""
|
||||
CLIP / SigLIP2 embedder using ONNX Runtime.
|
||||
|
||||
Supports two model families:
|
||||
- OpenCLIP ViT-B/32 (512-d) — legacy, config name "openclip_vitb32"
|
||||
- SigLIP2 ViT-B/16 (768-d) — default, config name "siglip2_vitb16"
|
||||
|
||||
Expects two ONNX files under {models_dir}/embed/:
|
||||
- visual.onnx (image encoder)
|
||||
- textual.onnx (text encoder)
|
||||
|
||||
These are exported from open_clip via export_models.py / bootstrap_models.py.
|
||||
OpenCLIP ViT-B/32 visual encoder (ONNX). Produces 512-d image features
|
||||
consumed by the content classifier. Not exposed as a standalone service;
|
||||
the classifier owns the lifecycle.
|
||||
"""
|
||||
import logging
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import onnxruntime as ort
|
||||
import onnxruntime as ort # noqa: F401 (provider plumbing relies on this)
|
||||
|
||||
from app.config import VisionSettings
|
||||
from app.services.vision.base import Embedder
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# ── Model-specific constants ──────────────────────────────────────────
|
||||
|
||||
# OpenCLIP ViT-B/32 (ImageNet norm, 224px)
|
||||
_OPENCLIP_MEAN = np.array([0.48145466, 0.4578275, 0.40821073], dtype=np.float32)
|
||||
_OPENCLIP_STD = np.array([0.26862954, 0.26130258, 0.27577711], dtype=np.float32)
|
||||
_OPENCLIP_SIZE = 224
|
||||
|
||||
# SigLIP2 ViT-B/16 (SigLIP norm, 384px)
|
||||
_SIGLIP2_MEAN = np.array([0.5, 0.5, 0.5], dtype=np.float32)
|
||||
_SIGLIP2_STD = np.array([0.5, 0.5, 0.5], dtype=np.float32)
|
||||
_SIGLIP2_SIZE = 384
|
||||
_MEAN = np.array([0.48145466, 0.4578275, 0.40821073], dtype=np.float32)
|
||||
_STD = np.array([0.26862954, 0.26130258, 0.27577711], dtype=np.float32)
|
||||
_SIZE = 224
|
||||
|
||||
|
||||
def _preprocess_image(
|
||||
image: np.ndarray,
|
||||
input_size: int,
|
||||
mean: np.ndarray,
|
||||
std: np.ndarray,
|
||||
) -> np.ndarray:
|
||||
"""Resize, center-crop, normalize an RGB uint8 image to NCHW float32."""
|
||||
def _preprocess(image: np.ndarray) -> np.ndarray:
|
||||
from PIL import Image
|
||||
|
||||
img = Image.fromarray(image).convert("RGB")
|
||||
w, h = img.size
|
||||
scale = input_size / min(w, h)
|
||||
scale = _SIZE / min(w, h)
|
||||
img = img.resize((int(w * scale), int(h * scale)), Image.BICUBIC)
|
||||
w, h = img.size
|
||||
left = (w - input_size) // 2
|
||||
top = (h - input_size) // 2
|
||||
img = img.crop((left, top, left + input_size, top + input_size))
|
||||
left = (w - _SIZE) // 2
|
||||
top = (h - _SIZE) // 2
|
||||
img = img.crop((left, top, left + _SIZE, top + _SIZE))
|
||||
|
||||
arr = np.array(img, dtype=np.float32) / 255.0
|
||||
arr = (arr - mean) / std
|
||||
arr = arr.transpose(2, 0, 1) # HWC → CHW
|
||||
return arr[np.newaxis] # NCHW
|
||||
arr = (arr - _MEAN) / _STD
|
||||
arr = arr.transpose(2, 0, 1)
|
||||
return arr[np.newaxis]
|
||||
|
||||
|
||||
class OpenCLIPEmbedder(Embedder):
|
||||
"""Legacy OpenCLIP ViT-B/32 embedder (512-d)."""
|
||||
class CLIPVisualEncoder:
|
||||
"""OpenCLIP ViT-B/32 image encoder, 512-d normalized output."""
|
||||
|
||||
def __init__(self, settings: VisionSettings):
|
||||
model_dir = Path(settings.models_dir) / "embed"
|
||||
visual_path = model_dir / "visual.onnx"
|
||||
textual_path = model_dir / "textual.onnx"
|
||||
|
||||
model_path = Path(settings.models_dir) / "embed" / "visual.onnx"
|
||||
from app.services.vision.providers import create_session
|
||||
from app.config import settings as app_settings
|
||||
providers = app_settings.vision.execution_providers
|
||||
|
||||
logger.info("Loading OpenCLIP visual encoder from %s", visual_path)
|
||||
self._visual = create_session(str(visual_path), configured_providers=providers)
|
||||
logger.info("Loading CLIP visual encoder from %s", model_path)
|
||||
self._session = create_session(
|
||||
str(model_path),
|
||||
configured_providers=app_settings.vision.execution_providers,
|
||||
)
|
||||
|
||||
logger.info("Loading OpenCLIP textual encoder from %s", textual_path)
|
||||
self._textual = create_session(str(textual_path), configured_providers=providers)
|
||||
|
||||
def embed_image(self, image: np.ndarray) -> np.ndarray:
|
||||
inp = _preprocess_image(image, _OPENCLIP_SIZE, _OPENCLIP_MEAN, _OPENCLIP_STD)
|
||||
input_name = self._visual.get_inputs()[0].name
|
||||
out = self._visual.run(None, {input_name: inp})[0][0]
|
||||
def encode(self, image: np.ndarray) -> np.ndarray:
|
||||
inp = _preprocess(image)
|
||||
name = self._session.get_inputs()[0].name
|
||||
out = self._session.run(None, {name: inp})[0][0]
|
||||
out = out / np.linalg.norm(out)
|
||||
return out.astype(np.float32)
|
||||
|
||||
def embed_text(self, text: str) -> np.ndarray:
|
||||
import open_clip
|
||||
tokenizer = open_clip.get_tokenizer("ViT-B-32")
|
||||
tokens = tokenizer([text]).numpy().astype(np.int64)
|
||||
eot_indices = tokens.argmax(axis=-1).astype(np.int64)
|
||||
inputs = self._textual.get_inputs()
|
||||
out = self._textual.run(None, {
|
||||
inputs[0].name: tokens,
|
||||
inputs[1].name: eot_indices,
|
||||
})[0][0]
|
||||
out = out / np.linalg.norm(out)
|
||||
return out.astype(np.float32)
|
||||
|
||||
@property
|
||||
def dim(self) -> int:
|
||||
return 512
|
||||
|
||||
|
||||
class SigLIP2Embedder(Embedder):
|
||||
"""SigLIP2 ViT-B/16 embedder (768-d) — higher recall than OpenCLIP."""
|
||||
|
||||
def __init__(self, settings: VisionSettings):
|
||||
model_dir = Path(settings.models_dir) / "embed_siglip2"
|
||||
visual_path = model_dir / "visual.onnx"
|
||||
textual_path = model_dir / "textual.onnx"
|
||||
|
||||
from app.services.vision.providers import create_session
|
||||
from app.config import settings as app_settings
|
||||
providers = app_settings.vision.execution_providers
|
||||
|
||||
logger.info("Loading SigLIP2 visual encoder from %s", visual_path)
|
||||
self._visual = create_session(str(visual_path), configured_providers=providers)
|
||||
|
||||
logger.info("Loading SigLIP2 textual encoder from %s", textual_path)
|
||||
self._textual = create_session(str(textual_path), configured_providers=providers)
|
||||
|
||||
def embed_image(self, image: np.ndarray) -> np.ndarray:
|
||||
inp = _preprocess_image(image, _SIGLIP2_SIZE, _SIGLIP2_MEAN, _SIGLIP2_STD)
|
||||
input_name = self._visual.get_inputs()[0].name
|
||||
out = self._visual.run(None, {input_name: inp})[0][0]
|
||||
out = out / np.linalg.norm(out)
|
||||
return out.astype(np.float32)
|
||||
|
||||
def embed_text(self, text: str) -> np.ndarray:
|
||||
import open_clip
|
||||
tokenizer = open_clip.get_tokenizer("ViT-B-16-SigLIP-384")
|
||||
tokens = tokenizer([text]).numpy().astype(np.int64)
|
||||
inputs = self._textual.get_inputs()
|
||||
feed = {inputs[0].name: tokens}
|
||||
# SigLIP2 text encoder may need attention mask
|
||||
if len(inputs) > 1:
|
||||
attention_mask = (tokens != 0).astype(np.int64)
|
||||
feed[inputs[1].name] = attention_mask
|
||||
out = self._textual.run(None, feed)[0][0]
|
||||
out = out / np.linalg.norm(out)
|
||||
return out.astype(np.float32)
|
||||
|
||||
@property
|
||||
def dim(self) -> int:
|
||||
return 768
|
||||
|
||||
Reference in New Issue
Block a user