""" SigLIP2 zero-shot classifier using ONNX Runtime. Uses onnx-community/siglip2-large-patch16-256-ONNX (separate vision + text models). Zero-shot: text prompts only, no reference images needed (folder names used for class labels). """ import time from pathlib import Path import numpy as np import onnxruntime as ort from PIL import Image from huggingface_hub import hf_hub_download from transformers import AutoProcessor from jina_fewshot import IMAGE_EXTS REPO_ID = "onnx-community/siglip2-large-patch16-256-ONNX" # Use quantized models to save memory; full fp32 text_model is 2.3GB VISION_ONNX = "onnx/vision_model_quantized.onnx" TEXT_ONNX = "onnx/text_model_quantized.onnx" def _download(repo_id, filename): print(f" Downloading {filename} from {repo_id}...") path = hf_hub_download(repo_id=repo_id, filename=filename) print(f" Downloaded: {path}") return path def _make_session(onnx_path, device): available = ort.get_available_providers() if "CUDAExecutionProvider" in available and device == "cuda": providers = ["CUDAExecutionProvider", "CPUExecutionProvider"] else: providers = ["CPUExecutionProvider"] print(f" ONNX providers: {providers}") return ort.InferenceSession(onnx_path, providers=providers) class SigLIP2ONNXClassifier: """Zero-shot crop classifier using SigLIP2 ONNX (separate vision + text encoders).""" def __init__(self, device="cuda"): print("[*] Loading SigLIP2 ONNX (siglip2-large-patch16-256)...") t0 = time.perf_counter() self.device = device # Download and load vision model vision_path = _download(REPO_ID, VISION_ONNX) self.vision_session = _make_session(vision_path, device) # Download and load text model text_path = _download(REPO_ID, TEXT_ONNX) self.text_session = _make_session(text_path, device) # Processor handles both image preprocessing and tokenization self.processor = AutoProcessor.from_pretrained(REPO_ID, use_fast=False) # Map I/O names self._vision_input_names = [i.name for i in self.vision_session.get_inputs()] self._vision_output_names = [o.name for o in self.vision_session.get_outputs()] self._text_input_names = [i.name for i in self.text_session.get_inputs()] self._text_output_names = [o.name for o in self.text_session.get_outputs()] print(f" Vision inputs: {self._vision_input_names}") print(f" Vision outputs: {self._vision_output_names}") print(f" Text inputs: {self._text_input_names}") print(f" Text outputs: {self._text_output_names}") self.labels = [] self._text_embeds = None # Sanity check dummy = Image.new("RGB", (256, 256), color=(255, 0, 0)) v_emb = self._encode_image(dummy) print(f" [SANITY] vision embed shape={v_emb.shape}, norm={np.linalg.norm(v_emb):.4f}") t_emb = self._encode_texts(["a red square"]) print(f" [SANITY] text embed shape={t_emb.shape}, norm={np.linalg.norm(t_emb):.4f}") print(f"[*] SigLIP2 ONNX loaded in {time.perf_counter() - t0:.1f}s") def _encode_image(self, image): """Encode a single PIL image, return [1, D] embedding.""" processed = self.processor(images=image, return_tensors="np") pixel_values = processed["pixel_values"].astype(np.float32) feeds = {} for name in self._vision_input_names: if "pixel" in name.lower(): feeds[name] = pixel_values outputs = self.vision_session.run(self._vision_output_names, feeds) # Pick the pooler_output or last_hidden_state[:,0,:] — typically first 2D output for out in outputs: if out.ndim == 2: return out # Fallback: CLS token from 3D for out in outputs: if out.ndim == 3: return out[:, 0, :] raise RuntimeError(f"No usable vision output. Shapes: {[o.shape for o in outputs]}") def _encode_texts(self, texts): """Encode text strings, return [N, D] embeddings.""" processed = self.processor(text=texts, return_tensors="np", padding=True, truncation=True) feeds = {} for name in self._text_input_names: nl = name.lower() if "input_id" in nl and "input_ids" in processed: feeds[name] = processed["input_ids"].astype(np.int64) elif ("attention" in nl or "mask" in nl) and "attention_mask" in processed: feeds[name] = processed["attention_mask"].astype(np.int64) outputs = self.text_session.run(self._text_output_names, feeds) # Pick pooler_output (2D) or CLS from 3D for out in outputs: if out.ndim == 2: return out for out in outputs: if out.ndim == 3: return out[:, 0, :] raise RuntimeError(f"No usable text output. Shapes: {[o.shape for o in outputs]}") def build_refs(self, refs_dir, **kwargs): """Extract class names from refs_dir subfolders and precompute text embeddings.""" refs_dir = Path(refs_dir) self.labels = sorted(d.name for d in refs_dir.iterdir() if d.is_dir()) if not self.labels: raise ValueError(f"No subfolders in {refs_dir}") self._text_embeds = self._encode_texts(self.labels) print(f" SigLIP2 ONNX labels: {self.labels}") print(f" Text embeds shape: {self._text_embeds.shape}") def classify_crop(self, crop, conf_threshold, gap_threshold): """ Classify a single crop image using zero-shot SigLIP2. Computes image-text similarity via dot product + sigmoid (SigLIP style). Returns dict matching jina_fewshot.classify() format. """ image_emb = self._encode_image(crop) # [1, D] text_emb = self._text_embeds # [N, D] # SigLIP2 uses sigmoid on logits (dot product scaled by model) logits = (image_emb @ text_emb.T).squeeze(0).astype(np.float64) probs = 1.0 / (1.0 + np.exp(-logits)) # sigmoid probs = np.nan_to_num(probs, nan=0.0) sorted_idx = np.argsort(probs)[::-1] best_idx = sorted_idx[0] second_idx = sorted_idx[1] conf = float(probs[best_idx]) gap = float(probs[best_idx] - probs[second_idx]) if conf >= conf_threshold: prediction = self.labels[best_idx] status = "accepted" else: prediction = "unknown" status = f"rejected: conf {conf:.4f} < {conf_threshold}" return { "prediction": prediction, "raw_prediction": self.labels[best_idx], "confidence": conf, "gap": gap, "second_best": self.labels[second_idx], "second_conf": float(probs[second_idx]), "status": status, "all_sims": {self.labels[j]: float(probs[j]) for j in range(len(self.labels))}, }