diff --git a/src/mediapipe_vision/face_landmark/detector.py b/src/mediapipe_vision/face_landmark/detector.py index c828bcb..32bad27 100644 --- a/src/mediapipe_vision/face_landmark/detector.py +++ b/src/mediapipe_vision/face_landmark/detector.py @@ -1,4 +1,4 @@ -from typing import Dict, List, Optional, Tuple +from typing import Any, Dict, List, Optional, Tuple import mediapipe as mp import numpy as np @@ -8,12 +8,11 @@ from mediapipe.tasks.python import vision from mediapipe.tasks.python.core.base_options import BaseOptions from ...utils.timing import TimestampProvider # Import TimestampProvider - -# Import new types +from ..common import BaseDetector from ..types import LandmarkPoint -class FaceLandmarkDetector: +class FaceLandmarkDetector(BaseDetector[List[LandmarkPoint]]): """Detects face landmarks and blendshapes using MediaPipe FaceLandmarker.""" def __init__(self, model_path: str): @@ -22,11 +21,11 @@ class FaceLandmarkDetector: Args: model_path: Path to the MediaPipe FaceLandmarker .task file. """ - if not model_path: - raise ValueError("A valid model_path must be provided.") - - self.model_path = model_path - self._detector_instance = None + super().__init__(model_path) + self._output_blendshapes = False + self._output_transform_matrix = False + self._blendshapes = [] + self._transform_matrices = [] self._current_options = None # Store current options for comparison self._timestamp_provider = None # Added @@ -61,6 +60,81 @@ class FaceLandmarkDetector: self._timestamp_provider = TimestampProvider() if mode_enum == vision.RunningMode.VIDEO else None return vision.FaceLandmarker.create_from_options(options) + def _create_detector_options(self, base_options: python.BaseOptions, + mode_enum: vision.RunningMode, **kwargs) -> vision.FaceLandmarkerOptions: + """Create FaceLandmarker-specific options with parameters: + - num_faces: Maximum number of faces to detect + - min_detection_confidence: Minimum confidence for face detection + - min_presence_confidence: Minimum confidence for face presence + - min_tracking_confidence: Minimum confidence for face tracking + - output_blendshapes: Whether to output face blendshapes + - output_transform_matrix: Whether to output facial transformation matrices + """ + self._output_blendshapes = kwargs.get('output_blendshapes', False) + self._output_transform_matrix = kwargs.get('output_transform_matrix', False) + + return vision.FaceLandmarkerOptions( + base_options=base_options, + running_mode=mode_enum, + num_faces=kwargs.get('num_faces', 1), + min_face_detection_confidence=kwargs.get('min_detection_confidence', 0.5), + min_face_presence_confidence=kwargs.get('min_presence_confidence', 0.5), + min_tracking_confidence=kwargs.get('min_tracking_confidence', 0.5), + output_face_blendshapes=self._output_blendshapes, + output_facial_transformation_matrixes=self._output_transform_matrix, + ) + + def _create_detector_instance(self, options: vision.FaceLandmarkerOptions) -> vision.FaceLandmarker: + return vision.FaceLandmarker.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.num_faces, + self._current_options.min_face_detection_confidence, + self._current_options.min_face_presence_confidence, + self._current_options.min_tracking_confidence, + self._current_options.output_face_blendshapes, + self._current_options.output_facial_transformation_matrixes, + running_mode, + delegate, + ) + else: + return ( + kwargs.get('num_faces', 1), + kwargs.get('min_detection_confidence', 0.5), + kwargs.get('min_presence_confidence', 0.5), + kwargs.get('min_tracking_confidence', 0.5), + kwargs.get('output_blendshapes', False), + kwargs.get('output_transform_matrix', False), + running_mode, + delegate, + ) + + def _process_detection_result(self, detection_result: Any) -> List[List[LandmarkPoint]]: + """Process detection result and extract blendshapes and transforms if requested.""" + self._blendshapes = [] + self._transform_matrices = [] + + current_image_landmarks = [] + if detection_result and detection_result.face_landmarks: + for face_idx, face_landmarks_mp in enumerate(detection_result.face_landmarks): + landmarks = [LandmarkPoint(index=lm_idx, x=lm.x, y=lm.y, z=lm.z) for lm_idx, lm in enumerate(face_landmarks_mp)] + current_image_landmarks.append(landmarks) + + if self._output_blendshapes and detection_result.face_blendshapes and face_idx < len(detection_result.face_blendshapes): + blendshapes = {bs.category_name: bs.score for bs in detection_result.face_blendshapes[face_idx]} + self._blendshapes.append(blendshapes) + + if ( + self._output_transform_matrix + and detection_result.facial_transformation_matrixes + and face_idx < len(detection_result.facial_transformation_matrixes) + ): + self._transform_matrices.append(detection_result.facial_transformation_matrixes[face_idx]) + + return current_image_landmarks + def detect( self, image: torch.Tensor, @@ -72,7 +146,7 @@ class FaceLandmarkDetector: output_transform_matrix: bool = False, running_mode: str = "video", delegate: str = "cpu", - ) -> Tuple[List[List[LandmarkPoint]], Optional[List[Dict[str, float]]], Optional[List[np.ndarray]]]: # Added running_mode + ) -> Tuple[List[List[List[LandmarkPoint]]], Optional[List[List[Dict[str, float]]]], Optional[List[List[np.ndarray]]]]: """Detects face landmarks in the input image tensor. Returns: