Files
mule-image/backend/app/services/vision/bootstrap_models.py
dtoro fa9b21856f feat: replace face pipeline with InsightFace, add content classifier
Face detection/recognition:
- Replace YuNet + SFace with InsightFace buffalo_l (RetinaFace + ArcFace)
- 512-d ArcFace embeddings (was 128-d SFace), migration 0006 resizes column
- Remove YOLO person-bbox workaround — RetinaFace is accurate enough
- Detection threshold 0.65 cleanly separates real faces (0.72+) from
  false positives on dogs/paintings (0.56-0.61)

Content-type classification:
- CLIP zero-shot classifier using native PyTorch text encoder + ONNX
  image encoder for high-quality text-image similarity
- Categories: photograph, screenshot, document, receipt, meme, artwork
- Writes Tag(kind=content_type) per photo via photo_tags
- Margin-based confidence: top-1 vs top-2 score difference
- New ClassifierSettings in config (enabled, min_confidence)
- Wired into vision_fanout pipeline

Tested: 6 real faces from 4 photos (zero false positives), 11/13 photos
classified (8 photograph, 2 artwork, 1 meme).

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-04-10 13:49:02 +02:00

86 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 are skipped.
For models that require export (OpenCLIP, YOLOv8n), see export_models.py.
Those must be exported once on any machine with pip, then placed in
the models volume before the worker starts.
"""
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, description)
# Models with url=None must be pre-exported via export_models.py.
# InsightFace (RetinaFace + ArcFace) auto-downloads via the insightface
# package on first use — no manual download entries needed.
DOWNLOADS = []
# Models that need manual export via export_models.py
EXPORTS = [
("embed/visual.onnx", "OpenCLIP ViT-B/32 visual encoder"),
("embed/textual.onnx", "OpenCLIP ViT-B/32 textual encoder"),
("detect/yolov8n.onnx", "YOLOv8n object detector"),
]
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)
# Download auto-downloadable models
for rel_path, url, desc in DOWNLOADS:
dest = base / rel_path
dest.parent.mkdir(parents=True, exist_ok=True)
if dest.exists():
logger.debug("Already exists: %s (%s)", dest, desc)
continue
logger.info("Downloading %s%s", desc, dest)
try:
urlretrieve(url, str(dest))
size_kb = dest.stat().st_size / 1024
logger.info("Downloaded %s (%.0f KB)", desc, size_kb)
except Exception as e:
logger.error("Failed to download %s: %s", desc, e)
if dest.exists():
dest.unlink()
# Check for manually-exported models
missing = []
for rel_path, desc in EXPORTS:
dest = base / rel_path
if not dest.exists():
missing.append((rel_path, desc))
if missing:
logger.warning(
"Missing %d model(s) that require manual export via export_models.py:",
len(missing),
)
for rel_path, desc in missing:
logger.warning(" %s%s", base / rel_path, desc)
logger.warning(
"Run: python -m app.services.vision.export_models --models-dir %s",
base,
)
else:
logger.info("All model files present in %s", base)
if __name__ == "__main__":
logging.basicConfig(level=logging.INFO)
bootstrap()