segmenter
This commit is contained in:
RyanOnTheInside
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user