dtoro
2a6661f779
- Add export_models.py for OpenCLIP ViT-B/32 and YOLOv8n ONNX export - Fix ArgMax(13) ORT ARM64 incompatibility by passing eot_indices as a separate ONNX input (computed outside the graph in embed.py) - Use legacy TorchScript exporter (dynamo=False) for IR version 9 compat - Upgrade onnxruntime to 1.18.1 - Rewrite bootstrap_models.py with clear separation of auto-downloadable models (YuNet, SFace) vs manually-exported ones (OpenCLIP, YOLOv8n) - Wire bootstrap into worker CMD (runs before Celery) Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
87 lines
2.9 KiB
Python
87 lines
2.9 KiB
Python
"""
|
|
OpenCLIP ViT-B/32 embedder using ONNX Runtime.
|
|
|
|
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.
|
|
"""
|
|
import logging
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import onnxruntime as ort
|
|
|
|
from app.config import VisionSettings
|
|
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
|
|
|
|
|
|
def _preprocess_image(image: 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)
|
|
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))
|
|
|
|
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
|
|
|
|
|
|
class OpenCLIPEmbedder(Embedder):
|
|
def __init__(self, settings: VisionSettings):
|
|
model_dir = Path(settings.models_dir) / "embed"
|
|
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 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)
|
|
self._textual = ort.InferenceSession(str(textual_path), opts, providers=["CPUExecutionProvider"])
|
|
|
|
def embed_image(self, image: np.ndarray) -> np.ndarray:
|
|
inp = _preprocess_image(image)
|
|
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-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, {
|
|
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
|