feat: upgrade to SigLIP2 ViT-B/16 for semantic search
Replace OpenCLIP ViT-B/32 (512-d, ~78% recall) with SigLIP2 ViT-B/16 (768-d, ~84% recall) as the default embedding model for significantly better image-text retrieval quality. - New SigLIP2Embedder class with 384px input and SigLIP normalization - ONNX export pipeline for SigLIP2 visual + textual encoders - Migration 0010: resize embeddings.vector from 512 to 768 dimensions - Config-driven model selection: "siglip2_vitb16" (default) or "openclip_vitb32" (legacy) — both models can coexist - Content classifier follows the configured embedder family - Existing embeddings cleared on migration; vision backfill regenerates Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -1,11 +1,15 @@
|
||||
"""
|
||||
OpenCLIP ViT-B/32 embedder using ONNX Runtime.
|
||||
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 bootstrap_models.py.
|
||||
These are exported from open_clip via export_models.py / bootstrap_models.py.
|
||||
"""
|
||||
import logging
|
||||
from pathlib import Path
|
||||
@@ -18,33 +22,46 @@ from app.services.vision.base import Embedder
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# OpenCLIP ViT-B/32 preprocessing constants (ImageNet norm)
|
||||
_MEAN = np.array([0.48145466, 0.4578275, 0.40821073], dtype=np.float32)
|
||||
_STD = np.array([0.26862954, 0.26130258, 0.27577711], dtype=np.float32)
|
||||
_INPUT_SIZE = 224
|
||||
# ── 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
|
||||
|
||||
|
||||
def _preprocess_image(image: np.ndarray) -> np.ndarray:
|
||||
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."""
|
||||
from PIL import Image
|
||||
|
||||
img = Image.fromarray(image).convert("RGB")
|
||||
# Resize shortest edge to _INPUT_SIZE, then center crop
|
||||
w, h = img.size
|
||||
scale = _INPUT_SIZE / min(w, h)
|
||||
scale = input_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 - input_size) // 2
|
||||
top = (h - input_size) // 2
|
||||
img = img.crop((left, top, left + input_size, top + input_size))
|
||||
|
||||
arr = np.array(img, dtype=np.float32) / 255.0
|
||||
arr = (arr - _MEAN) / _STD
|
||||
arr = (arr - mean) / std
|
||||
arr = arr.transpose(2, 0, 1) # HWC → CHW
|
||||
return arr[np.newaxis] # NCHW
|
||||
|
||||
|
||||
class OpenCLIPEmbedder(Embedder):
|
||||
"""Legacy OpenCLIP ViT-B/32 embedder (512-d)."""
|
||||
|
||||
def __init__(self, settings: VisionSettings):
|
||||
model_dir = Path(settings.models_dir) / "embed"
|
||||
visual_path = model_dir / "visual.onnx"
|
||||
@@ -54,14 +71,14 @@ class OpenCLIPEmbedder(Embedder):
|
||||
opts.inter_op_num_threads = 2
|
||||
opts.intra_op_num_threads = 2
|
||||
|
||||
logger.info("Loading visual encoder from %s", visual_path)
|
||||
logger.info("Loading OpenCLIP visual encoder from %s", visual_path)
|
||||
self._visual = ort.InferenceSession(str(visual_path), opts, providers=["CPUExecutionProvider"])
|
||||
|
||||
logger.info("Loading textual encoder from %s", textual_path)
|
||||
logger.info("Loading OpenCLIP textual encoder from %s", textual_path)
|
||||
self._textual = ort.InferenceSession(str(textual_path), opts, providers=["CPUExecutionProvider"])
|
||||
|
||||
def embed_image(self, image: np.ndarray) -> np.ndarray:
|
||||
inp = _preprocess_image(image)
|
||||
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]
|
||||
out = out / np.linalg.norm(out)
|
||||
@@ -71,7 +88,6 @@ class OpenCLIPEmbedder(Embedder):
|
||||
import open_clip
|
||||
tokenizer = open_clip.get_tokenizer("ViT-B-32")
|
||||
tokens = tokenizer([text]).numpy().astype(np.int64)
|
||||
# Compute EOT indices outside ONNX (avoids ArgMax(13) op)
|
||||
eot_indices = tokens.argmax(axis=-1).astype(np.int64)
|
||||
inputs = self._textual.get_inputs()
|
||||
out = self._textual.run(None, {
|
||||
@@ -84,3 +100,47 @@ class OpenCLIPEmbedder(Embedder):
|
||||
@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"
|
||||
|
||||
opts = ort.SessionOptions()
|
||||
opts.inter_op_num_threads = 2
|
||||
opts.intra_op_num_threads = 2
|
||||
|
||||
logger.info("Loading SigLIP2 visual encoder from %s", visual_path)
|
||||
self._visual = ort.InferenceSession(str(visual_path), opts, providers=["CPUExecutionProvider"])
|
||||
|
||||
logger.info("Loading SigLIP2 textual encoder from %s", textual_path)
|
||||
self._textual = ort.InferenceSession(str(textual_path), opts, providers=["CPUExecutionProvider"])
|
||||
|
||||
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-SigLIP2")
|
||||
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