134 lines
6.1 KiB
Python
134 lines
6.1 KiB
Python
"""Base node wrapper for MediaPipe Vision detectors."""
|
|
|
|
import torch
|
|
import logging
|
|
from typing import Any, Dict, List, Tuple, Optional, Type, Callable, Union
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
class BaseMediaPipeDetectorNode:
|
|
"""Base class for all MediaPipe detector nodes in ComfyUI.
|
|
|
|
This class captures common patterns across detector nodes like:
|
|
- Detector instance management
|
|
- Model validation
|
|
- Common input parameters
|
|
- Cleanup in __del__
|
|
|
|
Parameter Naming Standards for Subclasses:
|
|
-----------------------------------------
|
|
To maintain consistency across detector types, use these standardized parameter names:
|
|
|
|
- min_confidence: Main confidence threshold for detection (replaces various min_detection_confidence,
|
|
score_threshold, etc.)
|
|
- min_tracking_confidence: Confidence threshold for tracking objects across video frames
|
|
- min_presence_confidence: Confidence threshold for presence of an object in a frame
|
|
- max_results: Maximum number of detections to return (replaces num_poses, num_hands, etc.)
|
|
|
|
When implementing specific detector nodes, map these standardized parameter names to
|
|
the appropriate underlying MediaPipe API parameters in the detect() method.
|
|
"""
|
|
|
|
# To be overridden by subclasses
|
|
DETECTOR_CLASS = None # Subclasses must set to their specific detector class
|
|
MODEL_INFO_TYPE = "MEDIAPIPE_MODEL_INFO" # Override with specific model info type
|
|
EXPECTED_TASK_TYPE = None # Must be set to match the MODEL_INFO_TYPE task_type field
|
|
CATEGORY = "MediaPipeVision/Base" # Override with specific category
|
|
RETURN_TYPES = tuple() # Must be set by subclass
|
|
RETURN_NAMES = tuple() # Must be set by subclass
|
|
FUNCTION = "detect" # Default function name, can be overridden
|
|
DESCRIPTION = "" # Override with specific node description
|
|
|
|
def __init__(self):
|
|
"""Initialize with empty detector and model path."""
|
|
self._detector = None
|
|
self._model_path = None
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
"""Define common input types for detector nodes."""
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE", {"tooltip": "Input image for detection"}),
|
|
"model_info": (cls.MODEL_INFO_TYPE, {"tooltip": f"Model loaded from the appropriate loader node"}),
|
|
"running_mode": (["image", "video"], {
|
|
"default": "video",
|
|
"tooltip": "Processing mode: 'image' for single images, 'video' for sequences/streams"
|
|
}),
|
|
"delegate": (["cpu", "gpu"], {
|
|
"default": "cpu",
|
|
"tooltip": "Processing device - GPU is faster but CPU is more compatible with all systems"
|
|
}),
|
|
}
|
|
}
|
|
# Subclasses should extend this with additional parameters
|
|
|
|
def validate_model_info(self, model_info: Dict) -> str:
|
|
"""Validate model_info and return model path."""
|
|
if not isinstance(model_info, dict):
|
|
raise ValueError(f"Invalid model_info provided. Expected dictionary, got {type(model_info)}")
|
|
|
|
task_type = model_info.get('task_type')
|
|
if task_type != self.EXPECTED_TASK_TYPE:
|
|
raise ValueError(f"Invalid model_info provided. Expected task_type '{self.EXPECTED_TASK_TYPE}', got '{task_type}'")
|
|
|
|
model_path = model_info.get('model_path')
|
|
if not model_path:
|
|
raise ValueError("Model path not found or invalid in model_info")
|
|
|
|
return model_path
|
|
|
|
def initialize_or_update_detector(self, model_path: str):
|
|
"""Initialize detector or update if model path changed."""
|
|
if self.DETECTOR_CLASS is None:
|
|
raise NotImplementedError("Subclasses must set DETECTOR_CLASS to their specific detector class")
|
|
|
|
if self._detector is None or self._model_path != model_path:
|
|
# Close existing detector if needed
|
|
if self._detector and hasattr(self._detector, 'close'):
|
|
try:
|
|
logger.info(f"Closing existing detector for {self._model_path}")
|
|
self._detector.close()
|
|
except Exception as e:
|
|
logger.warning(f"Error closing detector: {e}")
|
|
|
|
# Create new detector
|
|
logger.info(f"Creating new {self.DETECTOR_CLASS.__name__} instance for {model_path}")
|
|
self._detector = self.DETECTOR_CLASS(model_path)
|
|
self._model_path = model_path
|
|
elif not hasattr(self._detector, '_detector_instance'):
|
|
# Re-initialize if internal instance is missing
|
|
logger.warning(f"Re-initializing detector wrapper due to missing internal instance")
|
|
self._detector = self.DETECTOR_CLASS(model_path)
|
|
self._model_path = model_path
|
|
|
|
return self._detector
|
|
|
|
def detect(self, **kwargs):
|
|
"""
|
|
Main detection function, should be implemented by subclasses.
|
|
This is a placeholder that subclasses can use as a template.
|
|
"""
|
|
# Basic implementation that subclasses can override
|
|
|
|
# 1. Extract and validate model_info
|
|
model_info = kwargs.get('model_info')
|
|
model_path = self.validate_model_info(model_info)
|
|
|
|
# 2. Initialize or update detector
|
|
detector = self.initialize_or_update_detector(model_path)
|
|
|
|
# 3. Call detector's detect method with appropriate parameters
|
|
# This part should be implemented by subclasses
|
|
|
|
raise NotImplementedError("Subclasses must implement or override the detect method")
|
|
|
|
def __del__(self):
|
|
"""Clean up resources by closing detector."""
|
|
if hasattr(self, '_detector') and self._detector and hasattr(self._detector, 'close'):
|
|
try:
|
|
logger.info(f"Closing detector instance in __del__ for {self._model_path}")
|
|
self._detector.close()
|
|
except Exception as e:
|
|
logger.warning(f"Error closing detector in __del__: {e}")
|
|
self._detector = None |