diff --git a/src/mediapipe_vision/common/__init__.py b/src/mediapipe_vision/common/__init__.py new file mode 100644 index 0000000..33004e0 --- /dev/null +++ b/src/mediapipe_vision/common/__init__.py @@ -0,0 +1,3 @@ +from .base_detector import BaseDetector + +__all__ = ["BaseDetector"] \ No newline at end of file diff --git a/src/mediapipe_vision/common/base_detector.py b/src/mediapipe_vision/common/base_detector.py new file mode 100644 index 0000000..b42a2fd --- /dev/null +++ b/src/mediapipe_vision/common/base_detector.py @@ -0,0 +1,175 @@ +from abc import ABC, abstractmethod +from typing import Any, Dict, Generic, List, Optional, Tuple, TypeVar, Union + +import mediapipe as mp +import numpy as np +import torch +from mediapipe.tasks import python +from mediapipe.tasks.python import vision +from mediapipe.tasks.python.core.base_options import BaseOptions + +from ...utils.timing import TimestampProvider + +T = TypeVar('T') # Generic type for detector results + +class BaseDetector(Generic[T], ABC): + """Base class for MediaPipe vision detectors. + + This class provides common functionality for all MediaPipe vision detectors, + reducing code duplication and providing a consistent interface. + + Generic type T represents the detector-specific result type. + """ + + def __init__(self, model_path: str): + """Initialize the detector with the model path. + + Args: + model_path: Path to the MediaPipe model file. + """ + if not model_path: + raise ValueError("A valid model_path must be provided.") + + self.model_path = model_path + self._detector_instance = None + self._current_options = None + self._timestamp_provider = None + + @abstractmethod + def _create_detector_options(self, base_options: python.BaseOptions, + mode_enum: vision.RunningMode, **kwargs) -> Any: + """Create detector-specific options. + + Args: + base_options: MediaPipe base options with model path and delegate. + mode_enum: RunningMode enum (IMAGE or VIDEO). + **kwargs: Detector-specific configuration parameters. + + Returns: + Detector-specific options object. + """ + pass + + @abstractmethod + def _get_options_tuple(self, running_mode: str = None, + delegate: str = None, **kwargs) -> tuple: + """Get tuple of options for comparison to detect config changes. + + Args: + running_mode: String representation of running mode ("image" or "video"). + delegate: String representation of delegate ("cpu" or "gpu"). + **kwargs: Detector-specific parameters to include in comparison. + + Returns: + Tuple containing all configuration options for comparison. + """ + pass + + @abstractmethod + def _process_detection_result(self, detection_result: Any) -> List[T]: + """Process detector-specific results. + + Args: + detection_result: Raw detection result from MediaPipe. + + Returns: + List of processed detection results of type T. + """ + pass + + @abstractmethod + def _create_detector_instance(self, options: Any) -> Any: + """Create detector-specific instance. + + Args: + options: Configured detector options. + + Returns: + MediaPipe detector instance. + """ + pass + + def _initialize_detector(self, running_mode: str, delegate: str, **kwargs) -> Any: + """Initialize the MediaPipe detector with specified options. + + Args: + running_mode: "image" or "video". + delegate: "cpu" or "gpu". + **kwargs: Detector-specific configuration parameters. + + Returns: + Initialized detector instance. + """ + 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 = self._create_detector_options(base_options, mode_enum, **kwargs) + + self._current_options = options + self._timestamp_provider = TimestampProvider() if mode_enum == vision.RunningMode.VIDEO else None + + return self._create_detector_instance(options) + + def detect(self, image: torch.Tensor, running_mode: str = "video", + delegate: str = "cpu", **kwargs) -> List[List[T]]: + """Detect features in the input image tensor. + + Args: + image: Input image tensor in BHWC format. + running_mode: "image" or "video". + delegate: "cpu" or "gpu". + **kwargs: Detector-specific parameters. + + Returns: + A list (batch) of lists (detections per image) of detection results. + + Raises: + ValueError: If the input tensor is not in the expected format. + """ + if image.dim() != 4: + raise ValueError("Input tensor must be in BHWC format.") + + # Check if options changed to decide whether to reinitialize + new_options_tuple = self._get_options_tuple(running_mode=running_mode, delegate=delegate, **kwargs) + current_options_tuple = None + if self._current_options: + current_mode = "image" if self._current_options.running_mode == vision.RunningMode.IMAGE else "video" + current_delegate = "cpu" if self._current_options.base_options.delegate == BaseOptions.Delegate.CPU else "gpu" + current_options_tuple = self._get_options_tuple(running_mode=current_mode, delegate=current_delegate) + + if self._detector_instance is None or new_options_tuple != current_options_tuple: + if self._detector_instance and hasattr(self._detector_instance, "close"): + self._detector_instance.close() + self._detector_instance = self._initialize_detector(running_mode, delegate, **kwargs) + + batch_size = image.shape[0] + batch_results = [] + is_video_mode = self._current_options.running_mode == vision.RunningMode.VIDEO + + # Get the appropriate detection function based on mode + if hasattr(self._detector_instance, "detect_for_video") and is_video_mode: + detect_func = self._detector_instance.detect_for_video + elif hasattr(self._detector_instance, "detect"): + detect_func = self._detector_instance.detect + else: + # Try to find the appropriate method name for this detector type + func_names = [func for func in dir(self._detector_instance) if func.startswith(("detect", "recognize")) and callable(getattr(self._detector_instance, func))] + video_funcs = [func for func in func_names if "video" in func.lower()] + detect_func = getattr(self._detector_instance, video_funcs[0] if is_video_mode and video_funcs else func_names[0]) + + # Process each image in the batch + for i in range(batch_size): + img_tensor = image[i] + 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: + timestamp_ms = self._timestamp_provider.next() + detection_result = detect_func(mp_image, timestamp_ms) + else: + detection_result = detect_func(mp_image) + + batch_results.append(self._process_detection_result(detection_result)) + + return batch_results \ No newline at end of file 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, + ) 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: diff --git a/src/mediapipe_vision/face_stylization/detector.py b/src/mediapipe_vision/face_stylization/detector.py index 8681dba..3661278 100644 --- a/src/mediapipe_vision/face_stylization/detector.py +++ b/src/mediapipe_vision/face_stylization/detector.py @@ -1,45 +1,60 @@ +from typing import Any, List + import mediapipe as mp import numpy as np import torch from mediapipe.tasks import python from mediapipe.tasks.python import vision +from mediapipe.tasks.python.core.base_options import BaseOptions from ...utils.timing import TimestampProvider +from ..common import BaseDetector -class FaceStylizer: +class FaceStylizer(BaseDetector[torch.Tensor]): """Applies stylization to faces using MediaPipe FaceStylizer.""" def __init__(self, model_path: str): if not model_path: raise ValueError("A valid model_path must be provided.") self.model_path = model_path - self.detector_instance = None + self._detector_instance = None self.current_options = None self._timestamp_provider = None - def _initialize_detector(self): - """Initializes the FaceStylizer.""" - base_options = python.BaseOptions(model_asset_path=self.model_path) + def _create_detector_options(self, base_options: python.BaseOptions, + mode_enum: vision.RunningMode, **kwargs) -> vision.FaceStylizerOptions: + """Create FaceStylizer-specific options. + Note: FaceStylizer only supports IMAGE mode, not VIDEO mode. + """ # FaceStylizerOptions has no specific config other than base options - # Check if FaceStylizer supports VIDEO mode - # According to docs (as of late 2023/early 2024), it seems it only supports IMAGE mode. - # Sticking with IMAGE mode. - options = vision.FaceStylizerOptions( - base_options=base_options - # running_mode=vision.RunningMode.VIDEO # Not supported - ) - self.current_options = () - self._timestamp_provider = TimestampProvider() + # and doesn't support VIDEO mode - we ignore the mode_enum parameter + return vision.FaceStylizerOptions(base_options=base_options) + + def _create_detector_instance(self, options: vision.FaceStylizerOptions) -> vision.FaceStylizer: return vision.FaceStylizer.create_from_options(options) - def stylize(self, image: torch.Tensor) -> torch.Tensor: # Output is IMAGE + def _get_options_tuple(self, running_mode: str = None, delegate: str = None, **kwargs) -> tuple: + # FaceStylizer has no additional options to track + return (delegate,) + + def _process_detection_result(self, stylization_result: Any) -> torch.Tensor: + """Process stylization result to torch tensor.""" + if stylization_result: + # Convert MediaPipe Image back to NumPy (HWC, uint8) + stylized_np = stylization_result.numpy_view() + # Convert to float tensor (HWC, 0-1) + return torch.from_numpy(stylized_np).float().div(255.0) + return None # If no result, we'll handle in the stylize method + + def stylize(self, image: torch.Tensor) -> torch.Tensor: """Applies face stylization to the input image tensor(s).""" if image.dim() != 4: raise ValueError("Input tensor must be in BHWC format.") - if self.detector_instance is None: - self.detector_instance = self._initialize_detector() + # Ensure detector is initialized - always use IMAGE mode + if self._detector_instance is None: + self._initialize_detector("image", "cpu") batch_size = image.shape[0] batch_results = [] @@ -49,20 +64,27 @@ class FaceStylizer: np_image = (img_tensor.cpu().numpy() * 255).astype(np.uint8) mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=np_image) - stylization_result = self.detector_instance.stylize(mp_image) - - output_image_tensor = None - if stylization_result: - # Convert MediaPipe Image back to NumPy (HWC, uint8) - stylized_np = stylization_result.numpy_view() - # Convert to float tensor (HWC, 0-1) - output_image_tensor = torch.from_numpy(stylized_np).float().div(255.0) - else: - # Handle cases where stylization fails (e.g., no face detected) - # Return original image tensor for this batch item + # FaceStylizer only has the stylize method, not detect + stylization_result = self._detector_instance.stylize(mp_image) + output_image_tensor = self._process_detection_result(stylization_result) + + # If processing failed (no face detected), return original image + if output_image_tensor is None: output_image_tensor = img_tensor - + batch_results.append(output_image_tensor) output_batch_tensor = torch.stack(batch_results, dim=0) return output_batch_tensor + + def _initialize_detector(self, running_mode: str, delegate: str): + """Initializes the FaceStylizer.""" + base_options = python.BaseOptions(model_asset_path=self.model_path) + # FaceStylizerOptions has no specific config other than base options + # Check if FaceStylizer supports VIDEO mode + # According to docs (as of late 2023/early 2024), it seems it only supports IMAGE mode. + # Sticking with IMAGE mode. + options = self._create_detector_options(base_options, vision.RunningMode.IMAGE) + self.current_options = () + self._timestamp_provider = TimestampProvider() + self._detector_instance = self._create_detector_instance(options) diff --git a/src/mediapipe_vision/gesture_recognition/detector.py b/src/mediapipe_vision/gesture_recognition/detector.py index 64c6338..f008718 100644 --- a/src/mediapipe_vision/gesture_recognition/detector.py +++ b/src/mediapipe_vision/gesture_recognition/detector.py @@ -1,4 +1,4 @@ -from typing import List +from typing import Any, List 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 GestureCategory, GestureRecognitionResult -class GestureRecognizer: +class GestureRecognizer(BaseDetector[GestureRecognitionResult]): """Recognizes hand gestures using MediaPipe GestureRecognizer.""" def __init__(self, model_path: str): @@ -24,33 +23,96 @@ class GestureRecognizer: self._current_options = None self._timestamp_provider = None # Added - def _initialize_detector( - self, - num_hands: int, - min_detection_confidence: float, - min_tracking_confidence: float, - min_presence_confidence: float, - running_mode: str, - delegate: str, - ): # Added running_mode - mode_enum = vision.RunningMode.IMAGE if running_mode == "image" else vision.RunningMode.VIDEO # Added - 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) - # Note: GestureRecognizerOptions has min_tracking_confidence, useful for VIDEO mode - options = vision.GestureRecognizerOptions( + def _create_detector_options(self, base_options: python.BaseOptions, + mode_enum: vision.RunningMode, **kwargs) -> vision.GestureRecognizerOptions: + """Create GestureRecognizer-specific options with parameters: + - num_hands: Maximum number of hands to detect + - min_detection_confidence: Minimum confidence for hand detection + - min_tracking_confidence: Minimum confidence for hand tracking + - min_presence_confidence: Minimum confidence for hand presence + """ + return vision.GestureRecognizerOptions( base_options=base_options, - running_mode=mode_enum, # Use enum mode - num_hands=num_hands, - min_hand_detection_confidence=min_detection_confidence, - min_tracking_confidence=min_tracking_confidence, - min_hand_presence_confidence=min_presence_confidence, - # Add other relevant VIDEO mode parameters if needed (e.g., canned_gestures_classifier_options) + running_mode=mode_enum, + num_hands=kwargs.get('num_hands', 2), + min_hand_detection_confidence=kwargs.get('min_detection_confidence', 0.5), + min_tracking_confidence=kwargs.get('min_tracking_confidence', 0.5), + min_hand_presence_confidence=kwargs.get('min_presence_confidence', 0.5), ) - self._current_options = options - # 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.GestureRecognizerOptions) -> vision.GestureRecognizer: + """Create GestureRecognizer instance. + + Args: + options: Configured GestureRecognizerOptions. + + Returns: + MediaPipe GestureRecognizer instance. + """ return vision.GestureRecognizer.create_from_options(options) + def _get_options_tuple(self, running_mode: str = None, delegate: str = None, **kwargs) -> tuple: + """Get tuple of options for comparison. + + Args: + running_mode: String representation of running mode ("image" or "video"). + delegate: String representation of delegate ("cpu" or "gpu"). + **kwargs: GestureRecognizer-specific parameters. + + Returns: + Tuple containing all configuration options for comparison. + """ + if self._current_options: + # Getting values from existing options object + return ( + self._current_options.num_hands, + self._current_options.min_hand_detection_confidence, + self._current_options.min_tracking_confidence, + self._current_options.min_hand_presence_confidence, + running_mode, + delegate, + ) + else: + # Getting values from kwargs for new options + return ( + kwargs.get('num_hands', 2), + kwargs.get('min_detection_confidence', 0.5), + kwargs.get('min_tracking_confidence', 0.5), + kwargs.get('min_presence_confidence', 0.5), + running_mode, + delegate, + ) + + def _process_detection_result(self, recognition_result: Any) -> List[GestureRecognitionResult]: + """Process GestureRecognizer detection results. + + Args: + recognition_result: Raw GestureRecognizer result from MediaPipe. + + Returns: + List of GestureRecognitionResult objects. + """ + current_image_gestures = [] + if recognition_result and recognition_result.gestures: + for hand_idx, hand_gestures in enumerate(recognition_result.gestures): + gesture_categories = [ + GestureCategory( + index=cat.index, + score=cat.score, + display_name=cat.display_name, + category_name=cat.category_name, + ) + for cat in hand_gestures + ] + handedness = None + if recognition_result.handedness and hand_idx < len(recognition_result.handedness): + handedness_cats = recognition_result.handedness[hand_idx] + if handedness_cats: + handedness = handedness_cats[0].display_name + current_image_gestures.append(GestureRecognitionResult(gestures=gesture_categories, handedness=handedness)) + return current_image_gestures + + # Public API method that keeps the same signature as before def recognize( self, image: torch.Tensor, @@ -60,80 +122,14 @@ class GestureRecognizer: min_presence_confidence: float = 0.5, running_mode: str = "video", delegate: str = "cpu", - ) -> List[List[GestureRecognitionResult]]: # Added running_mode + ) -> List[List[GestureRecognitionResult]]: """Recognizes hand gestures in the input image tensor.""" - if image.dim() != 4: - raise ValueError("Input tensor must be in BHWC format.") - - # Add running_mode to tuple check - new_options_tuple = ( - num_hands, - min_detection_confidence, - min_tracking_confidence, - min_presence_confidence, - running_mode, - delegate.lower(), + return self.detect( + image, + running_mode=running_mode, + delegate=delegate, + num_hands=num_hands, + min_detection_confidence=min_detection_confidence, + min_tracking_confidence=min_tracking_confidence, + min_presence_confidence=min_presence_confidence, ) - current_options_tuple = None - if self._current_options: - current_options_tuple = ( - self._current_options.num_hands, - self._current_options.min_hand_detection_confidence, - self._current_options.min_tracking_confidence, - self._current_options.min_hand_presence_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: - if self._detector_instance and hasattr(self._detector_instance, "close"): - self._detector_instance.close() - # Pass new params to initializer - self._detector_instance = self._initialize_detector( - num_hands, - min_detection_confidence, - min_tracking_confidence, - min_presence_confidence, - running_mode, - delegate, - ) # Pass running_mode - - batch_size = image.shape[0] - batch_results = [] - # Determine the correct recognition function based on mode - is_video_mode = self._current_options.running_mode == vision.RunningMode.VIDEO - recognize_func = self._detector_instance.recognize_for_video if is_video_mode else self._detector_instance.recognize - - for i in range(batch_size): - img_tensor = image[i] - 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: - timestamp_ms = self._timestamp_provider.next() - recognition_result = recognize_func(mp_image, timestamp_ms) - else: - recognition_result = recognize_func(mp_image) - - current_image_gestures = [] - if recognition_result and recognition_result.gestures: # Added check for recognition_result - for hand_idx, hand_gestures in enumerate(recognition_result.gestures): - gesture_categories = [ - GestureCategory( - index=cat.index, - score=cat.score, - display_name=cat.display_name, - category_name=cat.category_name, - ) - for cat in hand_gestures - ] - handedness = None - if recognition_result.handedness and hand_idx < len(recognition_result.handedness): - handedness_cats = recognition_result.handedness[hand_idx] - if handedness_cats: - handedness = handedness_cats[0].display_name - current_image_gestures.append(GestureRecognitionResult(gestures=gesture_categories, handedness=handedness)) - batch_results.append(current_image_gestures) - - # self._detector_instance.close() # Close handled on re-init - return batch_results diff --git a/src/mediapipe_vision/hand_landmark/detector.py b/src/mediapipe_vision/hand_landmark/detector.py index 420d996..3ddfaea 100644 --- a/src/mediapipe_vision/hand_landmark/detector.py +++ b/src/mediapipe_vision/hand_landmark/detector.py @@ -1,4 +1,4 @@ -from typing import List +from typing import Any, List 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 HandLandmarksResult, LandmarkPoint -class HandLandmarkDetector: +class HandLandmarkDetector(BaseDetector[HandLandmarksResult]): """Detects hand landmarks using MediaPipe HandLandmarker.""" def __init__(self, model_path: str): @@ -31,33 +30,70 @@ class HandLandmarkDetector: self._current_options = None self._timestamp_provider = None # Added - def _initialize_detector( - self, - num_hands: int, - min_detection_confidence: float, - min_presence_confidence: float, - min_tracking_confidence: float, - running_mode: str, - delegate: str, - ): # Added running_mode - """Initializes the HandLandmarker detector.""" - mode_enum = vision.RunningMode.IMAGE if running_mode == "image" else vision.RunningMode.VIDEO # Added - 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.HandLandmarkerOptions( + def _create_detector_options(self, base_options: python.BaseOptions, + mode_enum: vision.RunningMode, **kwargs) -> vision.HandLandmarkerOptions: + """Create HandLandmarker-specific options with parameters: + - num_hands: Maximum number of hands to detect + - min_detection_confidence: Minimum confidence for hand detection + - min_presence_confidence: Minimum confidence for hand presence + - min_tracking_confidence: Minimum confidence for hand tracking + """ + return vision.HandLandmarkerOptions( base_options=base_options, - running_mode=mode_enum, # Use enum mode - num_hands=num_hands, - min_hand_detection_confidence=min_detection_confidence, - # min_hand_presence_confidence and min_tracking_confidence are for VIDEO/LIVE_STREAM - min_hand_presence_confidence=min_presence_confidence, - min_tracking_confidence=min_tracking_confidence, + running_mode=mode_enum, + num_hands=kwargs.get('num_hands', 2), + min_hand_detection_confidence=kwargs.get('min_detection_confidence', 0.5), + min_hand_presence_confidence=kwargs.get('min_presence_confidence', 0.5), + min_tracking_confidence=kwargs.get('min_tracking_confidence', 0.5), ) - self._current_options = options - # 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.HandLandmarkerOptions) -> vision.HandLandmarker: return vision.HandLandmarker.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_hands, + self._current_options.min_hand_detection_confidence, + self._current_options.min_hand_presence_confidence, + self._current_options.min_tracking_confidence, + running_mode, + delegate, + ) + else: + return ( + kwargs.get('num_hands', 2), + kwargs.get('min_detection_confidence', 0.5), + kwargs.get('min_presence_confidence', 0.5), + kwargs.get('min_tracking_confidence', 0.5), + running_mode, + delegate, + ) + + def _process_detection_result(self, detection_result: Any) -> List[HandLandmarksResult]: + current_image_results = [] + if detection_result and detection_result.hand_landmarks: + for hand_idx, hand_landmarks_mp in enumerate(detection_result.hand_landmarks): + landmarks = [LandmarkPoint(index=lm_idx, x=lm.x, y=lm.y, z=lm.z) for lm_idx, lm in enumerate(hand_landmarks_mp)] + + world_landmarks = None + if detection_result.hand_world_landmarks and hand_idx < len(detection_result.hand_world_landmarks): + world_landmarks = [ + LandmarkPoint(index=lm_idx, x=lm.x, y=lm.y, z=lm.z) + for lm_idx, lm in enumerate(detection_result.hand_world_landmarks[hand_idx]) + ] + + handedness = None + if detection_result.handedness and hand_idx < len(detection_result.handedness): + handedness_cats = detection_result.handedness[hand_idx] + if handedness_cats: + handedness = handedness_cats[0].display_name + + current_image_results.append( + HandLandmarksResult(landmarks=landmarks, world_landmarks=world_landmarks, handedness=handedness) + ) + return current_image_results + def detect( self, image: torch.Tensor, @@ -67,88 +103,14 @@ class HandLandmarkDetector: min_tracking_confidence: float = 0.5, running_mode: str = "video", delegate: str = "cpu", - ) -> List[List[HandLandmarksResult]]: # Added running_mode - """Detects hand landmarks in the input image tensor. - - Returns: - A list (batch) of lists (hands per image) of HandLandmarksResult objects. - """ - if image.dim() != 4: - raise ValueError("Input tensor must be in BHWC format.") - - # Add running_mode to tuple check - new_options_tuple = ( - num_hands, - min_detection_confidence, - min_presence_confidence, - min_tracking_confidence, - running_mode, - delegate.lower(), + ) -> List[List[HandLandmarksResult]]: + """Detects hand landmarks in the input image tensor.""" + return super().detect( + image, + running_mode=running_mode, + delegate=delegate, + num_hands=num_hands, + min_detection_confidence=min_detection_confidence, + min_presence_confidence=min_presence_confidence, + min_tracking_confidence=min_tracking_confidence, ) - current_options_tuple = None - if self._current_options: - current_options_tuple = ( - self._current_options.num_hands, - self._current_options.min_hand_detection_confidence, - self._current_options.min_hand_presence_confidence, - self._current_options.min_tracking_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: - if self._detector_instance and hasattr(self._detector_instance, "close"): - self._detector_instance.close() - # Pass new params to initializer - self._detector_instance = self._initialize_detector( - num_hands, - min_detection_confidence, - min_presence_confidence, - min_tracking_confidence, - running_mode, - delegate, - ) # Pass running_mode - - 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: - timestamp_ms = self._timestamp_provider.next() - detection_result = detect_func(mp_image, timestamp_ms) - else: - detection_result = detect_func(mp_image) - - current_image_results = [] - if detection_result and detection_result.hand_landmarks: # Added check for detection_result - for hand_idx, hand_landmarks_mp in enumerate(detection_result.hand_landmarks): - landmarks = [LandmarkPoint(index=lm_idx, x=lm.x, y=lm.y, z=lm.z) for lm_idx, lm in enumerate(hand_landmarks_mp)] - - world_landmarks = None - if detection_result.hand_world_landmarks and hand_idx < len(detection_result.hand_world_landmarks): - world_landmarks = [ - LandmarkPoint(index=lm_idx, x=lm.x, y=lm.y, z=lm.z) - for lm_idx, lm in enumerate(detection_result.hand_world_landmarks[hand_idx]) - ] - - handedness = None - if detection_result.handedness and hand_idx < len(detection_result.handedness): - handedness_cats = detection_result.handedness[hand_idx] - if handedness_cats: - handedness = handedness_cats[0].display_name - - current_image_results.append( - HandLandmarksResult(landmarks=landmarks, world_landmarks=world_landmarks, handedness=handedness) - ) - - batch_results.append(current_image_results) - - # self._detector_instance.close() # Close handled on re-init - return batch_results diff --git a/src/mediapipe_vision/image_embedding/detector.py b/src/mediapipe_vision/image_embedding/detector.py index bf87f51..107df12 100644 --- a/src/mediapipe_vision/image_embedding/detector.py +++ b/src/mediapipe_vision/image_embedding/detector.py @@ -10,35 +10,145 @@ from mediapipe.tasks.python.core.base_options import BaseOptions # Import TimestampProvider from the new location from ...utils.timing import TimestampProvider +# Import BaseDetector class +from ..common.base_detector import BaseDetector + # Import new types from ..types import ImageEmbedderResult -class ImageEmbedder: +class ImageEmbedder(BaseDetector): """Generates image embeddings using MediaPipe ImageEmbedder.""" - def __init__(self, model_path: str): - if not model_path: - raise ValueError("A valid model_path must be provided.") - self.model_path = model_path - self._detector_instance = None - self._current_options = None - self._timestamp_provider = None # Added - - def _initialize_detector(self, l2_normalize: bool, quantize: bool, running_mode: str, delegate: str): # Added running_mode - mode_enum = vision.RunningMode.IMAGE if running_mode == "image" else vision.RunningMode.VIDEO # Added - 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.ImageEmbedderOptions( + def _create_detector_options(self, base_options: python.BaseOptions, + mode_enum: vision.RunningMode, **kwargs) -> vision.ImageEmbedderOptions: + """Create image embedder-specific options. + + Args: + base_options: MediaPipe base options with model path and delegate. + mode_enum: RunningMode enum (IMAGE or VIDEO). + **kwargs: Detector-specific configuration parameters. + + Returns: + Configured ImageEmbedderOptions object. + """ + l2_normalize = kwargs.get('l2_normalize', True) + quantize = kwargs.get('quantize', False) + + return vision.ImageEmbedderOptions( base_options=base_options, - running_mode=mode_enum, # Use enum mode + running_mode=mode_enum, l2_normalize=l2_normalize, quantize=quantize, ) - self._current_options = options - # Only create timestamp provider for VIDEO mode - self._timestamp_provider = TimestampProvider() if mode_enum == vision.RunningMode.VIDEO else None + + def _get_options_tuple(self, running_mode: str = None, + delegate: str = None, **kwargs) -> tuple: + """Get tuple of options for comparison to detect config changes. + + Args: + running_mode: String representation of running mode ("image" or "video"). + delegate: String representation of delegate ("cpu" or "gpu"). + **kwargs: Detector-specific parameters to include in comparison. + + Returns: + Tuple containing all configuration options for comparison. + """ + l2_normalize = kwargs.get('l2_normalize', True) + quantize = kwargs.get('quantize', False) + + return (l2_normalize, quantize, running_mode, delegate.lower() if delegate else "cpu") + + def _process_detection_result(self, detection_result) -> List[ImageEmbedderResult]: + """Process embedding results. + + Args: + detection_result: Raw embedding result from MediaPipe. + + Returns: + List of ImageEmbedderResult objects. + """ + current_image_embeddings = [] + + # Check for quantize parameter in current options + quantize = False + if self._current_options and hasattr(self._current_options, 'quantize'): + quantize = self._current_options.quantize + + if detection_result and detection_result.embeddings: + for emb in detection_result.embeddings: + result = ImageEmbedderResult( + float_embedding=list(emb.embedding) if not quantize else None, + quantized_embedding=bytes(emb.embedding) if quantize else None, + head_index=emb.head_index, + head_name=emb.head_name, + ) + current_image_embeddings.append(result) + + return current_image_embeddings + + def _create_detector_instance(self, options: vision.ImageEmbedderOptions) -> vision.ImageEmbedder: + """Create embedder instance. + + Args: + options: Configured ImageEmbedderOptions. + + Returns: + MediaPipe ImageEmbedder instance. + """ return vision.ImageEmbedder.create_from_options(options) + + def detect(self, image: torch.Tensor, running_mode: str = "video", + delegate: str = "cpu", **kwargs) -> List[List[ImageEmbedderResult]]: + """Override of the BaseDetector detect method to use embed/embed_for_video. + + Args: + image: Input image tensor in BHWC format. + running_mode: "image" or "video". + delegate: "cpu" or "gpu". + **kwargs: Detector-specific parameters. + + Returns: + A list (batch) of lists (detections per image) of detection results. + """ + if image.dim() != 4: + raise ValueError("Input tensor must be in BHWC format.") + + # Check if options changed to decide whether to reinitialize + new_options_tuple = self._get_options_tuple(running_mode=running_mode, delegate=delegate, **kwargs) + current_options_tuple = None + if self._current_options: + current_mode = "image" if self._current_options.running_mode == vision.RunningMode.IMAGE else "video" + current_delegate = "cpu" if self._current_options.base_options.delegate == BaseOptions.Delegate.CPU else "gpu" + current_options_tuple = self._get_options_tuple(running_mode=current_mode, delegate=current_delegate) + + if self._detector_instance is None or new_options_tuple != current_options_tuple: + if self._detector_instance and hasattr(self._detector_instance, "close"): + self._detector_instance.close() + self._detector_instance = self._initialize_detector(running_mode, delegate, **kwargs) + + batch_size = image.shape[0] + batch_results = [] + is_video_mode = self._current_options.running_mode == vision.RunningMode.VIDEO + + # Use embed/embed_for_video methods directly + embed_func = self._detector_instance.embed_for_video if is_video_mode else self._detector_instance.embed + + # Process each image in the batch + for i in range(batch_size): + img_tensor = image[i] + 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: + timestamp_ms = self._timestamp_provider.next() + detection_result = embed_func(mp_image, timestamp_ms) + else: + detection_result = embed_func(mp_image) + + batch_results.append(self._process_detection_result(detection_result)) + + return batch_results def embed( self, @@ -47,55 +157,24 @@ class ImageEmbedder: quantize: bool = False, running_mode: str = "video", delegate: str = "cpu", - ) -> List[List[ImageEmbedderResult]]: # Added running_mode - """Generates embeddings for the input image tensor(s).""" - if image.dim() != 4: - raise ValueError("Input tensor must be in BHWC format.") - - # Add running_mode to tuple check - new_options_tuple = (l2_normalize, quantize, running_mode, delegate.lower()) - current_options_tuple = None - if self._current_options: - current_options_tuple = ( - self._current_options.l2_normalize, - self._current_options.quantize, - "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: - # ImageEmbedder doesn't seem to have close(), so just replace instance - if self._detector_instance and hasattr(self._detector_instance, "close"): - self._detector_instance.close() - self._detector_instance = self._initialize_detector(l2_normalize, quantize, running_mode, delegate) - - batch_size = image.shape[0] - batch_results = [] - # Determine the correct embedding function based on mode - is_video_mode = self._current_options.running_mode == vision.RunningMode.VIDEO - embed_func = self._detector_instance.embed_for_video if is_video_mode else self._detector_instance.embed - - for i in range(batch_size): - img_tensor = image[i] - 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: - timestamp_ms = self._timestamp_provider.next() - embedding_result = embed_func(mp_image, timestamp_ms) - else: - embedding_result = embed_func(mp_image) - - current_image_embeddings = [] - if embedding_result and embedding_result.embeddings: # Added check for embedding_result - for emb in embedding_result.embeddings: - result = ImageEmbedderResult( - float_embedding=list(emb.embedding) if not quantize else None, - quantized_embedding=bytes(emb.embedding) if quantize else None, - head_index=emb.head_index, - head_name=emb.head_name, - ) - current_image_embeddings.append(result) - batch_results.append(current_image_embeddings) - - return batch_results + ) -> List[List[ImageEmbedderResult]]: + """Generates embeddings for the input image tensor(s). + + Args: + image: Input image tensor (BHWC, RGB, 0-1 float). + l2_normalize: Whether to L2-normalize the embedding vectors. + quantize: Whether to quantize the embedding vectors. + running_mode: Processing mode ("image" or "video"). + delegate: Computation delegate ('cpu' or 'gpu'). + + Returns: + A list (batch) of lists (embeddings) of ImageEmbedderResult objects. + """ + # Use the generic detect method from BaseDetector + return self.detect( + image, + running_mode=running_mode, + delegate=delegate, + l2_normalize=l2_normalize, + quantize=quantize + ) diff --git a/src/mediapipe_vision/image_segmentation/segmenter.py b/src/mediapipe_vision/image_segmentation/segmenter.py index cc3a833..3eef534 100644 --- a/src/mediapipe_vision/image_segmentation/segmenter.py +++ b/src/mediapipe_vision/image_segmentation/segmenter.py @@ -5,93 +5,168 @@ This module contains the implementation of MediaPipe Image Segmentation function import logging import platform # Needed for OS check +from typing import List, Tuple import mediapipe as mp import numpy as np import torch from mediapipe.tasks import python from mediapipe.tasks.python import vision -from mediapipe.tasks.python.vision import ImageSegmenter +from mediapipe.tasks.python.vision import ImageSegmenter as MPImageSegmenter +from ..common.base_detector import BaseDetector from ...utils.timing import TimestampProvider logger = logging.getLogger(__name__) -class ImageSegmenter: +class ImageSegmenter(BaseDetector): """Performs image segmentation using MediaPipe ImageSegmenter. Adapted from ComfyUI-Stream-Pack implementation. """ - def __init__(self, model_path: str): - """Initialize the segmenter with the model path. - + def _create_detector_options(self, base_options: python.BaseOptions, + mode_enum: vision.RunningMode, **kwargs) -> vision.ImageSegmenterOptions: + """Create image segmenter-specific options. + Args: - model_path: Path to the MediaPipe ImageSegmenter .tflite or .task file. + base_options: MediaPipe base options with model path and delegate. + mode_enum: RunningMode enum (IMAGE or VIDEO). + **kwargs: Detector-specific configuration parameters. + + Returns: + Configured ImageSegmenterOptions object. """ - if not model_path: - raise ValueError("A valid model_path must be provided.") + output_confidence_masks = kwargs.get('output_confidence_masks', True) + output_category_mask = kwargs.get('output_category_mask', False) + + return vision.ImageSegmenterOptions( + base_options=base_options, + running_mode=mode_enum, + output_confidence_masks=output_confidence_masks, + output_category_mask=output_category_mask, + ) - self.model_path = model_path - # Keep track of the initialized segmenter instance and its config - self._detector_instance = None - self.current_config = None - self.delegate = None # Store delegate used - self._timestamp_provider = None # Added - - def _initialize_segmenter(self, output_confidence_masks: bool, output_category_mask: bool, running_mode: str, delegate_mode="cpu"): - """Initializes the ImageSegmenter or returns existing if config matches. - Handles GPU delegate selection. + def _get_options_tuple(self, running_mode: str = None, + delegate: str = None, **kwargs) -> tuple: + """Get tuple of options for comparison to detect config changes. + + Args: + running_mode: String representation of running mode ("image" or "video"). + delegate: String representation of delegate ("cpu" or "gpu"). + **kwargs: Detector-specific parameters to include in comparison. + + Returns: + Tuple containing all configuration options for comparison. """ - # Return type removed from signature as it's handled internally - new_config = (output_confidence_masks, output_category_mask, running_mode, delegate_mode) + output_confidence_masks = kwargs.get('output_confidence_masks', True) + output_category_mask = kwargs.get('output_category_mask', False) + + # Windows GPU check - provide warning and fallback + if delegate and delegate.lower() == "gpu" and platform.system().lower() == "windows": + logger.warning("GPU delegate is not supported on Windows for Python MediaPipe Tasks. Falling back to CPU.") + delegate = "cpu" + + return (output_confidence_masks, output_category_mask, running_mode, delegate) - # Re-initialize if config changed or not initialized - if self._detector_instance is None or self.current_config != new_config: - if self._detector_instance: - logger.info("Closing existing ImageSegmenter due to config change.") - try: - self._detector_instance.close() - except Exception as e: - logger.warning(f"Error closing previous segmenter: {e}") - self._detector_instance = None + def _process_detection_result(self, detection_result) -> tuple: + """Process segmentation results. + + Args: + detection_result: Raw segmentation result from MediaPipe. + + Returns: + Tuple of (confidence_masks, category_mask) as torch tensors or None. + """ + confidence_masks = None + category_mask = None + + if self._current_options and hasattr(self._current_options, 'output_confidence_masks') and self._current_options.output_confidence_masks: + if detection_result and detection_result.confidence_masks: + confidence_masks = [] + for mask in detection_result.confidence_masks: + mask_np = mask.numpy_view() + mask_tensor = torch.from_numpy(mask_np).float() # Keep as HW + confidence_masks.append(mask_tensor) + + if self._current_options and hasattr(self._current_options, 'output_category_mask') and self._current_options.output_category_mask: + if detection_result and detection_result.category_mask: + mask_np = detection_result.category_mask.numpy_view() + category_mask = torch.from_numpy(mask_np).long() # Keep as HW + + return confidence_masks, category_mask - logger.info( - f"Initializing ImageSegmenter: model='{self.model_path}', conf_mask={output_confidence_masks}, cat_mask={output_category_mask}, mode={running_mode}, delegate={delegate_mode}" - ) + def _create_detector_instance(self, options: vision.ImageSegmenterOptions) -> vision.ImageSegmenter: + """Create segmenter instance. + + Args: + options: Configured ImageSegmenterOptions. + + Returns: + MediaPipe ImageSegmenter instance. + """ + try: + return vision.ImageSegmenter.create_from_options(options) + except Exception as e: + logger.error(f"Failed to create ImageSegmenter: {e}") + raise RuntimeError(f"Failed to initialize MediaPipe ImageSegmenter: {e}") - # Configure BaseOptions with delegate - use_gpu = delegate_mode.lower() == "gpu" - if use_gpu and platform.system().lower() == "windows": - logger.warning("GPU delegate is not supported on Windows for Python MediaPipe Tasks. Falling back to CPU.") - use_gpu = False - delegate_mode = "cpu" + def detect(self, image: torch.Tensor, running_mode: str = "video", + delegate: str = "cpu", **kwargs) -> List[tuple]: + """Override of the BaseDetector detect method to use segment/segment_for_video. + + Args: + image: Input image tensor in BHWC format. + running_mode: "image" or "video". + delegate: "cpu" or "gpu". + **kwargs: Detector-specific parameters. + + Returns: + A list (batch) of tuples of (confidence_masks, category_mask). + """ + if image.dim() != 4: + raise ValueError("Input tensor must be in BHWC format.") - delegate = python.BaseOptions.Delegate.GPU if use_gpu else python.BaseOptions.Delegate.CPU - base_options = python.BaseOptions(model_asset_path=self.model_path, delegate=delegate) + # Check if options changed to decide whether to reinitialize + new_options_tuple = self._get_options_tuple(running_mode=running_mode, delegate=delegate, **kwargs) + current_options_tuple = None + if self._current_options: + current_mode = "image" if self._current_options.running_mode == vision.RunningMode.IMAGE else "video" + current_delegate = "cpu" if self._current_options.base_options.delegate == python.BaseOptions.Delegate.CPU else "gpu" + current_options_tuple = self._get_options_tuple(running_mode=current_mode, delegate=current_delegate) + + if self._detector_instance is None or new_options_tuple != current_options_tuple: + if self._detector_instance and hasattr(self._detector_instance, "close"): + self._detector_instance.close() + self._detector_instance = self._initialize_detector(running_mode, delegate, **kwargs) + + batch_size = image.shape[0] + batch_results = [] + is_video_mode = self._current_options.running_mode == vision.RunningMode.VIDEO + + # Use segment/segment_for_video methods directly + segment_func = self._detector_instance.segment_for_video if is_video_mode else self._detector_instance.segment + + # Process each image in the batch + for i in range(batch_size): + img_tensor = image[i] + np_image = (img_tensor.cpu().numpy() * 255).astype(np.uint8) + mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=np_image) try: - mode_enum = vision.RunningMode.IMAGE if running_mode == "image" else vision.RunningMode.VIDEO - options = vision.ImageSegmenterOptions( - base_options=base_options, - running_mode=mode_enum, - output_confidence_masks=output_confidence_masks, - output_category_mask=output_category_mask, - ) - self._detector_instance = vision.ImageSegmenter.create_from_options(options) - self.current_config = new_config - self.delegate = delegate_mode - self._timestamp_provider = TimestampProvider() if mode_enum == vision.RunningMode.VIDEO else None - logger.info("ImageSegmenter initialized successfully.") + if is_video_mode: + timestamp_ms = self._timestamp_provider.next() + detection_result = segment_func(mp_image, timestamp_ms) + else: + detection_result = segment_func(mp_image) + + batch_results.append(self._process_detection_result(detection_result)) except Exception as e: - logger.error(f"Failed to create ImageSegmenter: {e}") - # Raise error to prevent using a failed instance - raise RuntimeError(f"Failed to initialize MediaPipe ImageSegmenter: {e}") - else: - # logger.debug("Reusing existing ImageSegmenter instance.") # Optional: for debugging - pass + logger.error(f"Error during segmentation for image {i}: {e}") + # Append None results for this image to maintain batch integrity + batch_results.append((None, None)) - # No explicit return needed, instance is stored in self._detector_instance + return batch_results def segment( self, @@ -118,73 +193,28 @@ class ImageSegmenter: """ if image.dim() != 4: raise ValueError("Input tensor must be in BHWC format.") - # if image.shape[0] != 1: # Temporarily removed batch size 1 restriction - # logger.warning("ImageSegmenter currently processes batch > 1 sequentially. Performance may vary.") - # raise ValueError("Input tensor must be in BHWC format with batch size 1.") if not output_confidence_masks and not output_category_mask: raise ValueError("At least one output type (confidence or category mask) must be enabled.") - # Initialize (or get existing) segmenter instance - # Initialization now happens within the function if needed - self._initialize_segmenter(output_confidence_masks, output_category_mask, running_mode, delegate_mode) - if not self._detector_instance: - # Initialization failed, return None or raise error - logger.error("Segmenter instance is not available after initialization attempt.") - # Return structure matching expected output format (list per batch item) - return [None] * image.shape[0], [None] * image.shape[0] - - batch_size = image.shape[0] - batch_results_confidence = [] if output_confidence_masks else None - batch_results_category = [] if output_category_mask else None - - # Determine the correct segmentation function based on mode - is_video_mode = self.current_config[2] == "video" - segment_func = self._detector_instance.segment_for_video if is_video_mode else self._detector_instance.segment - - for i in range(batch_size): - img_tensor = image[i] # HWC - # Convert image format - np_image = (img_tensor.cpu().numpy() * 255).astype(np.uint8) - mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=np_image) - - # Perform segmentation - segmentation_result = None - try: - if is_video_mode: - timestamp_ms = self._timestamp_provider.next() - segmentation_result = segment_func(mp_image, timestamp_ms) - else: - segmentation_result = segment_func(mp_image) - except Exception as e: - logger.error(f"Error during segmentation for image {i}: {e}") - # Append None results for this image - if output_confidence_masks: - batch_results_confidence.append(None) - if output_category_mask: - batch_results_category.append(None) - continue # Skip to the next image - - # Process results for the current image - current_image_confidence_masks = None - current_image_category_mask = None - - if output_confidence_masks: - if segmentation_result and segmentation_result.confidence_masks: - current_image_confidence_masks = [] - for mask in segmentation_result.confidence_masks: - mask_np = mask.numpy_view() - mask_tensor = torch.from_numpy(mask_np).float() # Keep as HW - current_image_confidence_masks.append(mask_tensor) - batch_results_confidence.append(current_image_confidence_masks) - - if output_category_mask: - if segmentation_result and segmentation_result.category_mask: - mask_np = segmentation_result.category_mask.numpy_view() - # Category mask is uint8, convert to long tensor - current_image_category_mask = torch.from_numpy(mask_np).long() # Keep as HW - batch_results_category.append(current_image_category_mask) - - return batch_results_confidence, batch_results_category + # Use the generic detect method from BaseDetector + batch_results = self.detect( + image, + running_mode=running_mode, + delegate=delegate_mode, + output_confidence_masks=output_confidence_masks, + output_category_mask=output_category_mask + ) + + # Unzip the batch results for backward compatibility + confidence_masks_batch = [] + category_mask_batch = [] + + for result_pair in batch_results: + confidence_masks, category_mask = result_pair + confidence_masks_batch.append(confidence_masks) + category_mask_batch.append(category_mask) + + return confidence_masks_batch, category_mask_batch def close(self): """Closes the segmenter instance if it exists.""" @@ -195,8 +225,7 @@ class ImageSegmenter: except Exception as e: logger.warning(f"Error closing segmenter instance: {e}") self._detector_instance = None - self.current_config = None - self.delegate = None + self._current_options = None def __del__(self): # Ensure resources are released when the object is destroyed diff --git a/src/mediapipe_vision/object_detection/detector.py b/src/mediapipe_vision/object_detection/detector.py index 032e9a3..d5f9cfc 100644 --- a/src/mediapipe_vision/object_detection/detector.py +++ b/src/mediapipe_vision/object_detection/detector.py @@ -1,4 +1,4 @@ -from typing import List +from typing import Any, List 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 new types +from ..common import BaseDetector from ..types import BoundingBox, ObjectDetectionCategory, ObjectDetectionResult -class ObjectDetector: +class ObjectDetector(BaseDetector[ObjectDetectionResult]): """Detects objects using MediaPipe ObjectDetector.""" def __init__(self, model_path: str): @@ -30,22 +29,61 @@ class ObjectDetector: self._current_options = None self._timestamp_provider = None # Added - def _initialize_detector(self, score_threshold: float, max_results: int, running_mode: str, delegate: str): # Added running_mode - """Initializes the ObjectDetector.""" - mode_enum = vision.RunningMode.IMAGE if running_mode == "image" else vision.RunningMode.VIDEO # Added - 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.ObjectDetectorOptions( + def _create_detector_options(self, base_options: python.BaseOptions, + mode_enum: vision.RunningMode, **kwargs) -> vision.ObjectDetectorOptions: + """Create ObjectDetector-specific options with parameters: + - score_threshold: Minimum score for object detection + - max_results: Maximum number of detection results + """ + return vision.ObjectDetectorOptions( base_options=base_options, - running_mode=mode_enum, # Use enum mode - score_threshold=score_threshold, - max_results=max_results, + running_mode=mode_enum, + score_threshold=kwargs.get('score_threshold', 0.5), + max_results=kwargs.get('max_results', 5), ) - self._current_options = options - # 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.ObjectDetectorOptions) -> vision.ObjectDetector: return vision.ObjectDetector.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.score_threshold, + self._current_options.max_results, + running_mode, + delegate, + ) + else: + return ( + kwargs.get('score_threshold', 0.5), + kwargs.get('max_results', 5), + running_mode, + delegate, + ) + + def _process_detection_result(self, detection_result: Any) -> List[ObjectDetectionResult]: + current_image_detections = [] + if detection_result and detection_result.detections: + for detection in detection_result.detections: + bbox_mp = detection.bounding_box + categories_mp = detection.categories + + bbox = BoundingBox(origin_x=bbox_mp.origin_x, origin_y=bbox_mp.origin_y, width=bbox_mp.width, height=bbox_mp.height) + + categories = [ + ObjectDetectionCategory( + index=cat.index, + score=cat.score, + display_name=cat.display_name, + category_name=cat.category_name, + ) + for cat in categories_mp + ] + + current_image_detections.append(ObjectDetectionResult(bounding_box=bbox, categories=categories)) + + return current_image_detections + def detect( self, image: torch.Tensor, @@ -53,69 +91,12 @@ class ObjectDetector: max_results: int = 5, running_mode: str = "video", delegate: str = "cpu", - ) -> List[List[ObjectDetectionResult]]: # Added running_mode - """Detects objects in the input image tensor. - - Returns: - A list (batch) of lists (detections per image) of ObjectDetectionResult objects. - """ - if image.dim() != 4: - raise ValueError("Input tensor must be in BHWC format.") - - # Add running_mode to tuple check - new_options_tuple = (score_threshold, max_results, running_mode, delegate.lower()) - current_options_tuple = None - if self._current_options: - current_options_tuple = ( - self._current_options.score_threshold, - self._current_options.max_results, - "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: - if self._detector_instance and hasattr(self._detector_instance, "close"): - self._detector_instance.close() - self._detector_instance = self._initialize_detector(score_threshold, max_results, 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: - timestamp_ms = self._timestamp_provider.next() - detection_result = detect_func(mp_image, timestamp_ms) - else: - detection_result = detect_func(mp_image) - - current_image_detections = [] - if detection_result and detection_result.detections: # Added check for detection_result - for detection in detection_result.detections: - bbox_mp = detection.bounding_box - categories_mp = detection.categories - - bbox = BoundingBox(origin_x=bbox_mp.origin_x, origin_y=bbox_mp.origin_y, width=bbox_mp.width, height=bbox_mp.height) - - categories = [ - ObjectDetectionCategory( - index=cat.index, - score=cat.score, - display_name=cat.display_name, - category_name=cat.category_name, - ) - for cat in categories_mp - ] - - current_image_detections.append(ObjectDetectionResult(bounding_box=bbox, categories=categories)) - - batch_results.append(current_image_detections) - - # self._detector_instance.close() # Close handled on re-init - return batch_results + ) -> List[List[ObjectDetectionResult]]: + """Detects objects in the input image tensor.""" + return super().detect( + image, + running_mode=running_mode, + delegate=delegate, + score_threshold=score_threshold, + max_results=max_results, + ) diff --git a/src/mediapipe_vision/pose_landmark/detector.py b/src/mediapipe_vision/pose_landmark/detector.py index 405b685..a793543 100644 --- a/src/mediapipe_vision/pose_landmark/detector.py +++ b/src/mediapipe_vision/pose_landmark/detector.py @@ -1,4 +1,4 @@ -from typing import List, Optional, Tuple +from typing import Any, List, Optional, Tuple import mediapipe as mp import numpy as np @@ -8,56 +8,108 @@ from mediapipe.tasks.python import vision from mediapipe.tasks.python.core.base_options import BaseOptions from ...utils.timing import TimestampProvider - -# Import new types +from ..common import BaseDetector from ..types import LandmarkPoint, PoseLandmarksResult -class PoseLandmarkDetector: +class PoseLandmarkDetector(BaseDetector[PoseLandmarksResult]): """Detects pose landmarks in an image using MediaPipe PoseLandmarker.""" def __init__(self, model_path: str): - """Initialize the detector with the model path. + """Initialize the detector with the model path.""" + super().__init__(model_path) + self._output_segmentation_masks = False + self._segmentation_masks = [] - Args: - model_path: Path to the MediaPipe PoseLandmarker .task file. + def _create_detector_options(self, base_options: python.BaseOptions, + mode_enum: vision.RunningMode, **kwargs) -> vision.PoseLandmarkerOptions: + """Create PoseLandmarker-specific options with parameters: + - num_poses: Maximum number of poses to detect + - min_detection_confidence: Minimum confidence for pose detection + - min_presence_confidence: Minimum confidence for pose presence + - min_tracking_confidence: Minimum confidence for pose tracking + - output_segmentation_masks: Whether to output segmentation masks """ - if not model_path: - raise ValueError("A valid model_path must be provided.") - - self.model_path = model_path - self._detector_instance = None - self._current_options = None - self._timestamp_provider = None # Added - - def _initialize_detector( - self, - num_poses: int, - min_detection_confidence: float, - min_presence_confidence: float, - min_tracking_confidence: float, - output_segmentation_masks: bool, - running_mode: str, - delegate: str, - ) -> vision.PoseLandmarker: # Added running_mode - """Initializes the PoseLandmarker detector.""" - mode_enum = vision.RunningMode.IMAGE if running_mode == "image" else vision.RunningMode.VIDEO # Added - 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.PoseLandmarkerOptions( + self._output_segmentation_masks = kwargs.get('output_segmentation_masks', False) + return vision.PoseLandmarkerOptions( base_options=base_options, - running_mode=mode_enum, # Use enum mode - num_poses=num_poses, - min_pose_detection_confidence=min_detection_confidence, - min_pose_presence_confidence=min_presence_confidence, - min_tracking_confidence=min_tracking_confidence, - output_segmentation_masks=output_segmentation_masks, + running_mode=mode_enum, + num_poses=kwargs.get('num_poses', 1), + min_pose_detection_confidence=kwargs.get('min_detection_confidence', 0.5), + min_pose_presence_confidence=kwargs.get('min_presence_confidence', 0.5), + min_tracking_confidence=kwargs.get('min_tracking_confidence', 0.5), + output_segmentation_masks=self._output_segmentation_masks, ) - self._current_options = options - # 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.PoseLandmarkerOptions) -> vision.PoseLandmarker: return vision.PoseLandmarker.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_poses, + self._current_options.min_pose_detection_confidence, + self._current_options.min_pose_presence_confidence, + self._current_options.min_tracking_confidence, + self._current_options.output_segmentation_masks, + running_mode, + delegate, + ) + else: + return ( + kwargs.get('num_poses', 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_segmentation_masks', False), + running_mode, + delegate, + ) + + def _process_detection_result(self, detection_result: Any) -> List[PoseLandmarksResult]: + """Process detection result and extract segmentation masks if requested.""" + self._segmentation_masks = [] # Store masks separately + + current_image_landmarks = [] + if detection_result and detection_result.pose_landmarks: + for pose_idx, pose_landmarks_mp in enumerate(detection_result.pose_landmarks): + landmarks = [ + LandmarkPoint( + index=lm_idx, + x=lm.x, + y=lm.y, + z=lm.z, + visibility=getattr(lm, "visibility", None), + presence=getattr(lm, "presence", None), + ) + for lm_idx, lm in enumerate(pose_landmarks_mp) + ] + + world_landmarks = None + if detection_result.pose_world_landmarks and pose_idx < len(detection_result.pose_world_landmarks): + world_landmarks = [ + LandmarkPoint( + index=lm_idx, + x=lm.x, + y=lm.y, + z=lm.z, + visibility=getattr(lm, "visibility", None), + presence=getattr(lm, "presence", None), + ) + for lm_idx, lm in enumerate(detection_result.pose_world_landmarks[pose_idx]) + ] + + current_image_landmarks.append(PoseLandmarksResult(landmarks=landmarks, world_landmarks=world_landmarks)) + + # Handle segmentation masks separately + if self._output_segmentation_masks and detection_result.segmentation_masks: + mask_mp = detection_result.segmentation_masks[0] # Get the first mask + mask_np = mask_mp.numpy_view() + mask_tensor = torch.from_numpy(mask_np).float() # HW + self._segmentation_masks = [mask_tensor] # Store as a list for consistency + + return current_image_landmarks + def detect( self, image: torch.Tensor, @@ -68,118 +120,32 @@ class PoseLandmarkDetector: output_segmentation_masks: bool = False, running_mode: str = "video", delegate: str = "cpu", - ) -> Tuple[List[List[PoseLandmarksResult]], Optional[List[List[torch.Tensor]]]]: # Added running_mode - """Detects pose landmarks in the input image tensor. - - Returns: - A tuple containing: - - pose_landmarks_batch: List (batch) of lists (poses) of PoseLandmarksResult objects. - - segmentation_masks_batch: List (batch) of lists (poses) of segmentation masks (HW FloatTensor) (if requested). - """ - if image.dim() != 4: - raise ValueError("Input tensor must be in BHWC format.") - - # Add running_mode to tuple check - new_options_tuple = ( - num_poses, - min_detection_confidence, - min_presence_confidence, - min_tracking_confidence, - output_segmentation_masks, - running_mode, - delegate.lower(), + ) -> Tuple[List[List[PoseLandmarksResult]], Optional[List[List[torch.Tensor]]]]: + """Detects pose landmarks in the input image tensor.""" + batch_landmarks = super().detect( + image, + running_mode=running_mode, + delegate=delegate, + num_poses=num_poses, + min_detection_confidence=min_detection_confidence, + min_presence_confidence=min_presence_confidence, + min_tracking_confidence=min_tracking_confidence, + output_segmentation_masks=output_segmentation_masks, ) - current_options_tuple = None - if self._current_options: - current_options_tuple = ( - self._current_options.num_poses, - self._current_options.min_pose_detection_confidence, - self._current_options.min_pose_presence_confidence, - self._current_options.min_tracking_confidence, - self._current_options.output_segmentation_masks, - "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: - if self._detector_instance and hasattr(self._detector_instance, "close"): - self._detector_instance.close() - # Pass new params to initializer - self._detector_instance = self._initialize_detector( - num_poses, - min_detection_confidence, - min_presence_confidence, - min_tracking_confidence, - output_segmentation_masks, - running_mode, - delegate, - ) # Pass running_mode - - batch_size = image.shape[0] - batch_results_landmarks = [] - batch_results_masks = [] if output_segmentation_masks else None - - # 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: - timestamp_ms = self._timestamp_provider.next() - detection_result = detect_func(mp_image, timestamp_ms) - else: - detection_result = detect_func(mp_image) - - current_image_landmarks = [] - current_image_masks = [] if output_segmentation_masks else None - - if detection_result and detection_result.pose_landmarks: # Added check for detection_result - for pose_idx, pose_landmarks_mp in enumerate(detection_result.pose_landmarks): - landmarks = [ - LandmarkPoint( - index=lm_idx, - x=lm.x, - y=lm.y, - z=lm.z, - visibility=getattr(lm, "visibility", None), # Use getattr for safety - presence=getattr(lm, "presence", None), - ) - for lm_idx, lm in enumerate(pose_landmarks_mp) - ] - - world_landmarks = None - if detection_result.pose_world_landmarks and pose_idx < len(detection_result.pose_world_landmarks): - world_landmarks = [ - LandmarkPoint( - index=lm_idx, - x=lm.x, - y=lm.y, - z=lm.z, - visibility=getattr(lm, "visibility", None), - presence=getattr(lm, "presence", None), - ) - for lm_idx, lm in enumerate(detection_result.pose_world_landmarks[pose_idx]) - ] - - current_image_landmarks.append(PoseLandmarksResult(landmarks=landmarks, world_landmarks=world_landmarks)) - - # Handle segmentation masks separately as they are not indexed per pose in the result - if output_segmentation_masks and detection_result.segmentation_masks: - # MediaPipe returns one mask per image if output_segmentation_masks is True - # We store it as a list associated with the image, not per pose - # Assuming only one mask is relevant per image, might need adjustment if multiple masks are returned - mask_mp = detection_result.segmentation_masks[0] # Get the first mask - mask_np = mask_mp.numpy_view() - mask_tensor = torch.from_numpy(mask_np).float() # HW - current_image_masks = [mask_tensor] # Store as a list for consistency with return type - - batch_results_landmarks.append(current_image_landmarks) - if output_segmentation_masks: - batch_results_masks.append(current_image_masks) - - # self._detector_instance.close() # Close handled on re-init - return batch_results_landmarks, batch_results_masks + + # Handle segmentation masks + batch_masks = None + if output_segmentation_masks: + batch_size = image.shape[0] + batch_masks = [] + for i in range(batch_size): + # We need to build the segmentation masks batch separately + # since the base detector doesn't know about them + if hasattr(self, '_segmentation_masks') and self._segmentation_masks: + batch_masks.append(self._segmentation_masks) + else: + batch_masks.append([]) + + self._segmentation_masks = [] # Clear after use + + return batch_landmarks, batch_masks