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>
84 lines
2.7 KiB
Python
84 lines
2.7 KiB
Python
"""
|
|
Download vision model weights on first worker boot.
|
|
|
|
Run as: python -m app.services.vision.bootstrap_models
|
|
|
|
Or called from the vision worker entrypoint before Celery starts.
|
|
Downloads are idempotent — existing files with matching sizes are skipped.
|
|
"""
|
|
import logging
|
|
import os
|
|
from pathlib import Path
|
|
from urllib.request import urlretrieve
|
|
|
|
from app.config import settings
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# (relative_path, url, expected_size_bytes_approx)
|
|
# Sizes are approximate — used only for skip-if-exists checks, not integrity.
|
|
MODELS = [
|
|
# OpenCLIP ViT-B/32 — visual and textual encoders (ONNX)
|
|
# These must be exported manually via export_openclip.py (see below).
|
|
# Placeholder entries — bootstrap will warn if missing.
|
|
("embed/visual.onnx", None, None),
|
|
("embed/textual.onnx", None, None),
|
|
|
|
# YOLOv8n — object detection
|
|
# Export: `yolo export model=yolov8n.pt format=onnx imgsz=640`
|
|
# Placeholder — must be exported from ultralytics offline.
|
|
("detect/yolov8n.onnx", None, None),
|
|
|
|
# YuNet — face detection (Apache 2.0, opencv_zoo)
|
|
(
|
|
"face/yunet.onnx",
|
|
"https://github.com/opencv/opencv_zoo/raw/main/models/face_detection_yunet/face_detection_yunet_2023mar.onnx",
|
|
233_000,
|
|
),
|
|
|
|
# SFace — face recognition (Apache 2.0, opencv_zoo)
|
|
(
|
|
"face/sface.onnx",
|
|
"https://github.com/opencv/opencv_zoo/raw/main/models/face_recognition_sface/face_recognition_sface_2021dec.onnx",
|
|
37_000_000,
|
|
),
|
|
]
|
|
|
|
|
|
def bootstrap(models_dir: str | None = None):
|
|
"""Ensure all model files are present. Download what we can, warn about
|
|
files that need manual export."""
|
|
base = Path(models_dir or settings.vision.models_dir)
|
|
base.mkdir(parents=True, exist_ok=True)
|
|
|
|
for rel_path, url, expected_size in MODELS:
|
|
dest = base / rel_path
|
|
dest.parent.mkdir(parents=True, exist_ok=True)
|
|
|
|
if dest.exists():
|
|
logger.debug("Model already exists: %s", dest)
|
|
continue
|
|
|
|
if url is None:
|
|
logger.warning(
|
|
"Model file %s not found and has no auto-download URL. "
|
|
"See bootstrap_models.py for export instructions.",
|
|
dest,
|
|
)
|
|
continue
|
|
|
|
logger.info("Downloading %s → %s", url, dest)
|
|
try:
|
|
urlretrieve(url, str(dest))
|
|
actual = dest.stat().st_size
|
|
logger.info("Downloaded %s (%d bytes)", rel_path, actual)
|
|
except Exception as e:
|
|
logger.error("Failed to download %s: %s", rel_path, e)
|
|
if dest.exists():
|
|
dest.unlink()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
logging.basicConfig(level=logging.INFO)
|
|
bootstrap()
|