Spaces:
Running
Running
File size: 7,020 Bytes
81de9b1 100dbc1 81de9b1 7905374 81de9b1 100dbc1 81de9b1 100dbc1 81de9b1 1319df4 81de9b1 1319df4 81de9b1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 | """
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))},
}
|