nodes/ryanontheinside-ComfyUI_RealtimeNodes
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Merge pull request #33 from ryanontheinside/refactor/base-detector

Browse Source 
Refactor/base detector
This commit is contained in:
RyanOnTheInside
2025-05-16 15:46:55 -04:00
committed by GitHub
parent 49c966ab04 0e11043a72
commit 91b155a115
11 changed files with 1021 additions and 763 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/mediapipe_vision/common/__init__.py
+3
View File
@@ -0,0 +1,3 @@
from .base_detector import BaseDetector
__all__ = ["BaseDetector"]
src/mediapipe_vision/common/base_detector.py
+175
View File
@@ -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
src/mediapipe_vision/face_detection/detector.py
+52 -81
View File
@@ -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,
)
src/mediapipe_vision/face_landmark/detector.py
+84 -10
View File
@@ -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:
src/mediapipe_vision/face_stylization/detector.py
+51 -29
View File
@@ -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)
src/mediapipe_vision/gesture_recognition/detector.py
+98 -102
View File
@@ -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
src/mediapipe_vision/hand_landmark/detector.py
+73 -111
View File
@@ -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
src/mediapipe_vision/image_embedding/detector.py
+149 -70
View File
@@ -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
)
src/mediapipe_vision/image_segmentation/segmenter.py
+154 -125
View File
@@ -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
src/mediapipe_vision/object_detection/detector.py
+63 -82
View File
@@ -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,
)
src/mediapipe_vision/pose_landmark/detector.py
+119 -153
View File
@@ -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