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