feat: add vision pipeline scaffolding with ONNX backend
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>
This commit is contained in:
153
backend/app/services/vision/faces.py
Normal file
153
backend/app/services/vision/faces.py
Normal file
@@ -0,0 +1,153 @@
|
||||
"""
|
||||
Face detection (YuNet) + recognition (SFace) using ONNX Runtime.
|
||||
|
||||
Both models are from opencv_zoo (Apache 2.0 license).
|
||||
Expects {models_dir}/face/:
|
||||
- yunet.onnx (~75 KB)
|
||||
- sface.onnx (~37 MB, 128-d embeddings)
|
||||
"""
|
||||
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 FaceProcessor, FaceDetection
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_YUNET_INPUT_SIZE = 640
|
||||
|
||||
|
||||
def _align_face(image: np.ndarray, landmarks: np.ndarray) -> np.ndarray:
|
||||
"""Align and crop a 112x112 face patch using 5-point landmarks.
|
||||
landmarks shape: (5, 2) — left_eye, right_eye, nose, left_mouth, right_mouth."""
|
||||
from PIL import Image
|
||||
import math
|
||||
|
||||
left_eye = landmarks[0]
|
||||
right_eye = landmarks[1]
|
||||
|
||||
dx = right_eye[0] - left_eye[0]
|
||||
dy = right_eye[1] - left_eye[1]
|
||||
angle = math.degrees(math.atan2(dy, dx))
|
||||
eye_center = ((left_eye[0] + right_eye[0]) / 2, (left_eye[1] + right_eye[1]) / 2)
|
||||
eye_dist = math.sqrt(dx * dx + dy * dy)
|
||||
scale = 64.0 / max(eye_dist, 1e-6) # target: eyes at ~64px apart in 112x112
|
||||
|
||||
img = Image.fromarray(image)
|
||||
img = img.rotate(-angle, center=eye_center, resample=Image.BICUBIC)
|
||||
|
||||
cx, cy = eye_center
|
||||
half = 56.0 / scale
|
||||
crop = img.crop((int(cx - half), int(cy - half * 0.8), int(cx + half), int(cy + half * 1.2)))
|
||||
crop = crop.resize((112, 112), Image.BICUBIC)
|
||||
return np.array(crop, dtype=np.float32)
|
||||
|
||||
|
||||
class YuNetSFaceProcessor(FaceProcessor):
|
||||
def __init__(self, settings: VisionSettings):
|
||||
face_dir = Path(settings.models_dir) / "face"
|
||||
yunet_path = face_dir / "yunet.onnx"
|
||||
sface_path = face_dir / "sface.onnx"
|
||||
|
||||
opts = ort.SessionOptions()
|
||||
opts.inter_op_num_threads = 2
|
||||
opts.intra_op_num_threads = 2
|
||||
|
||||
logger.info("Loading YuNet from %s", yunet_path)
|
||||
self._detector = ort.InferenceSession(str(yunet_path), opts, providers=["CPUExecutionProvider"])
|
||||
|
||||
logger.info("Loading SFace from %s", sface_path)
|
||||
self._recognizer = ort.InferenceSession(str(sface_path), opts, providers=["CPUExecutionProvider"])
|
||||
|
||||
self._min_face_size = settings.faces.min_face_size
|
||||
self._score_threshold = settings.faces.recognition_threshold
|
||||
|
||||
def process(self, image: np.ndarray) -> list[FaceDetection]:
|
||||
orig_h, orig_w = image.shape[:2]
|
||||
|
||||
# Scale image for YuNet (expects fixed input size)
|
||||
scale = min(_YUNET_INPUT_SIZE / orig_w, _YUNET_INPUT_SIZE / orig_h)
|
||||
new_w = int(orig_w * scale)
|
||||
new_h = int(orig_h * scale)
|
||||
|
||||
from PIL import Image as PILImage
|
||||
resized = np.array(
|
||||
PILImage.fromarray(image).resize((new_w, new_h), PILImage.BICUBIC),
|
||||
dtype=np.uint8,
|
||||
)
|
||||
|
||||
# YuNet expects BGR, uint8, NHWC
|
||||
bgr = resized[:, :, ::-1].copy()
|
||||
|
||||
# Run detection
|
||||
det_input = self._detector.get_inputs()[0]
|
||||
# YuNet uses dynamic input — reshape
|
||||
blob = bgr.astype(np.float32)[np.newaxis] # (1, H, W, 3)
|
||||
# Some YuNet ONNX exports expect (1, 3, H, W)
|
||||
if det_input.shape and len(det_input.shape) == 4 and det_input.shape[1] == 3:
|
||||
blob = blob.transpose(0, 3, 1, 2)
|
||||
|
||||
detections_raw = self._detector.run(None, {det_input.name: blob})
|
||||
dets = detections_raw[0] # (N, 15): x,y,w,h,score, 5x landmark pairs
|
||||
|
||||
if dets is None or len(dets) == 0:
|
||||
return []
|
||||
|
||||
results = []
|
||||
for det in dets:
|
||||
score = float(det[4]) if len(det) > 4 else float(det[-1])
|
||||
if score < self._score_threshold:
|
||||
continue
|
||||
|
||||
x, y, w, h = det[0], det[1], det[2], det[3]
|
||||
|
||||
# Filter small faces
|
||||
face_size = max(w, h) / scale
|
||||
if face_size < self._min_face_size:
|
||||
continue
|
||||
|
||||
# Rescale to original image coords
|
||||
x1 = x / scale
|
||||
y1 = y / scale
|
||||
x2 = (x + w) / scale
|
||||
y2 = (y + h) / scale
|
||||
|
||||
bbox = [
|
||||
max(0, x1 / orig_w),
|
||||
max(0, y1 / orig_h),
|
||||
min(1, x2 / orig_w),
|
||||
min(1, y2 / orig_h),
|
||||
]
|
||||
|
||||
# Extract landmarks (5 points) for alignment
|
||||
if len(det) >= 15:
|
||||
landmarks = det[5:15].reshape(5, 2) / scale
|
||||
else:
|
||||
# Fallback: no landmarks, skip recognition
|
||||
continue
|
||||
|
||||
# Align face for recognition
|
||||
face_crop = _align_face(image, landmarks)
|
||||
|
||||
# SFace expects (1, 3, 112, 112) float32, BGR, normalized
|
||||
face_bgr = face_crop[:, :, ::-1].copy()
|
||||
face_blob = (face_bgr / 255.0).transpose(2, 0, 1)[np.newaxis].astype(np.float32)
|
||||
|
||||
rec_input = self._recognizer.get_inputs()[0].name
|
||||
embedding = self._recognizer.run(None, {rec_input: face_blob})[0][0]
|
||||
embedding = embedding / np.linalg.norm(embedding)
|
||||
|
||||
results.append(FaceDetection(
|
||||
bbox=bbox,
|
||||
embedding=embedding.astype(np.float32),
|
||||
quality=score,
|
||||
))
|
||||
|
||||
return results
|
||||
|
||||
@property
|
||||
def embedding_dim(self) -> int:
|
||||
return 128
|
||||
Reference in New Issue
Block a user