5e2094f2ad
FastAPI service in infra/media-worker implementing every MediaInferenceProvider operation behind HTTP API v1, plus capabilities. Bearer auth, MEDIA_ROOT realpath enforcement, strict request schemas, cancellation on disconnect, GPU/CPU gates with 503 overflow, input size/duration limits, lazy SHA256-pinned model downloads, fail-closed redaction with atomic output, offline pytest suite with fake backends, CPU/CUDA/ROCm Dockerfile and a license review of every model. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_0138YqDj3ttfayRwuMKKrqVj
266 lines
10 KiB
Python
266 lines
10 KiB
Python
"""detectFaces, trackFaces, faceLandmarks.
|
|
|
|
Detection is OpenCV YuNet (MIT). Tracking is a ByteTrack-style IoU tracker
|
|
implemented here: high-confidence detections associate first, low-confidence
|
|
ones only extend existing tracks, and a track ends after 500 ms without a match.
|
|
Identity embeddings (for local person clustering) come from an opt-in backend
|
|
and are null otherwise.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import math
|
|
import threading
|
|
from dataclasses import dataclass, field
|
|
from typing import Any
|
|
|
|
import cv2
|
|
import numpy as np
|
|
|
|
from .. import media_io
|
|
from ..schemas import DetectFacesRequest, TimesRequest, TrackFacesRequest
|
|
from .common import Ctx, clamp01, norm_box, norm_point, open_input, unique_sorted
|
|
|
|
DETECT_MAX_SIDE = 1280
|
|
DETECT_THRESHOLD = 0.6
|
|
TRACK_LOW_THRESHOLD = 0.3
|
|
TRACK_MATCH_IOU = 0.3
|
|
TRACK_MAX_GAP_MS = 500
|
|
|
|
|
|
@dataclass
|
|
class Detection:
|
|
"""Pixel-space face: box (x, y, w, h), score, five landmarks (x, y)."""
|
|
|
|
box: tuple[float, float, float, float]
|
|
score: float
|
|
landmarks: list[tuple[float, float]]
|
|
|
|
|
|
class YuNetDetector:
|
|
def __init__(self, model_path: str) -> None:
|
|
self.model_path = model_path
|
|
self._local = threading.local()
|
|
|
|
def _detector(self, width: int, height: int) -> Any:
|
|
detector = getattr(self._local, "detector", None)
|
|
if detector is None:
|
|
detector = cv2.FaceDetectorYN.create(
|
|
self.model_path, "", (width, height), TRACK_LOW_THRESHOLD, 0.3, 5000
|
|
)
|
|
self._local.detector = detector
|
|
detector.setInputSize((width, height))
|
|
return detector
|
|
|
|
def detect(self, frame: np.ndarray) -> list[Detection]:
|
|
height, width = frame.shape[:2]
|
|
_, faces = self._detector(width, height).detect(frame)
|
|
result: list[Detection] = []
|
|
for row in faces if faces is not None else []:
|
|
points = [(float(row[4 + 2 * i]), float(row[5 + 2 * i])) for i in range(5)]
|
|
result.append(
|
|
Detection(
|
|
(float(row[0]), float(row[1]), float(row[2]), float(row[3])), float(row[14]), points
|
|
)
|
|
)
|
|
return result
|
|
|
|
|
|
def load_face_detector(files: dict[str, str], _settings: Any) -> YuNetDetector:
|
|
return YuNetDetector(files["face_detection_yunet_2023mar.onnx"])
|
|
|
|
|
|
class SFaceEmbedder:
|
|
def __init__(self, model_path: str) -> None:
|
|
self.model_path = model_path
|
|
self._local = threading.local()
|
|
|
|
def embed(self, frame: np.ndarray, detection: Detection) -> np.ndarray:
|
|
recognizer = getattr(self._local, "recognizer", None)
|
|
if recognizer is None:
|
|
recognizer = cv2.FaceRecognizerSF.create(self.model_path, "")
|
|
self._local.recognizer = recognizer
|
|
row = np.array(
|
|
[*detection.box, *[c for p in detection.landmarks for c in p], detection.score],
|
|
dtype=np.float32,
|
|
)
|
|
aligned = recognizer.alignCrop(frame, row)
|
|
return recognizer.feature(aligned).reshape(-1).astype(np.float32)
|
|
|
|
|
|
def load_identity_embedder(files: dict[str, str], _settings: Any) -> SFaceEmbedder:
|
|
return SFaceEmbedder(files["face_recognition_sface_2021dec.onnx"])
|
|
|
|
|
|
def _observation(det: Detection, width: int, height: int) -> dict[str, Any]:
|
|
return {
|
|
"box": norm_box(*det.box, width, height),
|
|
"confidence": clamp01(det.score),
|
|
"landmarks": [norm_point(x, y, width, height) for x, y in det.landmarks],
|
|
}
|
|
|
|
|
|
def detect_run(ctx: Ctx, body: DetectFacesRequest, cancel: threading.Event) -> dict[str, Any]:
|
|
real, info = open_input(ctx, body.path, cancel)
|
|
detector = ctx.registry.runtime("detectFaces")
|
|
frames = (
|
|
media_io.frames_at(real, info, unique_sorted(body.timesMs), cancel, DETECT_MAX_SIDE)
|
|
if body.timesMs is not None
|
|
else media_io.stream_frames(real, info, cancel, body.sampleFps, DETECT_MAX_SIDE)
|
|
)
|
|
out = []
|
|
for t_ms, frame in frames:
|
|
height, width = frame.shape[:2]
|
|
faces = [
|
|
_observation(d, width, height) for d in detector.detect(frame) if d.score >= DETECT_THRESHOLD
|
|
]
|
|
out.append({"tMs": t_ms, "faces": faces})
|
|
return {"frames": out}
|
|
|
|
|
|
def landmarks_run(ctx: Ctx, body: TimesRequest, cancel: threading.Event) -> dict[str, Any]:
|
|
real, info = open_input(ctx, body.path, cancel)
|
|
detector = ctx.registry.runtime("faceLandmarks")
|
|
out = []
|
|
for t_ms, frame in media_io.frames_at(real, info, unique_sorted(body.timesMs), cancel, DETECT_MAX_SIDE):
|
|
height, width = frame.shape[:2]
|
|
faces = [
|
|
{
|
|
"box": norm_box(*d.box, width, height),
|
|
"landmarks": [norm_point(x, y, width, height) for x, y in d.landmarks],
|
|
}
|
|
for d in detector.detect(frame)
|
|
if d.score >= DETECT_THRESHOLD
|
|
]
|
|
out.append({"tMs": t_ms, "faces": faces})
|
|
return {"frames": out}
|
|
|
|
|
|
def iou(a: tuple[float, float, float, float], b: tuple[float, float, float, float]) -> float:
|
|
ax1, ay1, ax2, ay2 = a[0], a[1], a[0] + a[2], a[1] + a[3]
|
|
bx1, by1, bx2, by2 = b[0], b[1], b[0] + b[2], b[1] + b[3]
|
|
iw = max(0.0, min(ax2, bx2) - max(ax1, bx1))
|
|
ih = max(0.0, min(ay2, by2) - max(ay1, by1))
|
|
inter = iw * ih
|
|
union = a[2] * a[3] + b[2] * b[3] - inter
|
|
return inter / union if union > 0 else 0.0
|
|
|
|
|
|
@dataclass
|
|
class Track:
|
|
track_id: str
|
|
box: tuple[float, float, float, float]
|
|
velocity: tuple[float, float] = (0.0, 0.0)
|
|
last_ms: int = 0
|
|
keyframes: list[dict[str, Any]] = field(default_factory=list)
|
|
# Best (score, frame crop source, detection) for identity embeddings.
|
|
best: list[tuple[float, np.ndarray, Detection]] = field(default_factory=list)
|
|
|
|
def predicted(self, t_ms: int) -> tuple[float, float, float, float]:
|
|
dt = t_ms - self.last_ms
|
|
return (
|
|
self.box[0] + self.velocity[0] * dt,
|
|
self.box[1] + self.velocity[1] * dt,
|
|
self.box[2],
|
|
self.box[3],
|
|
)
|
|
|
|
|
|
class IouTracker:
|
|
"""Two-stage greedy IoU association (ByteTrack's idea without the Kalman filter)."""
|
|
|
|
def __init__(self, max_gap_ms: int = TRACK_MAX_GAP_MS) -> None:
|
|
self.max_gap_ms = max_gap_ms
|
|
self.active: list[Track] = []
|
|
self.finished: list[Track] = []
|
|
self._next = 1
|
|
|
|
def _associate(
|
|
self, tracks: list[Track], dets: list[Detection], t_ms: int
|
|
) -> tuple[list[tuple[Track, Detection]], list[Track], list[Detection]]:
|
|
pairs = sorted(
|
|
(
|
|
(iou(tr.predicted(t_ms), d.box), ti, di)
|
|
for ti, tr in enumerate(tracks)
|
|
for di, d in enumerate(dets)
|
|
),
|
|
reverse=True,
|
|
)
|
|
used_t: set[int] = set()
|
|
used_d: set[int] = set()
|
|
matches: list[tuple[Track, Detection]] = []
|
|
for score, ti, di in pairs:
|
|
if score < TRACK_MATCH_IOU:
|
|
break
|
|
if ti in used_t or di in used_d:
|
|
continue
|
|
used_t.add(ti)
|
|
used_d.add(di)
|
|
matches.append((tracks[ti], dets[di]))
|
|
rest_t = [t for i, t in enumerate(tracks) if i not in used_t]
|
|
rest_d = [d for i, d in enumerate(dets) if i not in used_d]
|
|
return matches, rest_t, rest_d
|
|
|
|
def update(
|
|
self, t_ms: int, dets: list[Detection], frame: np.ndarray | None, keep_crops: bool
|
|
) -> list[tuple[Track, Detection]]:
|
|
expired = [t for t in self.active if t_ms - t.last_ms > self.max_gap_ms]
|
|
self.finished.extend(expired)
|
|
self.active = [t for t in self.active if t not in expired]
|
|
high = [d for d in dets if d.score >= DETECT_THRESHOLD]
|
|
low = [d for d in dets if TRACK_LOW_THRESHOLD <= d.score < DETECT_THRESHOLD]
|
|
matches, rest_tracks, rest_high = self._associate(self.active, high, t_ms)
|
|
low_matches, _, _ = self._associate(rest_tracks, low, t_ms)
|
|
matches += low_matches
|
|
for track, det in matches:
|
|
dt = max(t_ms - track.last_ms, 1)
|
|
track.velocity = ((det.box[0] - track.box[0]) / dt, (det.box[1] - track.box[1]) / dt)
|
|
track.box = det.box
|
|
track.last_ms = t_ms
|
|
for det in rest_high:
|
|
track = Track(f"t{self._next}", det.box, last_ms=t_ms)
|
|
self._next += 1
|
|
self.active.append(track)
|
|
matches.append((track, det))
|
|
if keep_crops and frame is not None:
|
|
for track, det in matches:
|
|
if det.score >= DETECT_THRESHOLD:
|
|
track.best.append((det.score, frame, det))
|
|
track.best.sort(key=lambda item: -item[0])
|
|
del track.best[3:]
|
|
return matches
|
|
|
|
def close(self) -> list[Track]:
|
|
tracks = self.finished + self.active
|
|
self.active, self.finished = [], []
|
|
return sorted(tracks, key=lambda t: int(t.track_id[1:]))
|
|
|
|
|
|
def track_run(ctx: Ctx, body: TrackFacesRequest, cancel: threading.Event) -> dict[str, Any]:
|
|
real, info = open_input(ctx, body.path, cancel)
|
|
detector = ctx.registry.runtime("trackFaces")
|
|
embedder = (
|
|
ctx.registry.optional_runtime("faceIdentity")
|
|
if ctx.registry.optional_enabled("faceIdentity")
|
|
else None
|
|
)
|
|
tracker = IouTracker()
|
|
for t_ms, frame in media_io.stream_frames(real, info, cancel, body.sampleFps, DETECT_MAX_SIDE):
|
|
height, width = frame.shape[:2]
|
|
matches = tracker.update(t_ms, detector.detect(frame), frame, embedder is not None)
|
|
for track, det in matches:
|
|
track.keyframes.append(
|
|
{"tMs": t_ms, "box": norm_box(*det.box, width, height), "confidence": clamp01(det.score)}
|
|
)
|
|
tracks = []
|
|
for track in tracker.close():
|
|
embedding = None
|
|
if embedder is not None and track.best:
|
|
vectors = np.stack([embedder.embed(frame, det) for _, frame, det in track.best])
|
|
mean = vectors.mean(axis=0)
|
|
norm = float(np.linalg.norm(mean))
|
|
if norm > 0 and math.isfinite(norm):
|
|
embedding = [float(v) for v in mean / norm]
|
|
tracks.append({"trackId": track.track_id, "keyframes": track.keyframes, "embedding": embedding})
|
|
return {"sampleFps": body.sampleFps, "tracks": tracks}
|