dtoro
9282a5c734
Introduce the app/services/vision/ module with ABC interfaces, ONNX Runtime backend, model registry, and per-task implementations: - OpenCLIP ViT-B/32 embedder (image + text, 512-d) - RapidOCR engine (PP-OCRv4 via ONNX, no PaddlePaddle) - YOLOv8n object detector (raw ONNX, no ultralytics runtime) - YuNet + SFace face processor (Apache 2.0, opencv_zoo, 128-d) - DBSCAN face clustering helper Add VisionSettings to config (mulita.yml + Pydantic), bootstrap_models.py for first-boot weight downloads, models_data Docker volume, and ROCm backend stub for future GPU acceleration. No Celery tasks wired yet — models load but nothing invokes them. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
82 lines
2.8 KiB
Python
82 lines
2.8 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)
|
|
input_name = self._textual.get_inputs()[0].name
|
|
out = self._textual.run(None, {input_name: tokens})[0][0]
|
|
out = out / np.linalg.norm(out)
|
|
return out.astype(np.float32)
|
|
|
|
@property
|
|
def dim(self) -> int:
|
|
return 512
|