From b80b92d8f832ef74748471e05a4efdbe9eb35416 Mon Sep 17 00:00:00 2001 From: RyanOnTheInside <7623207+ryanontheinside@users.noreply.github.com> Date: Fri, 16 May 2025 15:21:50 -0400 Subject: [PATCH] face detector --- .../face_detection/detector.py | 133 +++++++----------- 1 file changed, 52 insertions(+), 81 deletions(-) diff --git a/src/mediapipe_vision/face_detection/detector.py b/src/mediapipe_vision/face_detection/detector.py index f9919c3..d8f76d4 100644 --- a/src/mediapipe_vision/face_detection/detector.py +++ b/src/mediapipe_vision/face_detection/detector.py @@ -3,7 +3,7 @@ This module contains the implementation of MediaPipe Face Detection functionality. """ -from typing import List +from typing import Any, List import mediapipe as mp import numpy as np @@ -13,12 +13,11 @@ from mediapipe.tasks.python import vision from mediapipe.tasks.python.core.base_options import BaseOptions from ...utils.timing import TimestampProvider # Import from utils - -# Import new types +from ..common import BaseDetector from ..types import BoundingBox, FaceDetectionResult, FaceKeypoint -class FaceDetector: +class FaceDetector(BaseDetector[FaceDetectionResult]): """Detects faces in an image using MediaPipe FaceDetector.""" def __init__(self, model_path: str): @@ -35,90 +34,62 @@ class FaceDetector: self._current_options = None # Store current options for comparison self._timestamp_provider = None # Initialize TimestampProvider later - def _initialize_detector(self, min_detection_confidence: float, running_mode: str, delegate: str): - """Initializes the FaceDetector.""" - # Map string mode to MediaPipe enum - mode_enum = vision.RunningMode.IMAGE if running_mode == "image" else vision.RunningMode.VIDEO - delegate_enum = BaseOptions.Delegate.CPU if delegate.lower() == "cpu" else BaseOptions.Delegate.GPU - - base_options = python.BaseOptions(model_asset_path=self.model_path, delegate=delegate_enum) - options = vision.FaceDetectorOptions( + def _create_detector_options(self, base_options: python.BaseOptions, + mode_enum: vision.RunningMode, **kwargs) -> vision.FaceDetectorOptions: + """Create FaceDetector-specific options with parameters: + - min_detection_confidence: Minimum confidence for face detection + """ + return vision.FaceDetectorOptions( base_options=base_options, - running_mode=mode_enum, # Use enum mode - min_detection_confidence=min_detection_confidence, + running_mode=mode_enum, + min_detection_confidence=kwargs.get('min_detection_confidence', 0.5), ) - self._current_options = options # Store for comparison - # Only create timestamp provider for VIDEO mode - self._timestamp_provider = TimestampProvider() if mode_enum == vision.RunningMode.VIDEO else None + + def _create_detector_instance(self, options: vision.FaceDetectorOptions) -> vision.FaceDetector: return vision.FaceDetector.create_from_options(options) + def _get_options_tuple(self, running_mode: str = None, delegate: str = None, **kwargs) -> tuple: + if self._current_options: + return ( + self._current_options.min_detection_confidence, + running_mode, + delegate, + ) + else: + return ( + kwargs.get('min_detection_confidence', 0.5), + running_mode, + delegate, + ) + + def _process_detection_result(self, detection_result: Any) -> List[FaceDetectionResult]: + current_image_detections = [] + if detection_result and detection_result.detections: + for detection in detection_result.detections: + bbox_mp = detection.bounding_box + keypoints_mp = detection.keypoints + score = detection.categories[0].score if detection.categories else None + + bbox = BoundingBox(origin_x=bbox_mp.origin_x, origin_y=bbox_mp.origin_y, width=bbox_mp.width, height=bbox_mp.height) + + keypoints = None + if keypoints_mp: + keypoints = [FaceKeypoint(label=kp.label, x=kp.x, y=kp.y) for kp in keypoints_mp] + + current_image_detections.append(FaceDetectionResult(bounding_box=bbox, keypoints=keypoints, score=score)) + return current_image_detections + def detect( self, image: torch.Tensor, min_detection_confidence: float = 0.5, running_mode: str = "video", delegate: str = "cpu", - ) -> List[List[FaceDetectionResult]]: # Added running_mode - """Detects faces in the input image tensor. - - Returns: - A list (batch) of lists (detections per image) of FaceDetectionResult objects. - """ - if image.dim() != 4: - raise ValueError("Input tensor must be in BHWC format.") - - # Check if options changed (model path handled implicitly by node) - # Add running_mode to tuple check - new_options_tuple = (min_detection_confidence, running_mode, delegate.lower()) - current_options_tuple = None - if self._current_options: - current_options_tuple = ( - self._current_options.min_detection_confidence, - "image" if self._current_options.running_mode == vision.RunningMode.IMAGE else "video", # Map enum back to string - "cpu" if self._current_options.base_options.delegate == BaseOptions.Delegate.CPU else "gpu", - ) - - if self._detector_instance is None or new_options_tuple != current_options_tuple: - # Close previous instance if it exists and has a close method - if self._detector_instance and hasattr(self._detector_instance, "close"): - self._detector_instance.close() - self._detector_instance = self._initialize_detector(min_detection_confidence, running_mode, delegate) - - batch_size = image.shape[0] - batch_results = [] - # Determine the correct detection function based on mode - is_video_mode = self._current_options.running_mode == vision.RunningMode.VIDEO - detect_func = self._detector_instance.detect_for_video if is_video_mode else self._detector_instance.detect - - for i in range(batch_size): - img_tensor = image[i] # HWC - np_image = (img_tensor.cpu().numpy() * 255).astype(np.uint8) - mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=np_image) - - if is_video_mode: - # Get timestamp for video mode - timestamp_ms = self._timestamp_provider.next() - detection_result = detect_func(mp_image, timestamp_ms) - else: - # Call image mode function - detection_result = detect_func(mp_image) - - current_image_detections = [] - if detection_result and detection_result.detections: # Check if result is not None - for detection in detection_result.detections: - bbox_mp = detection.bounding_box - keypoints_mp = detection.keypoints - score = detection.categories[0].score if detection.categories else None # Face detection usually has one score - - bbox = BoundingBox(origin_x=bbox_mp.origin_x, origin_y=bbox_mp.origin_y, width=bbox_mp.width, height=bbox_mp.height) - - keypoints = None - if keypoints_mp: - keypoints = [FaceKeypoint(label=kp.label, x=kp.x, y=kp.y) for kp in keypoints_mp] - - current_image_detections.append(FaceDetectionResult(bounding_box=bbox, keypoints=keypoints, score=score)) - - batch_results.append(current_image_detections) - - # Check MediaPipe documentation if resource cleanup is required here. - return batch_results + ) -> List[List[FaceDetectionResult]]: + """Detects faces in the input image tensor.""" + return super().detect( + image, + running_mode=running_mode, + delegate=delegate, + min_detection_confidence=min_detection_confidence, + )