diff --git a/.env.example b/.env.example index 23b10b0..4d9a061 100644 --- a/.env.example +++ b/.env.example @@ -18,6 +18,8 @@ AI_PRIVACY_BATCH_WAIT_MS=4 AI_PRIVACY_MAX_PENDING_FRAMES=128 AI_PRIVACY_FACE_IMGSZ=640 AI_PRIVACY_ENABLE_IDENTITY_EXCLUSION=true +AI_PRIVACY_IDENTITY_RECHECK_INTERVAL=12 +AI_PRIVACY_IDENTITY_NEGATIVE_RECHECK_INTERVAL=90 AI_PRIVACY_ME_IMAGE_PATH=me.jpeg AI_PRIVACY_FACE_MODEL_PATH=models/yolov8n-face.pt AI_PRIVACY_FACE_DETECTION_MODEL_PATH=models/face_detection_yunet_2023mar.onnx diff --git a/README.md b/README.md index b0aeb0e..8da9aac 100644 --- a/README.md +++ b/README.md @@ -32,13 +32,19 @@ Mac 로컬 실행은 위와 같이 CPU 모드를 사용한다. 운영 환경은 처리 순서는 다음과 같다. 1. YOLO FP16 배치로 얼굴을 탐지한다. -2. 신규 또는 주기적 확인 대상 얼굴을 YuNet/SFace 배치로 식별한다. +2. YOLO가 함께 출력한 5개 랜드마크로 얼굴을 정렬하고 SFace 배치로 식별한다. + 랜드마크가 유효하지 않은 얼굴만 YuNet을 폴백으로 사용한다. 3. 참조 얼굴(`AI_PRIVACY_ME_IMAGE_PATH` 또는 `/reference-face` API)을 제외한 얼굴에 확장 타원 마스크와 Gaussian blur를 GPU에서 적용한다. 4. 처리 결과를 원래 WebRTC 트랙으로 반환한다. `GET /health`에서 장치, ONNX provider, 큐 깊이, 배치 수, 평균 배치 크기와 평균 -처리 지연을 확인할 수 있다. +처리 지연을 확인할 수 있다. 얼굴 대조 수와 지연, YOLO 랜드마크 사용 수 및 YuNet +폴백 수도 별도 지표로 제공한다. + +한 프레임에서 새로 감지된 모든 얼굴은 ROI Align과 SFace의 단일 CUDA 배치로 +즉시 대조한다. 본인 트랙은 `AI_PRIVACY_IDENTITY_RECHECK_INTERVAL`, 비본인 트랙은 더 긴 +`AI_PRIVACY_IDENTITY_NEGATIVE_RECHECK_INTERVAL` 주기로 재확인한다. ### NVIDIA RTX 3090 배포 diff --git a/SERVER_STRUCTURE.md b/SERVER_STRUCTURE.md index fefdc49..c1e6448 100644 --- a/SERVER_STRUCTURE.md +++ b/SERVER_STRUCTURE.md @@ -183,12 +183,17 @@ Status: active → closing → closed | `ai_privacy_max_pending_frames` | 128 | 전체 스트림 공유 대기열 상한 | | `ai_privacy_face_imgsz` | 640 | YOLO 입력 크기 | | `ai_privacy_enable_identity_exclusion` | true | 참조 얼굴 블러 제외 여부 | +| `ai_privacy_identity_recheck_interval` | 12 | 본인 트랙 재대조 프레임 주기 | +| `ai_privacy_identity_negative_recheck_interval` | 90 | 비본인 트랙 재대조 프레임 주기 | | `ai_privacy_me_image_path` | null | 환경변수 기반 참조 얼굴 이미지 경로 | - 앱 프로세스마다 `SharedGpuRuntime` 하나가 YOLO/YuNet/SFace 모델과 배치 워커를 소유 - WebRTC 스트림마다 `FacePrivacyFilter` 하나를 생성해 얼굴 추적 상태를 격리 - 최대 4ms 동안 모은 서로 다른 스트림 프레임을 최대 8개까지 한 번에 추론 -- CUDA에서는 YOLO FP16, ONNX Runtime CUDA I/O binding, PyTorch 마스크/블러 사용 +- YOLO의 5개 얼굴 랜드마크를 SFace 정렬에 재사용하고, 유효하지 않을 때만 YuNet 실행 +- 한 프레임의 모든 신규 얼굴 crop·정렬·SFace 추론을 단일 CUDA 배치로 즉시 처리 +- 비본인 얼굴은 본인보다 긴 주기로 재대조해 군중 장면의 반복 추론량을 제한 +- CUDA에서는 YOLO FP16, ROI Align, ONNX Runtime I/O binding·TF32·비동기 실행 사용 - `ProtectedVideoTrack` 종료 시 스트림 필터 상태를 해제하고, 처리 실패 시 검은 프레임 송신 ## 6. WebSocket 시그널링 프로토콜 diff --git a/app/core/config.py b/app/core/config.py index cada175..8b36f40 100644 --- a/app/core/config.py +++ b/app/core/config.py @@ -24,6 +24,8 @@ class Settings(BaseSettings): ai_privacy_max_pending_frames: int = 128 ai_privacy_face_imgsz: int = 640 ai_privacy_enable_identity_exclusion: bool = True + ai_privacy_identity_recheck_interval: int = 12 + ai_privacy_identity_negative_recheck_interval: int = 90 ai_privacy_me_image_path: str | None = None ai_privacy_face_model_path: str = "models/yolov8n-face.pt" ai_privacy_face_detection_model_path: str = ( diff --git a/app/services/ai/filter.py b/app/services/ai/filter.py index 81141ec..b37c807 100644 --- a/app/services/ai/filter.py +++ b/app/services/ai/filter.py @@ -61,6 +61,10 @@ def build_privacy_config(settings: Settings) -> PrivacyBlurConfig: batch_wait_ms=settings.ai_privacy_batch_wait_ms, max_pending_frames=settings.ai_privacy_max_pending_frames, enable_identity_exclusion=identity_enabled, + identity_recheck_interval=settings.ai_privacy_identity_recheck_interval, + identity_negative_recheck_interval=( + settings.ai_privacy_identity_negative_recheck_interval + ), me_image_path=reference_path or _PROJECT_ROOT / "me.jpeg", face_detection_model_path=_asset_path( settings.ai_privacy_face_detection_model_path diff --git a/privacy_blur/config.py b/privacy_blur/config.py index 36be8f9..91bdfc8 100644 --- a/privacy_blur/config.py +++ b/privacy_blur/config.py @@ -34,6 +34,7 @@ class PrivacyBlurConfig: face_match_cosine_threshold: float = 0.40 face_match_lost_cosine_threshold: float = 0.30 identity_recheck_interval: int = 12 + identity_negative_recheck_interval: int = 90 identity_min_face_size: int = 36 detect_every_n_frames: int = 1 @@ -71,6 +72,13 @@ def __post_init__(self) -> None: raise ValueError("detect_every_n_frames must be at least 1") if self.identity_detector_size < 32 or self.identity_detector_size % 32: raise ValueError("identity_detector_size must be a multiple of 32") + if self.identity_recheck_interval < 1: + raise ValueError("identity_recheck_interval must be at least 1") + if self.identity_negative_recheck_interval < self.identity_recheck_interval: + raise ValueError( + "identity_negative_recheck_interval must be at least " + "identity_recheck_interval" + ) def with_overrides(self, **kwargs: Any) -> "PrivacyBlurConfig": return replace(self, **kwargs) diff --git a/privacy_blur/core.py b/privacy_blur/core.py index f1d90b3..87933f7 100644 --- a/privacy_blur/core.py +++ b/privacy_blur/core.py @@ -48,6 +48,13 @@ def track_match_score( return max(iou, center_score * 0.75) +@dataclass +class FaceDetection: + xywh: np.ndarray + confidence: float + landmarks: np.ndarray | None = None + + @dataclass class FaceTrack: track_id: int @@ -57,6 +64,7 @@ class FaceTrack: similarity: float last_seen_frame: int last_checked_frame: int + landmarks: np.ndarray | None = None class FacePrivacyEngine: @@ -94,26 +102,93 @@ def needs_detection(self) -> bool: or self.frame_index % self.config.detect_every_n_frames == 0 ) - def _find_matching_track( + def _match_detections_to_tracks( self, - face_xywh: np.ndarray, + detections: list[FaceDetection], tracks: list[FaceTrack], - matched_track_ids: set[int], - ) -> FaceTrack | None: - best_track: FaceTrack | None = None - best_score = 0.0 - for track in tracks: - if track.track_id in matched_track_ids: - continue - score = track_match_score(face_xywh, track.xywh, self.config) - if score > best_score: - best_score = score - best_track = track - return best_track if best_score > 0.0 else None + ) -> list[FaceTrack | None]: + if not detections or not tracks: + return [None] * len(detections) + + detection_boxes = np.stack([item.xywh[:4] for item in detections]).astype( + np.float32, copy=False + ) + track_boxes = np.stack([item.xywh[:4] for item in tracks]).astype( + np.float32, copy=False + ) + detection_x1 = detection_boxes[:, None, 0] + detection_y1 = detection_boxes[:, None, 1] + detection_x2 = detection_x1 + detection_boxes[:, None, 2] + detection_y2 = detection_y1 + detection_boxes[:, None, 3] + track_x1 = track_boxes[None, :, 0] + track_y1 = track_boxes[None, :, 1] + track_x2 = track_x1 + track_boxes[None, :, 2] + track_y2 = track_y1 + track_boxes[None, :, 3] + + intersection_width = np.maximum( + 0.0, np.minimum(detection_x2, track_x2) - np.maximum(detection_x1, track_x1) + ) + intersection_height = np.maximum( + 0.0, np.minimum(detection_y2, track_y2) - np.maximum(detection_y1, track_y1) + ) + intersection = intersection_width * intersection_height + detection_area = detection_boxes[:, None, 2] * detection_boxes[:, None, 3] + track_area = track_boxes[None, :, 2] * track_boxes[None, :, 3] + union = detection_area + track_area - intersection + iou = np.divide( + intersection, + union, + out=np.zeros_like(intersection), + where=union > 1e-6, + ) + + detection_centers = detection_boxes[:, :2] + detection_boxes[:, 2:] * 0.5 + track_centers = track_boxes[:, :2] + track_boxes[:, 2:] * 0.5 + center_distance = np.linalg.norm( + detection_centers[:, None, :] - track_centers[None, :, :], axis=2 + ) + allowed_distance = np.maximum.reduce( + [ + np.broadcast_to(detection_boxes[:, None, 2], iou.shape), + np.broadcast_to(detection_boxes[:, None, 3], iou.shape), + np.broadcast_to(track_boxes[None, :, 2], iou.shape), + np.broadcast_to(track_boxes[None, :, 3], iou.shape), + ] + ) * self.config.track_center_distance_scale + center_score = np.maximum( + 0.0, + 1.0 + - np.divide( + center_distance, + allowed_distance, + out=np.full_like(center_distance, np.inf), + where=allowed_distance > 1e-6, + ), + ) + scores = np.maximum(iou, center_score * 0.75) + scores = np.where( + (iou >= self.config.track_iou_threshold) | (center_score > 0.0), + scores, + 0.0, + ) + + assignments: list[FaceTrack | None] = [] + matched_indices: set[int] = set() + for detection_scores in scores: + if matched_indices: + detection_scores = detection_scores.copy() + detection_scores[list(matched_indices)] = -1.0 + best_index = int(np.argmax(detection_scores)) + if detection_scores[best_index] > 0.0: + assignments.append(tracks[best_index]) + matched_indices.add(best_index) + else: + assignments.append(None) + return assignments def begin_frame( self, - detections: list[tuple[np.ndarray, float]] | None, + detections: list[FaceDetection] | None, ) -> list[FaceTrack]: """Update boxes and return tracks that need a batched identity check.""" if detections is None: @@ -122,14 +197,16 @@ def begin_frame( with self._lock: self.last_detection_count = len(detections) active_tracks = self._active_tracks() + matched_tracks = self._match_detections_to_tracks( + detections, active_tracks + ) matched_track_ids: set[int] = set() updated_tracks: list[FaceTrack] = [] identity_checks: list[FaceTrack] = [] - for xywh, confidence in detections: - track = self._find_matching_track( - xywh, active_tracks, matched_track_ids - ) + for detection, track in zip(detections, matched_tracks): + xywh = detection.xywh + confidence = detection.confidence needs_identity_check = False if track is None: track = FaceTrack( @@ -139,7 +216,8 @@ def begin_frame( is_me=False, similarity=-1.0, last_seen_frame=self.frame_index, - last_checked_frame=self.frame_index, + last_checked_frame=-1, + landmarks=detection.landmarks, ) self.next_track_id += 1 needs_identity_check = self.config.enable_identity_exclusion @@ -147,13 +225,20 @@ def begin_frame( track.xywh = xywh track.confidence = confidence track.last_seen_frame = self.frame_index + track.landmarks = detection.landmarks + recheck_interval = ( + self.config.identity_recheck_interval + if track.is_me + else self.config.identity_negative_recheck_interval + ) needs_identity_check = ( self.config.enable_identity_exclusion - and self.frame_index - track.last_checked_frame - >= self.config.identity_recheck_interval + and ( + track.last_checked_frame < 0 + or self.frame_index - track.last_checked_frame + >= recheck_interval + ) ) - if needs_identity_check: - track.last_checked_frame = self.frame_index if needs_identity_check: if ( @@ -174,13 +259,21 @@ def begin_frame( if track.track_id not in matched_track_ids ) self.tracks = updated_tracks - return identity_checks + return sorted( + identity_checks, + key=lambda track: ( + not track.is_me, + track.last_checked_frame >= 0, + -track.confidence, + ), + ) def apply_identity_matches(self, matches: list[tuple[FaceTrack, float]]) -> None: with self._lock: for track, similarity in matches: previous_is_me = track.is_me track.similarity = similarity + track.last_checked_frame = self.frame_index track.is_me = similarity >= self.config.face_match_cosine_threshold or ( previous_is_me and similarity >= self.config.face_match_lost_cosine_threshold diff --git a/privacy_blur/runtime.py b/privacy_blur/runtime.py index 8087dc0..cb3bb9b 100644 --- a/privacy_blur/runtime.py +++ b/privacy_blur/runtime.py @@ -18,6 +18,7 @@ import torch import torch.nn.functional as F import onnxruntime as ort +from torchvision.ops import roi_align from ultralytics import YOLO from .config import PrivacyBlurConfig @@ -29,7 +30,7 @@ from .face_detection import detect_best_yunet_face if TYPE_CHECKING: - from .core import FacePrivacyEngine, FaceTrack + from .core import FaceDetection, FacePrivacyEngine, FaceTrack LOGGER = logging.getLogger(__name__) @@ -148,6 +149,10 @@ def __init__(self, config: PrivacyBlurConfig): self.cuda_enabled = self.device.type == "cuda" self.half_enabled = self.cuda_enabled and config.use_half self._device_index = self.device.index or 0 + self._ort_run_options = ort.RunOptions() + self._ort_run_options.add_run_config_entry( + "disable_synchronize_execution_providers", "1" + ) if self.cuda_enabled: torch.cuda.set_device(self.device) @@ -189,6 +194,15 @@ def __init__(self, config: PrivacyBlurConfig): self._frames = 0 self._last_batch_size = 0 self._total_batch_latency_ms = 0.0 + self._identity_batches = 0 + self._identity_faces = 0 + self._last_identity_batch_size = 0 + self._max_identity_batch_size = 0 + self._last_identity_latency_ms = 0.0 + self._max_identity_latency_ms = 0.0 + self._identity_yolo_landmark_faces = 0 + self._identity_yunet_fallback_faces = 0 + self._total_identity_latency_ms = 0.0 self._worker = threading.Thread( target=self._worker_loop, name="privacy-blur-gpu-batcher", @@ -212,8 +226,10 @@ def _create_ort_session(self, path: Path, kind: str) -> ort.InferenceSession: if self.cuda_enabled and "CUDAExecutionProvider" in available: provider_options = { "device_id": str(self._device_index), - "cudnn_conv_algo_search": "HEURISTIC", + "cudnn_conv_algo_search": "EXHAUSTIVE", + "cudnn_conv_use_max_workspace": "1", "do_copy_in_default_stream": "1", + "use_tf32": "1", } if self._cuda_stream is not None: provider_options["user_compute_stream"] = str( @@ -417,7 +433,7 @@ def _run_batch(self, requests: list[_FrameRequest]) -> list[np.ndarray]: for index, request in enumerate(requests) if request.engine.needs_detection() ] - detections: dict[int, list[tuple[np.ndarray, float]]] = {} + detections: dict[int, list[FaceDetection]] = {} if detect_indices: detected = self._predict_yolo( [requests[index].frame for index in detect_indices] @@ -432,9 +448,14 @@ def _run_batch(self, requests: list[_FrameRequest]) -> list[np.ndarray]: ) if identity_work: + identity_started = time.perf_counter() similarities = self._identity_similarities( - [(frame, track.xywh) for _, track, frame in identity_work] + [ + (frame, track.xywh, track.landmarks) + for _, track, frame in identity_work + ] ) + identity_latency_ms = (time.perf_counter() - identity_started) * 1000.0 grouped: dict[FacePrivacyEngine, list[tuple[FaceTrack, float]]] = ( defaultdict(list) ) @@ -443,12 +464,33 @@ def _run_batch(self, requests: list[_FrameRequest]) -> list[np.ndarray]: for engine, matches in grouped.items(): engine.apply_identity_matches(matches) + landmark_faces = sum( + track.landmarks is not None for _, track, _ in identity_work + ) + with self._stats_lock: + self._identity_batches += 1 + self._identity_faces += len(identity_work) + self._last_identity_batch_size = len(identity_work) + self._max_identity_batch_size = max( + self._max_identity_batch_size, len(identity_work) + ) + self._last_identity_latency_ms = identity_latency_ms + self._max_identity_latency_ms = max( + self._max_identity_latency_ms, identity_latency_ms + ) + self._identity_yolo_landmark_faces += landmark_faces + self._identity_yunet_fallback_faces += ( + len(identity_work) - landmark_faces + ) + self._total_identity_latency_ms += identity_latency_ms privacy_boxes = [request.engine.finish_frame() for request in requests] return self._blur_frames([request.frame for request in requests], privacy_boxes) def _predict_yolo( self, frames: list[np.ndarray] - ) -> list[list[tuple[np.ndarray, float]]]: + ) -> list[list[FaceDetection]]: + from .core import FaceDetection + results = self.detector.predict( source=frames, imgsz=self.config.face_imgsz, @@ -461,17 +503,41 @@ def _predict_yolo( batch=len(frames), ) - parsed: list[list[tuple[np.ndarray, float]]] = [] + parsed: list[list[FaceDetection]] = [] for result in results: - frame_detections: list[tuple[np.ndarray, float]] = [] + frame_detections: list[FaceDetection] = [] if result.boxes is not None and len(result.boxes): - boxes = ( - torch.as_tensor(result.boxes.xyxy).detach().float().cpu().numpy() - ) - confidences = ( - torch.as_tensor(result.boxes.conf).detach().float().cpu().numpy() + boxes_tensor = torch.as_tensor(result.boxes.xyxy) + confidences_tensor = torch.as_tensor(result.boxes.conf).reshape(-1, 1) + metadata_parts = [boxes_tensor, confidences_tensor] + landmarks: np.ndarray | None = None + landmark_confidences: np.ndarray | None = None + if result.keypoints is not None and len(result.keypoints): + keypoints_tensor = torch.as_tensor(result.keypoints.xy) + metadata_parts.append( + keypoints_tensor.reshape(keypoints_tensor.shape[0], -1) + ) + if result.keypoints.conf is not None: + metadata_parts.append( + torch.as_tensor(result.keypoints.conf).reshape( + keypoints_tensor.shape[0], -1 + ) + ) + + # One device-to-host synchronization for boxes, scores, and landmarks. + metadata = ( + torch.cat(metadata_parts, dim=1).detach().float().cpu().numpy() ) - for box, confidence in zip(boxes, confidences): + boxes = metadata[:, :4] + confidences = metadata[:, 4] + if len(metadata_parts) >= 3: + landmarks = metadata[:, 5:15].reshape(-1, 5, 2) + if len(metadata_parts) == 4: + landmark_confidences = metadata[:, 15:20] + + for detection_index, (box, confidence) in enumerate( + zip(boxes, confidences) + ): x1, y1, x2, y2 = box[:4] xywh = np.array( [x1, y1, max(0.0, x2 - x1), max(0.0, y2 - y1)], @@ -485,21 +551,181 @@ def _predict_yolo( continue ratio = max(width / max(height, 1e-6), height / max(width, 1e-6)) if ratio <= self.config.max_face_aspect_ratio: - frame_detections.append((xywh, float(confidence))) + face_landmarks: np.ndarray | None = None + if landmarks is not None and detection_index < len(landmarks): + candidate = landmarks[detection_index] + confidence_ok = ( + landmark_confidences is None + or detection_index >= len(landmark_confidences) + or bool( + np.all( + landmark_confidences[detection_index] + >= self.config.identity_detector_confidence + ) + ) + ) + if ( + candidate.shape == (5, 2) + and np.isfinite(candidate).all() + and float(np.min(np.ptp(candidate, axis=0))) >= 1.0 + and confidence_ok + ): + face_landmarks = candidate.astype( + np.float32, copy=True + ) + frame_detections.append( + FaceDetection( + xywh=xywh, + confidence=float(confidence), + landmarks=face_landmarks, + ) + ) parsed.append(frame_detections) return parsed def _identity_similarities( self, - faces: list[tuple[np.ndarray, np.ndarray]], + faces: list[tuple[np.ndarray, np.ndarray, np.ndarray | None]], ) -> list[float]: if self.reference_feature is None: return [-1.0] * len(faces) - embeddings, valid = self._face_embeddings(faces) + + embeddings = torch.zeros( + (len(faces), 128), device=self.device, dtype=torch.float32 + ) + valid = torch.zeros(len(faces), device=self.device, dtype=torch.bool) + landmark_indices = [ + index for index, (_, _, landmarks) in enumerate(faces) if landmarks is not None + ] + fallback_indices = [ + index for index, (_, _, landmarks) in enumerate(faces) if landmarks is None + ] + + if landmark_indices: + landmark_embeddings, landmark_valid = ( + self._face_embeddings_from_landmarks( + [faces[index] for index in landmark_indices] + ) + ) + target = torch.as_tensor(landmark_indices, device=self.device) + embeddings[target] = landmark_embeddings + valid[target] = landmark_valid + + if fallback_indices: + fallback_embeddings, fallback_valid = self._face_embeddings( + [(faces[index][0], faces[index][1]) for index in fallback_indices] + ) + target = torch.as_tensor(fallback_indices, device=self.device) + embeddings[target] = fallback_embeddings + valid[target] = fallback_valid + scores = embeddings @ self.reference_feature scores = torch.where(valid, scores, torch.full_like(scores, -1.0)) return [float(value) for value in scores.detach().float().cpu().tolist()] + def _face_embeddings_from_landmarks( + self, + faces: list[tuple[np.ndarray, np.ndarray, np.ndarray | None]], + ) -> tuple[torch.Tensor, torch.Tensor]: + """Create all SFace inputs with batched GPU ROI Align and alignment.""" + size = 112 + embeddings = torch.zeros( + (len(faces), 128), device=self.device, dtype=torch.float32 + ) + valid = torch.zeros(len(faces), device=self.device, dtype=torch.bool) + grouped: dict[ + tuple[int, int], + list[tuple[int, np.ndarray, np.ndarray, np.ndarray]], + ] = defaultdict(list) + for index, (frame, box, landmarks) in enumerate(faces): + if landmarks is not None: + grouped[frame.shape[:2]].append((index, frame, box, landmarks)) + + stream_context = ( + torch.cuda.stream(self._cuda_stream) + if self._cuda_stream is not None + else torch.no_grad() + ) + with stream_context: + crop_batches: list[torch.Tensor] = [] + landmark_batches: list[torch.Tensor] = [] + all_source_indices: list[int] = [] + for (frame_h, frame_w), group in grouped.items(): + unique_frames: list[np.ndarray] = [] + frame_indices: dict[int, int] = {} + rois: list[list[float]] = [] + mapped_landmarks: list[np.ndarray] = [] + source_indices: list[int] = [] + + for source_index, frame, box, landmarks in group: + frame_key = id(frame) + frame_index = frame_indices.get(frame_key) + if frame_index is None: + frame_index = len(unique_frames) + frame_indices[frame_key] = frame_index + unique_frames.append(frame) + + x, y, width, height = box[:4].astype(float) + pad_x, pad_y = width * 0.35, height * 0.35 + x1 = max(0, min(frame_w - 1, int(round(x - pad_x)))) + y1 = max(0, min(frame_h - 1, int(round(y - pad_y)))) + x2 = max(0, min(frame_w, int(round(x + width + pad_x)))) + y2 = max(0, min(frame_h, int(round(y + height + pad_y)))) + if x2 <= x1 or y2 <= y1: + continue + + rois.append([float(frame_index), x1, y1, x2, y2]) + scale = np.array( + [size / float(x2 - x1), size / float(y2 - y1)], + dtype=np.float32, + ) + # roi_align(aligned=True) shifts continuous ROI coordinates + # by -0.5 before sampling. Apply the same transform to points. + mapped = ( + landmarks.astype(np.float32, copy=False) + - np.array([x1 - 0.5, y1 - 0.5], dtype=np.float32) + ) * scale - 0.5 + mapped_landmarks.append(mapped) + source_indices.append(source_index) + + if not rois: + continue + + frame_batch = torch.from_numpy( + np.ascontiguousarray(np.stack(unique_frames)) + ).to(self.device) + frame_batch = frame_batch.permute(0, 3, 1, 2).float() + roi_batch = torch.as_tensor( + rois, device=self.device, dtype=torch.float32 + ) + crop_batch = roi_align( + frame_batch, + roi_batch, + output_size=(size, size), + spatial_scale=1.0, + sampling_ratio=2, + aligned=True, + ) + landmarks_batch = torch.from_numpy( + np.stack(mapped_landmarks) + ).to(self.device) + crop_batches.append(crop_batch) + landmark_batches.append(landmarks_batch) + all_source_indices.extend(source_indices) + + if crop_batches: + crop_batch = torch.cat(crop_batches, dim=0) + landmarks_batch = torch.cat(landmark_batches, dim=0) + aligned = self._align_faces(crop_batch, landmarks_batch) + aligned_rgb = aligned[:, [2, 1, 0], :, :].contiguous() + features = F.normalize(self._run_sface(aligned_rgb).float(), dim=1) + mapped_indices = torch.as_tensor( + all_source_indices, device=self.device + ) + embeddings[mapped_indices] = features + valid[mapped_indices] = True + return embeddings, valid + def _face_embeddings( self, faces: list[tuple[np.ndarray, np.ndarray]], @@ -587,12 +813,19 @@ def _yunet_landmarks( keypoints = outputs[level + 9].reshape(batch_size, anchors, 5, 2) score = torch.sqrt(cls * obj).squeeze(-1) - grid_y, grid_x = torch.meshgrid( - torch.arange(rows, device=self.device, dtype=torch.float32), - torch.arange(rows, device=self.device, dtype=torch.float32), - indexing="ij", - ) - grid = torch.stack((grid_x.flatten(), grid_y.flatten()), dim=-1) + grid_cache = getattr(self, "_yunet_grid_cache", None) + if grid_cache is None: + grid_cache = {} + self._yunet_grid_cache = grid_cache + grid = grid_cache.get((size, stride)) + if grid is None: + grid_y, grid_x = torch.meshgrid( + torch.arange(rows, device=self.device, dtype=torch.float32), + torch.arange(rows, device=self.device, dtype=torch.float32), + indexing="ij", + ) + grid = torch.stack((grid_x.flatten(), grid_y.flatten()), dim=-1) + grid_cache[(size, stride)] = grid width = torch.exp(bbox[..., 2].clamp(max=10)) * stride height = torch.exp(bbox[..., 3].clamp(max=10)) * stride rank = width * height * score @@ -670,62 +903,87 @@ def _run_ort( shape=tuple(output.shape), buffer_ptr=output.data_ptr(), ) - session.run_with_iobinding(binding) + session.run_with_iobinding(binding, self._ort_run_options) return outputs def _align_faces( self, crops: torch.Tensor, landmarks: torch.Tensor ) -> torch.Tensor: - destination = torch.tensor( - [ - [38.2946, 51.6963], - [73.5318, 51.5014], - [56.0252, 71.7366], - [41.5493, 92.3655], - [70.7299, 92.2041], - ], - device=self.device, - dtype=torch.float32, - ).expand(landmarks.shape[0], -1, -1) + destination_template = getattr(self, "_alignment_destination", None) + if destination_template is None: + destination_template = torch.tensor( + [ + [38.2946, 51.6963], + [73.5318, 51.5014], + [56.0252, 71.7366], + [41.5493, 92.3655], + [70.7299, 92.2041], + ], + device=self.device, + dtype=torch.float32, + ).unsqueeze(0) + self._alignment_destination = destination_template + destination = destination_template.expand(landmarks.shape[0], -1, -1) source_mean = landmarks.mean(dim=1, keepdim=True) destination_mean = destination.mean(dim=1, keepdim=True) source_centered = landmarks - source_mean destination_centered = destination - destination_mean - covariance = destination_centered.transpose(1, 2) @ source_centered / 5.0 - u, singular_values, vh = torch.linalg.svd(covariance) - signs = torch.ones_like(singular_values) - signs[:, 1] = torch.where( - torch.linalg.det(covariance) < 0, - -torch.ones_like(signs[:, 1]), - torch.ones_like(signs[:, 1]), + source_x, source_y = source_centered.unbind(dim=2) + destination_x, destination_y = destination_centered.unbind(dim=2) + denominator = source_centered.square().sum(dim=(1, 2)).clamp_min(1e-8) + coefficient_a = ( + source_x * destination_x + source_y * destination_y + ).sum(dim=1) / denominator + coefficient_b = ( + source_x * destination_y - source_y * destination_x + ).sum(dim=1) / denominator + + source_mean_xy = source_mean[:, 0] + destination_mean_xy = destination_mean[:, 0] + translation_x = ( + destination_mean_xy[:, 0] + - coefficient_a * source_mean_xy[:, 0] + + coefficient_b * source_mean_xy[:, 1] ) - rotation = u @ torch.diag_embed(signs) @ vh - variance = source_centered.square().sum(dim=(1, 2)) / 5.0 - scale = (singular_values * signs).sum(dim=1) / variance.clamp_min(1e-8) - translation = destination_mean[:, 0] - ( - scale[:, None] * (rotation @ source_mean.transpose(1, 2))[:, :, 0] + translation_y = ( + destination_mean_xy[:, 1] + - coefficient_b * source_mean_xy[:, 0] + - coefficient_a * source_mean_xy[:, 1] ) - transform = torch.zeros((landmarks.shape[0], 3, 3), device=self.device) - transform[:, :2, :2] = scale[:, None, None] * rotation - transform[:, :2, 2] = translation - transform[:, 2, 2] = 1.0 - inverse = torch.linalg.inv(transform) + determinant = ( + coefficient_a.square() + coefficient_b.square() + ).clamp_min(1e-8) + inverse_a = coefficient_a / determinant + inverse_b = coefficient_b / determinant output_size = 112 - y, x = torch.meshgrid( - torch.arange(output_size, device=self.device, dtype=torch.float32), - torch.arange(output_size, device=self.device, dtype=torch.float32), - indexing="ij", - ) - homogeneous = torch.stack((x, y, torch.ones_like(x)), dim=-1).reshape(-1, 3) - source = torch.einsum("bij,pj->bpi", inverse, homogeneous) crop_size = crops.shape[-1] - grid_x = source[..., 0] * (2.0 / (crop_size - 1)) - 1.0 - grid_y = source[..., 1] * (2.0 / (crop_size - 1)) - 1.0 - grid = torch.stack((grid_x, grid_y), dim=-1).reshape( - -1, output_size, output_size, 2 + output_scale = (output_size - 1) * 0.5 + input_scale = (crop_size - 1) * 0.5 + linear_scale = output_scale / input_scale + theta = torch.empty( + (landmarks.shape[0], 2, 3), + device=self.device, + dtype=torch.float32, + ) + theta[:, 0, 0] = inverse_a * linear_scale + theta[:, 0, 1] = inverse_b * linear_scale + theta[:, 1, 0] = -inverse_b * linear_scale + theta[:, 1, 1] = inverse_a * linear_scale + offset_x = output_scale - translation_x + offset_y = output_scale - translation_y + theta[:, 0, 2] = ( + inverse_a * offset_x + inverse_b * offset_y + ) / input_scale - 1.0 + theta[:, 1, 2] = ( + -inverse_b * offset_x + inverse_a * offset_y + ) / input_scale - 1.0 + grid = F.affine_grid( + theta, + (landmarks.shape[0], crops.shape[1], output_size, output_size), + align_corners=True, ) return F.grid_sample( crops, @@ -824,6 +1082,16 @@ def stats(self) -> dict[str, Any]: average_latency = ( self._total_batch_latency_ms / self._batches if self._batches else 0.0 ) + average_identity_batch = ( + self._identity_faces / self._identity_batches + if self._identity_batches + else 0.0 + ) + average_identity_latency = ( + self._total_identity_latency_ms / self._identity_batches + if self._identity_batches + else 0.0 + ) return { "device": str(self.device), "cuda": self.cuda_enabled, @@ -839,6 +1107,20 @@ def stats(self) -> dict[str, Any]: "last_batch_size": self._last_batch_size, "average_batch_size": round(average_batch, 2), "average_batch_latency_ms": round(average_latency, 2), + "identity_batches": self._identity_batches, + "identity_faces": self._identity_faces, + "last_identity_batch_size": self._last_identity_batch_size, + "max_identity_batch_size": self._max_identity_batch_size, + "last_identity_latency_ms": round( + self._last_identity_latency_ms, 2 + ), + "max_identity_latency_ms": round( + self._max_identity_latency_ms, 2 + ), + "identity_yolo_landmark_faces": self._identity_yolo_landmark_faces, + "identity_yunet_fallback_faces": self._identity_yunet_fallback_faces, + "average_identity_batch_size": round(average_identity_batch, 2), + "average_identity_latency_ms": round(average_identity_latency, 2), } def close(self) -> None: diff --git a/requirements.txt b/requirements.txt index 6406685..cd62e22 100644 --- a/requirements.txt +++ b/requirements.txt @@ -12,6 +12,7 @@ numpy>=1.24 opencv-python-headless>=4.8 ultralytics>=8.0 torch>=2.0 +torchvision>=0.15 onnx>=1.17 onnxruntime-gpu==1.24.3; sys_platform == "linux" onnxruntime==1.24.3; sys_platform != "linux" diff --git a/tests/test_filter.py b/tests/test_filter.py index 93d55c4..4f11e65 100644 --- a/tests/test_filter.py +++ b/tests/test_filter.py @@ -54,6 +54,8 @@ def test_runtime_config_passed_correctly( monkeypatch.setenv("AI_PRIVACY_REQUIRE_GPU", "true") monkeypatch.setenv("AI_PRIVACY_MAX_BATCH_SIZE", "16") monkeypatch.setenv("AI_PRIVACY_BATCH_WAIT_MS", "2.5") + monkeypatch.setenv("AI_PRIVACY_IDENTITY_RECHECK_INTERVAL", "6") + monkeypatch.setenv("AI_PRIVACY_IDENTITY_NEGATIVE_RECHECK_INTERVAL", "60") from app.core.config import get_settings get_settings.cache_clear() @@ -67,6 +69,8 @@ def test_runtime_config_passed_correctly( assert kwargs["require_gpu"] is True assert kwargs["max_batch_size"] == 16 assert kwargs["batch_wait_ms"] == 2.5 + assert kwargs["identity_recheck_interval"] == 6 + assert kwargs["identity_negative_recheck_interval"] == 60 assert kwargs["face_model_path"].is_absolute() diff --git a/tests/test_identity_performance.py b/tests/test_identity_performance.py new file mode 100644 index 0000000..be98939 --- /dev/null +++ b/tests/test_identity_performance.py @@ -0,0 +1,145 @@ +from unittest.mock import MagicMock + +import numpy as np +import torch + +from privacy_blur.config import PrivacyBlurConfig +from privacy_blur.core import ( + FaceDetection, + FacePrivacyEngine, + FaceTrack, + track_match_score, +) +from privacy_blur.runtime import SharedGpuRuntime + + +def _detection(confidence: float = 0.9) -> FaceDetection: + return FaceDetection( + xywh=np.array([20, 20, 80, 80], dtype=np.float32), + confidence=confidence, + landmarks=np.array( + [[45, 50], [75, 50], [60, 65], [48, 82], [72, 82]], + dtype=np.float32, + ), + ) + + +def _engine(**overrides) -> FacePrivacyEngine: + config = PrivacyBlurConfig( + enable_identity_exclusion=True, + identity_recheck_interval=3, + identity_negative_recheck_interval=10, + **overrides, + ) + return FacePrivacyEngine(config=config, runtime=MagicMock()) + + +def test_deferred_new_track_remains_due_until_compared() -> None: + engine = _engine() + + first = engine.begin_frame([_detection()]) + engine.finish_frame() + second = engine.begin_frame([_detection()]) + + assert len(first) == 1 + assert second[0] is first[0] + assert second[0].last_checked_frame == -1 + + +def test_non_reference_track_uses_longer_recheck_interval() -> None: + engine = _engine() + track = engine.begin_frame([_detection()])[0] + engine.apply_identity_matches([(track, 0.1)]) + engine.finish_frame() + + for _ in range(9): + assert engine.begin_frame([_detection()]) == [] + engine.finish_frame() + + assert engine.frame_index == 10 + assert engine.begin_frame([_detection()]) == [track] + + +def test_reference_track_keeps_short_recheck_interval() -> None: + engine = _engine() + track = engine.begin_frame([_detection()])[0] + engine.apply_identity_matches([(track, 0.9)]) + engine.finish_frame() + + for _ in range(2): + assert engine.begin_frame([_detection()]) == [] + engine.finish_frame() + + assert engine.frame_index == 3 + assert engine.begin_frame([_detection()]) == [track] + + +def test_vectorized_track_assignment_matches_scalar_greedy_result() -> None: + engine = _engine() + rng = np.random.default_rng(42) + boxes = rng.uniform(0, 500, size=(40, 4)).astype(np.float32) + boxes[:, 2:] = rng.uniform(30, 100, size=(40, 2)) + tracks = [ + FaceTrack( + track_id=index, + xywh=box, + confidence=0.9, + is_me=False, + similarity=0.0, + last_seen_frame=0, + last_checked_frame=0, + ) + for index, box in enumerate(boxes) + ] + shifted = boxes + rng.normal(0, 5, size=boxes.shape).astype(np.float32) + shifted[:, 2:] = boxes[:, 2:] + detections = [ + FaceDetection(xywh=box, confidence=0.9) for box in shifted[::-1] + ] + + expected: list[FaceTrack | None] = [] + matched_ids: set[int] = set() + for detection in detections: + candidates = [ + (track_match_score(detection.xywh, track.xywh, engine.config), track) + for track in tracks + if track.track_id not in matched_ids + ] + score, track = max(candidates, key=lambda item: item[0]) + expected.append(track if score > 0 else None) + if score > 0: + matched_ids.add(track.track_id) + + actual = engine._match_detections_to_tracks(detections, tracks) + + assert [track.track_id if track else None for track in actual] == [ + track.track_id if track else None for track in expected + ] + + +def test_all_detector_landmarks_use_one_sface_batch_without_yunet() -> None: + runtime = SharedGpuRuntime.__new__(SharedGpuRuntime) + runtime.device = torch.device("cpu") + runtime._cuda_stream = None + runtime.reference_feature = torch.nn.functional.normalize( + torch.ones(128, dtype=torch.float32), dim=0 + ) + runtime._run_sface = MagicMock( + side_effect=lambda inputs: torch.ones( + (inputs.shape[0], 128), dtype=torch.float32 + ) + ) + runtime._face_embeddings = MagicMock( + side_effect=AssertionError("YuNet fallback should not run") + ) + frame = np.zeros((120, 120, 3), dtype=np.uint8) + detection = _detection() + + faces = [(frame, detection.xywh, detection.landmarks)] * 32 + scores = runtime._identity_similarities(faces) + + assert len(scores) == 32 + assert min(scores) > 0.99 + runtime._run_sface.assert_called_once() + assert runtime._run_sface.call_args.args[0].shape == (32, 3, 112, 112) + runtime._face_embeddings.assert_not_called()