feat: add mediapipe vision task

This commit is contained in:
RyanOnTheInside
2025-04-19 16:04:47 -04:00
parent b2342bc19e
commit d38341f7ea
89 changed files with 6922 additions and 53 deletions
+100 -17
View File
@@ -4,7 +4,8 @@ A growing suite of nodes for real-time ComfyUI workflows. Features include value
The intention for this repository is to build a suite of nodes that can be used in the burgeoning real-time diffusion space. Contributions are welcome!
## Nodes
## Control Nodes
### Value Controls 🎚️
- **FloatControl**: Outputs a floating point value that changes over time using various patterns (sine wave, bounce, random walk, etc).
@@ -27,17 +28,11 @@ The intention for this repository is to build a suite of nodes that can be used
- **DTypeConverter**: Convert masks between different data types (float16, uint8, float32, float64).
- **FastWebcamCapture**: High-performance webcam capture node with resizing capabilities.
- **SimilarityFilter**: Filter out similar consecutive images and control downstream execution. Perfect for optimizing real-time workflows by skipping redundant processing of similar frames.
### Logic 🧠
- **LazyCondition**: Powerful conditional execution node that supports any input type. Uses lazy evaluation to truly skip execution of unused paths and maintains state to avoid feedback loops.
## Movement Patterns 🔄
All value and motion controls support various movement patterns:
- **Sine**: Smooth sinusoidal motion
- **Triangle**: Linear interpolation with smooth direction changes
- **Sawtooth**: Linear interpolation with sharp resets
- **Square**: Instant transitions between min/max values
- **Static**: No movement (constant value)
- **and more**
## Usage 📖
@@ -65,6 +60,74 @@ Use utility nodes to optimize and control your workflow:
- **SimilarityFilter**: Skip processing of similar frames by comparing consecutive images. Great for optimizing real-time workflows by only processing frames that have meaningful changes.
- **LazyCondition**: Create conditional execution paths that truly skip processing of unused branches. Works with any input type (images, latents, text, numbers) and maintains state of the last successful output to avoid feedback loops.
## 🔮 MediaPipe Vision
### ✨ Overview
This repository provides a complete implementation of Google MediaPipe vision tasks for ComfyUI. It enables computer vision capabilities that can be used for interactive AI art, responsive interfaces, motion tracking, and advanced masking workflows.
### 🚀 Features
| Category | Available Tools |
|----------|-------------|
| **Face Analysis** | Face detection, face mesh (478 points), blendshapes, head pose |
| **Body Tracking** | Pose estimation (33 landmarks), segmentation masks |
| **Hand Analysis** | Hand tracking (21 landmarks per hand), gesture recognition |
| **Image Processing** | Object detection, image segmentation, image embeddings |
| **Creative Tools** | Face stylization, interactive segmentation |
### 📋 Supported MediaPipe Tasks
* **Face Detection:** Face bounding boxes and keypoints
* **Face Landmark Detection:** Face mesh landmarks with expression analysis
* **Hand Landmark Detection:** Hand position tracking with 21 landmarks
* **Pose Landmark Detection:** Body pose tracking with 33 landmarks
* **Object Detection:** Common object detection using models like EfficientDet
* **Image Segmentation:** Category-based image segmentation
* **Gesture Recognition:** Recognition of common hand gestures
* **Image Embedding:** Feature vector generation for image similarity
* **Interactive Segmentation:** User-guided image masking
* **Face Stylization:** Artistic style application to faces
* **Holistic Landmark Detection:** Full-body landmark detection (legacy)
> **Note:** Holistic landmark detection uses the legacy MediaPipe API as we await the official Tasks API release.
### ⚙️ Landmark System
The project's landmark system allows extracting and using position data:
#### Position Extraction
**Landmark Position Extractors** access coordinate data from any landmark:
- Extract x, y, z positions from face, hand, or pose landmarks
- Access visibility and presence information where available
- Access world coordinates when available (hand and pose)
- Input landmark indices directly to access any point
- Process batches for multi-frame workflows
#### Position Processing
Several node types work with landmark position data:
- **Delta Controls** - Track movement and map changes to parameter values
- **Proximity Nodes** - Calculate distances between landmarks
- **Masking Nodes** - Generate masks centered at landmark positions
- **Head Pose Extraction** - Calculate yaw, pitch, roll from face landmarks
- **Blendshape Analysis** - Extract facial expression parameters
### Example Workflow
```
Load Face Landmarker → Face Landmarker ← Image Input
|
↓ landmarks
Face Landmark Position (Index: 1) → x,y,z coordinates
|
↓ x,y,z
Position Delta Float Control → value → ComfyUI Parameter
```
## Examples 🎬
### Value Control Demo
@@ -97,6 +160,7 @@ git clone https://github.com/ryanontheinside/ComfyUI_RealTimeNodes
cd ComfyUI_RealTimeNodes
pip install -r requirements.txt
```
> **Note:** For MediaPipe, GPU Support varies by platform. For Linux, see [these instructions](https://ai.google.dev/edge/mediapipe/framework/getting_started/gpu_support).
## Coming Soon 🚀
@@ -116,19 +180,37 @@ This is an evolving project that aims to expand the real-time capabilities of Co
### Contributing 🤝
Your feedback and contributions are more than welcome! This project grows stronger with community input.
This project provides flexible infrastructure for computer vision in ComfyUI. If you have ideas for:
- Creative AI interactions using vision
- Specific landmark tracking or detection needs
- Real-time vision workflows
- Improvements to the current implementation
Please open an issue, even if you're not sure how to implement it.
The aim is to **iterate quickly** to keep up with this burgeoning field of real-time ComfyUI
- Have an idea? Open an issue! 💡
- Found a bug? Open an issue! 🐛
- Made an improvement? Submit a PR! 🎉
- Want to help? Join the discussion! 💬
Please visit our [GitHub Issues](https://github.com/ryanontheinside/ComfyUI_RealTimeNodes/issues) page to contribute.
## Related Projects 🔗
### ComfyUI_RyanOnTheInside - Everything Reactivity ⚡
Make anything react to anything in your ComfyUI workflows. [ComfyUI_RyanOnTheInside](https://github.com/ryanontheinside/ComfyUI_RyanOnTheInside) - my main custom nodes suite that brings complete reactive control to standard ComfyUI workflows:
## 🔗 Related Projects
### [ComfyUI_ControlFreak](https://github.com/ryanontheinside/ComfyUI_ControlFreak)
Universal MIDI & Gamepad Mapping in ComfyUI. Map any MIDI controller or gamepad to any parameter in your ComfyUI workflow for intuitive, hands-on control of your generative art. Perfect for live performances, interactive installations, and streamlined creative workflows.
### [comfystream](https://github.com/yondonfu/comfystream)
A real-time streaming framework for ComfyUI that enables running workflows continuously on video streams, perfect for
combining with MediaPipe vision capabilities.
### [ComfyUI-Stream-Pack](https://github.com/livepeer/ComfyUI-Stream-Pack)
A collection of ComfyUI nodes for multimedia streaming applications. Combines video processing with generative models
for real-time media effects.
### [ComfyUI_RyanOnTheInside](https://github.com/ryanontheinside/ComfyUI_RyanOnTheInside) - Everything Reactivity ⚡
Make anything react to anything in your ComfyUI workflows. - my main custom nodes suite that brings complete reactive control to standard ComfyUI workflows:
- Dynamic node relationships
- React to audio, MIDI, motion, time, depth, color, Whisper, and more
@@ -143,4 +225,5 @@ Make anything react to anything in your ComfyUI workflows. [ComfyUI_RyanOnTheIns
- Reactive DepthFlow
- Actually more
Use it alongside these Control Nodes to master parameter control in both the batch and real-time paradigms in ComfyUI! The POWER!!
Use it alongside these Control Nodes to master parameter control in both the batch and real-time paradigms in ComfyUI! The POWER!!
+37 -16
View File
@@ -5,7 +5,7 @@ import re
import traceback
# Directory containing node wrapper modules relative to this __init__.py
WRAPPER_DIR = "node_wrappers.realtimenodes"
WRAPPER_DIR = "node_wrappers"
NODE_CLASS_MAPPINGS = {}
NODE_DISPLAY_NAME_MAPPINGS = {}
@@ -13,7 +13,7 @@ print("\033[93m" + "[ComfyUI_RealtimeNodes] Loading nodes...")
def load_nodes_from_directory(package_path):
"""Dynamically loads nodes from Python files in a directory."""
loaded_nodes = {}
loaded_nodes = {}
loaded_display_names = {}
# Calculate absolute path to the package directory
@@ -61,22 +61,43 @@ def load_nodes_from_directory(package_path):
# --- Main Loading Logic ---
# Define the subdirectories within WRAPPER_DIR to scan
# Order might matter if there are dependencies between modules, though ideally there shouldn't be
subdirs_to_scan = [
"controls",
"media",
"utils",
"specialized" # Add other subdirs like 'specialized' if needed
]
# Get base path for the node_wrappers directory
base_dir_path = os.path.join(os.path.dirname(__file__), WRAPPER_DIR)
# Iterate through specified subdirectories and load nodes
for subdir in subdirs_to_scan:
dir_path = f"{WRAPPER_DIR}.{subdir}"
nodes, display_names = load_nodes_from_directory(dir_path)
NODE_CLASS_MAPPINGS.update(nodes)
NODE_DISPLAY_NAME_MAPPINGS.update(display_names)
# Recursively scan all directories and subdirectories
def scan_for_modules(base_path, base_package):
"""Recursively scan directories for module files"""
try:
for item in os.listdir(base_path):
item_path = os.path.join(base_path, item)
# Skip __pycache__ and files/dirs starting with __
if item.startswith('__'):
continue
if os.path.isdir(item_path):
# Handle subdirectory
subdir_package = f"{base_package}.{item}"
# First load .py files in this directory
try:
nodes, display_names = load_nodes_from_directory(subdir_package)
NODE_CLASS_MAPPINGS.update(nodes)
NODE_DISPLAY_NAME_MAPPINGS.update(display_names)
except Exception as e:
print(f"\033[91m" + f"[ComfyUI_RealtimeNodes] Error processing directory {subdir_package}: {e}")
traceback.print_exc()
# Then recursively process subdirectories
subdir_path = os.path.join(base_path, item)
scan_for_modules(subdir_path, subdir_package)
except Exception as e:
print(f"\033[91m" + f"[ComfyUI_RealtimeNodes] Error scanning directory {base_path}: {e}")
traceback.print_exc()
# Start recursive scanning from the base directory
print(f"\033[92m" + f"[ComfyUI_RealtimeNodes] Scanning node_wrappers directory for modules...")
scan_for_modules(base_dir_path, WRAPPER_DIR)
suffix = " 🕒🅡🅣🅝"
@@ -0,0 +1,134 @@
"""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 = "Realtime Nodes/MediaPipe Vision/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
@@ -0,0 +1,293 @@
import logging
from typing import List, Dict, Tuple, Any, Optional, Type
import cv2
import numpy as np
from ....src.mediapipe_vision.utils.image_utils import convert_to_cv2, convert_to_tensor
logger = logging.getLogger(__name__)
# Default font setup - can be used by all visualization classes
DEFAULT_FONT = cv2.FONT_HERSHEY_SIMPLEX
DEFAULT_FONT_SCALE = 0.5
DEFAULT_FONT_THICKNESS = 1
class BaseLandmarkVisualizationNode:
"""Base class for landmark visualization nodes."""
CATEGORY = "Realtime Nodes/MediaPipe Vision/Visualization"
RETURN_TYPES = ("IMAGE",)
FUNCTION = "draw_visualizations"
# --- To be overridden by subclasses ---
LANDMARKS_TYPE = "ANY" # The specific ComfyUI landmark type string (e.g., "FACE_LANDMARKS")
CONNECTIONS = None # MediaPipe connections list (e.g., mp.solutions.face_mesh.FACEMESH_TESSELATION)
DEFAULT_DRAW_POINTS = True
DEFAULT_POINT_COLOR = "#FFFFFF"
DEFAULT_POINT_RADIUS = 3
DEFAULT_DRAW_CONNECTIONS = True
DEFAULT_CONNECTION_COLOR = "#00FF00"
DEFAULT_CONNECTION_THICKNESS = 2
SUPPORTS_VISIBILITY = False # Does this landmark type provide visibility scores?
# --- End Override Section ---
@classmethod
def INPUT_TYPES(cls):
"""Generate input types based on class properties."""
inputs = {
"required": {
"landmarks": (cls.LANDMARKS_TYPE,),
"image": ("IMAGE",),
},
"optional": {
"draw_points": ("BOOLEAN", {"default": cls.DEFAULT_DRAW_POINTS}),
"point_color": ("STRING", {"default": cls.DEFAULT_POINT_COLOR}),
"point_radius": ("INT", {"default": cls.DEFAULT_POINT_RADIUS, "min": 1, "max": 10}),
"draw_connections": ("BOOLEAN", {"default": cls.DEFAULT_DRAW_CONNECTIONS}),
"connection_color": ("STRING", {"default": cls.DEFAULT_CONNECTION_COLOR}),
"connection_thickness": ("INT", {"default": cls.DEFAULT_CONNECTION_THICKNESS, "min": 1, "max": 10}),
}
}
if cls.SUPPORTS_VISIBILITY:
inputs["optional"]["min_visibility"] = ("FLOAT", {
"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.05,
"tooltip": "Minimum visibility score to draw a point/connection endpoint"
})
return inputs
def draw_visualizations(self, landmarks, image, **kwargs):
"""Main execution function: handles batching, image conversion, and calls drawing logic."""
# Extract relevant optional arguments using class defaults
draw_points = kwargs.get('draw_points', self.DEFAULT_DRAW_POINTS)
point_color = kwargs.get('point_color', self.DEFAULT_POINT_COLOR)
point_radius = kwargs.get('point_radius', self.DEFAULT_POINT_RADIUS)
draw_connections = kwargs.get('draw_connections', self.DEFAULT_DRAW_CONNECTIONS)
connection_color = kwargs.get('connection_color', self.DEFAULT_CONNECTION_COLOR)
connection_thickness = kwargs.get('connection_thickness', self.DEFAULT_CONNECTION_THICKNESS)
min_visibility = kwargs.get('min_visibility', 0.5) # Default if not supported/provided
if not landmarks:
logger.warning(f"{self.__class__.__name__}: No landmarks provided.")
return (image,)
cv2_images = convert_to_cv2(image)
processed_cv2_images = []
batch_size = len(cv2_images)
if len(landmarks) != batch_size:
logger.error(f"{self.__class__.__name__}: Batch size mismatch. Images: {batch_size}, Landmark sets: {len(landmarks)}. Cannot proceed reliably.")
return (image,)
# Convert hex colors to BGR tuples
point_color_bgr = self._hex_to_bgr(point_color)
connection_color_bgr = self._hex_to_bgr(connection_color)
for i, cv2_image in enumerate(cv2_images):
height, width = cv2_image.shape[:2]
landmarks_for_image = landmarks[i] # List of landmark results for this image item
if not landmarks_for_image:
processed_cv2_images.append(cv2_image) # Append unmodified image
continue
# Iterate through all detections (faces, hands, poses) in this image item
for single_detection_result in landmarks_for_image:
self.draw_single_detection(cv2_image, width, height, single_detection_result,
draw_points, point_color_bgr, point_radius,
draw_connections, connection_color_bgr, connection_thickness,
min_visibility)
processed_cv2_images.append(cv2_image)
output_image = convert_to_tensor(processed_cv2_images)
return (output_image,)
def _hex_to_bgr(self, hex_color):
"""Convert hex color string to BGR tuple."""
hex_color = hex_color.lstrip('#')
r, g, b = tuple(int(hex_color[i:i+2], 16) for i in (0, 2, 4))
return (b, g, r) # OpenCV uses BGR
def draw_single_detection(self, cv2_image, width: int, height: int,
detection_result: Any,
draw_points: bool, point_color, point_radius: int,
draw_connections: bool, connection_color, connection_thickness: int,
min_visibility: float):
"""Draws landmarks and connections for a single detected object (face, hand, pose)."""
# Extract the actual list of landmarks (this might vary slightly based on MediaPipe result structure)
landmark_list = self.get_landmark_list(detection_result)
if not landmark_list:
return # Skip if no landmarks found for this detection
points = {}
# Process landmarks: calculate coordinates and optionally draw points
for lm in landmark_list:
# Check visibility if supported
is_visible = True
if self.SUPPORTS_VISIBILITY:
vis = lm.visibility if hasattr(lm, 'visibility') and lm.visibility is not None else 1.0
is_visible = vis >= min_visibility
if is_visible:
cx, cy = int(lm.x * width), int(lm.y * height)
points[lm.index] = (cx, cy) # Store denormalized point coordinates
if draw_points:
cv2.circle(cv2_image, (cx, cy), point_radius, point_color, -1)
# Draw connections if enabled and applicable
if draw_connections and self.CONNECTIONS and points:
for connection in self.CONNECTIONS:
start_idx = connection[0]
end_idx = connection[1]
# Draw connection only if both points are present (i.e., passed visibility check)
if start_idx in points and end_idx in points:
cv2.line(cv2_image, points[start_idx], points[end_idx], connection_color, connection_thickness)
def get_landmark_list(self, detection_result: Any) -> Optional[List[Any]]:
"""Abstract method to extract the list of landmark objects from a single detection result.
Needs to be implemented by subclasses based on the specific MediaPipe result structure.
e.g., for hands/pose: return detection_result.landmarks
for face: return detection_result (it's already the list)
"""
raise NotImplementedError("Subclasses must implement get_landmark_list")
class BaseDetectionVisualizationNode:
"""Base class for visualization nodes that need to draw bounding boxes and labels.
This can be used for object detection, face detection, and similar tasks
that require drawing boxes and text labels with a consistent font rendering approach.
"""
CATEGORY = "Realtime Nodes/MediaPipe Vision/Visualization"
RETURN_TYPES = ("IMAGE",)
FUNCTION = "draw_visualizations"
# Default font settings
FONT = DEFAULT_FONT
FONT_SCALE = DEFAULT_FONT_SCALE
FONT_THICKNESS = DEFAULT_FONT_THICKNESS
@classmethod
def get_text_dimensions(cls, text, font_scale=None, thickness=None):
"""Gets text dimensions using cv2.
Args:
text: The text to measure
font_scale: The font scale (defaults to cls.FONT_SCALE)
thickness: The font thickness (defaults to cls.FONT_THICKNESS)
Returns:
Tuple of (width, height) in pixels
"""
if font_scale is None:
font_scale = cls.FONT_SCALE
if thickness is None:
thickness = cls.FONT_THICKNESS
try:
(text_width, text_height), baseline = cv2.getTextSize(
text, cls.FONT, font_scale, thickness
)
return text_width, text_height + baseline
except Exception as e:
logger.error(f"Error measuring text dimensions: {e}")
return 100, 20 # Default fallback
def _hex_to_bgr(self, hex_color):
"""Convert hex color string to BGR tuple."""
hex_color = hex_color.lstrip('#')
r, g, b = tuple(int(hex_color[i:i+2], 16) for i in (0, 2, 4))
return (b, g, r) # OpenCV uses BGR
def draw_text_with_background(self, image, text, position,
text_color="#FFFFFF", bg_color="#000000",
padding=(4, 2), font_scale=None, thickness=None):
"""Draws text with a background rectangle.
Args:
image: cv2 image to draw on
text: Text to draw
position: (x, y) position for top-left of text
text_color: Hex color for text
bg_color: Hex color for background
padding: (x_pad, y_pad) padding around text
font_scale: Font scale (defaults to cls.FONT_SCALE)
thickness: Font thickness (defaults to cls.FONT_THICKNESS)
"""
if font_scale is None:
font_scale = self.FONT_SCALE
if thickness is None:
thickness = self.FONT_THICKNESS
# Convert hex colors to BGR tuples
text_color_bgr = self._hex_to_bgr(text_color)
bg_color_bgr = self._hex_to_bgr(bg_color)
x, y = position
text_width, text_height = self.get_text_dimensions(text, font_scale, thickness)
# Draw background rectangle
x_pad, y_pad = padding
cv2.rectangle(
image,
(x, y),
(x + text_width + x_pad*2, y + text_height + y_pad*2),
bg_color_bgr,
-1 # Filled rectangle
)
# Draw text
cv2.putText(
image,
text,
(x + x_pad, y + text_height),
self.FONT,
font_scale,
text_color_bgr,
thickness
)
return text_width, text_height
def process_batch(self, detections, image, drawing_function, **kwargs):
"""Process a batch of images and detections.
Args:
detections: List of detection results (one per image)
image: Batch of images as tensor
drawing_function: Function to call for each image
**kwargs: Additional arguments for the drawing function
Returns:
Tuple with processed image tensor
"""
if not detections:
logger.warning(f"{self.__class__.__name__}: No detections provided.")
return (image,)
cv2_images = convert_to_cv2(image)
processed_cv2_images = []
batch_size = len(cv2_images)
if len(detections) != batch_size:
logger.error(f"{self.__class__.__name__}: Batch size mismatch. Images: {batch_size}, Detection sets: {len(detections)}. Cannot proceed.")
return (image,)
for i, cv2_image in enumerate(cv2_images):
height, width = cv2_image.shape[:2]
detections_for_image = detections[i]
if not detections_for_image:
processed_cv2_images.append(cv2_image)
continue
# Call the drawing function with the specific parameters
drawing_function(cv2_image, width, height, detections_for_image, **kwargs)
processed_cv2_images.append(cv2_image)
output_image = convert_to_tensor(processed_cv2_images)
return (output_image,)
@@ -0,0 +1,113 @@
"""Node wrapper for MediaPipe Vision model loader."""
# Import necessary components from the source directory
from ....src.mediapipe_vision.model_loader import MediaPipeModelLoader, get_available_models
import logging
import traceback # For better error logging
# Add import for MEDIAPIPE_MODELS dictionary and the get_model_description function
from ....src.mediapipe_vision.model_loader import MEDIAPIPE_MODELS, get_model_description
logger = logging.getLogger(__name__)
# Define a base class for MediaPipe model loading nodes
class MediaPipeModelLoaderBaseNode:
"""Base ComfyUI node for loading MediaPipe Vision models.
Inheriting classes MUST define the class variable TASK_TYPE.
This node handles selecting the model variant for a specific task,
triggers the download/loading via the src.model_loader,
and returns a dictionary containing model information needed by other nodes.
"""
# Define common category for all MediaPipe loaders
CATEGORY = "Realtime Nodes/MediaPipe Vision/Loaders"
# Inheriting classes MUST set this
TASK_TYPE = None # e.g., "face_detector", "hand_landmarker"
@classmethod
def INPUT_TYPES(cls):
"""Define input types for the node."""
if cls.TASK_TYPE is None:
raise NotImplementedError("Subclasses of MediaPipeModelLoaderBaseNode must define TASK_TYPE")
task_variants = []
default_variant = "default" # Sensible default
try:
task_variants = get_available_models(cls.TASK_TYPE)
if not task_variants:
logger.warning(f"No model variants found for task '{cls.TASK_TYPE}'. Defaulting to ['default'].")
task_variants = ["default"] # Fallback variants
default_variant = task_variants[0] if task_variants else default_variant
except Exception as e:
logger.error(f"Failed to dynamically get available MediaPipe models for task {cls.TASK_TYPE}: {e}\n{traceback.format_exc()}")
# Use hardcoded fallbacks if dynamic loading fails
task_variants = ["default"]
default_variant = "default"
# Create tooltip with model descriptions
tooltip = f"Select model variant for {cls.TASK_TYPE}."
# Add model descriptions if available
tooltip += "\n\nModel Variants:"
for variant in task_variants:
description = get_model_description(cls.TASK_TYPE, variant)
tooltip += f"\n- {variant}: {description}"
return {
"required": {
# Dropdown for selecting the model variant for the specific task
"model_variant": (task_variants, {
"default": default_variant,
"tooltip": tooltip
}),
}
}
# Return a dictionary with model info, not the loader object itself
RETURN_TYPES = ("MEDIAPIPE_MODEL_INFO",) # Default, can be overridden by subclasses for specificity
RETURN_NAMES = ("model_info",)
FUNCTION = "load_and_get_info"
def load_and_get_info(self, model_variant: str): # Removed task_type argument
"""
Loads the specified MediaPipe model using the src loader
and returns a dictionary with model information (path, task, variant).
Uses the TASK_TYPE defined by the subclass.
"""
if self.TASK_TYPE is None:
raise NotImplementedError("TASK_TYPE not defined in the subclass.")
task_type = self.TASK_TYPE # Use the class variable
logger.info(f"Attempting to load MediaPipe model: Task='{task_type}', Variant='{model_variant}'")
# Use a temporary loader instance just to trigger the load/download logic
# The actual detector nodes will instantiate their own detectors using the path
loader = MediaPipeModelLoader()
# Pass the class's task_type here
success = loader.load_model(task_type, model_variant)
if not success:
error_msg = f"Failed to load MediaPipe model for Task='{task_type}', Variant='{model_variant}'"
logger.error(error_msg)
# It's often better to raise an error to stop the workflow
raise RuntimeError(error_msg)
# Get the essential model information
model_info = loader.get_model_info() # Should return {'model_path': ..., 'task_type': ..., 'model_variant': ...}
# Ensure the info dict is populated correctly - get_model_info should ideally return all needed keys
if 'model_path' not in model_info or not model_info['model_path']:
raise RuntimeError(f"Failed to get valid model path from loader for {task_type}/{model_variant}")
if 'task_type' not in model_info: model_info['task_type'] = task_type
if 'model_variant' not in model_info: model_info['model_variant'] = model_variant
logger.info(f"Successfully loaded model info: Path='{model_info.get('model_path')}', Task='{model_info.get('task_type')}', Variant='{model_info.get('model_variant')}'")
# Return the dictionary containing model path and other relevant info
# The specific RETURN_TYPES are handled by the subclass definition
return (model_info,)
@@ -0,0 +1,597 @@
import logging
from typing import Optional, List, Any, Tuple, Union
import math
import torch
import numpy as np
import cv2
import collections
import typing # Add typing for Dict hint
# Assuming types are correctly imported relative to this file's location
from ....src.mediapipe_vision.types import (
FACE_LANDMARKS, HAND_LANDMARKS, POSE_LANDMARKS,
LandmarkPoint, FaceLandmarksResult, HandLandmarksResult, PoseLandmarksResult
)
# Re-import needed utils if not available globally
from ....src.mediapipe_vision.utils.delta_utils import scale_value
from ....src.mediapipe_vision.utils.comfy_utils import AlwaysEqualProxy
logger = logging.getLogger(__name__)
MAX_POSITION_HISTORY = 50 # Matches MAX_LANDMARK_HISTORY for consistency
FLOAT_EQUALITY_TOLERANCE = 1e-6 # Same tolerance as base_delta_nodes
class LandmarkPositionBaseNode:
#NOTE: use x,y,and z for the landmark position rather than LandmarkPoint object for max flexibility
"""
Base class for extracting position lists from specific landmark types.
Outputs values corresponding to the input batch size, using defaults for missing items.
Includes an 'is_valid' flag indicating successful extraction for each batch item.
Processes only the first detected object (e.g., face, hand) per batch item.
RETURN_TYPES declare the semantic type (FLOAT/BOOLEAN); the actual return value may be a list.
"""
FUNCTION = "extract_positions"
# Use base types for semantic declaration
RETURN_TYPES = ("FLOAT", "FLOAT", "FLOAT", "FLOAT", "FLOAT", "BOOLEAN")
RETURN_NAMES = ("x", "y", "z", "visibility", "presence", "is_valid") # Singular names
CATEGORY = "Realtime Nodes/MediaPipe Vision/Common/Utils"
DEFAULT_VALUES = (0.0, 0.0, 0.0, 0.0, 0.0)
LANDMARKS_TYPE = "ANY_LANDMARKS" # Placeholder
# INPUT_TYPES are defined in subclasses
def extract_positions(self, landmarks: list, landmark_index: int, use_world_coordinates: bool = False) -> Tuple[List[float], List[float], List[float], List[float], List[float], List[bool]]:
"""Extracts position and properties for the specified landmark across the entire batch,
padding with defaults and providing validity flags for missing items."""
x_coords = []
y_coords = []
z_coords = []
vis_coords = []
pres_coords = []
is_valid_list = [] # New list to track validity
if not isinstance(landmarks, list):
raise TypeError(f"{self.__class__.__name__}: Input landmarks is not a list. Received type: {type(landmarks)}.")
batch_size = len(landmarks)
# Iterate through the batch (List[...])
for batch_idx in range(batch_size):
item_list = landmarks[batch_idx]
current_x, current_y, current_z, current_vis, current_pres = self.DEFAULT_VALUES
is_valid = False # Assume invalid until proven otherwise
landmark_point: Optional[LandmarkPoint] = None
target_landmarks_list: Optional[List[LandmarkPoint]] = None
# Check if item_list is valid and contains at least one result item
if isinstance(item_list, list) and item_list:
result_item = item_list[0] # Process only the first detected object
if result_item is not None:
try:
# --- Determine Structure and Extract Landmark List ---
# Case 1: result_item is the list itself (FACE_LANDMARKS)
if isinstance(result_item, list):
target_landmarks_list = result_item
if use_world_coordinates and self.LANDMARKS_TYPE == "FACE_LANDMARKS":
logger.warning(f"World coordinates requested but not available for FACE_LANDMARKS. Using image coordinates.")
# Case 2: result_item is an object with .landmarks (HAND/POSE_LANDMARKS)
elif hasattr(result_item, 'landmarks'):
landmarks_attr = getattr(result_item, 'landmarks', None)
if landmarks_attr is not None:
use_world = (use_world_coordinates and
self.LANDMARKS_TYPE != "FACE_LANDMARKS" and
hasattr(result_item, 'world_landmarks') and
getattr(result_item, 'world_landmarks', None))
if use_world:
target_landmarks_list = result_item.world_landmarks
else:
target_landmarks_list = landmarks_attr
# else: target_landmarks_list remains None
# --- Find Landmark Point by Index ---
if isinstance(target_landmarks_list, list) and target_landmarks_list:
landmark_point = next((lm for lm in target_landmarks_list if lm.index == landmark_index), None)
except Exception as e:
logger.error(f"{self.__class__.__name__}: Unexpected error processing batch item {batch_idx}: {e}", exc_info=True)
# Keep landmark_point as None, is_valid as False
# else: item_list was not a list or was empty, keep defaults/invalid
# --- Extract Values or Use Defaults, Set Validity ---
if landmark_point:
is_valid = True # Mark as valid since we found the point
current_x = landmark_point.x
current_y = landmark_point.y
current_z = landmark_point.z
current_vis = getattr(landmark_point, 'visibility', None)
current_pres = getattr(landmark_point, 'presence', None)
current_vis = current_vis if current_vis is not None else 0.0
current_pres = current_pres if current_pres is not None else 0.0
# else: defaults are used, is_valid remains False
# Append values for this batch item
x_coords.append(current_x)
y_coords.append(current_y)
z_coords.append(current_z)
vis_coords.append(current_vis)
pres_coords.append(current_pres)
is_valid_list.append(is_valid)
return (x_coords, y_coords, z_coords, vis_coords, pres_coords, is_valid_list)
# --- Base Position Delta Logic (Handling Multiple Histories) ---
class BasePositionDeltaNode:
"""Base class for nodes calculating delta from position history.
Handles single floats or lists of coordinates.
"""
FUNCTION = "execute"
def __init__(self):
# Store history deques keyed by item index from input list
# Key 0 is used for single float inputs
self._histories: typing.Dict[int, collections.deque] = {}
def _get_or_create_history(self, item_index: int) -> collections.deque:
"""Gets the deque for a specific index, creating it if needed."""
if item_index not in self._histories:
self._histories[item_index] = collections.deque(maxlen=MAX_POSITION_HISTORY)
return self._histories[item_index]
def _get_position_delta(self, x: float, y: float, z: float, window_size: int, item_index: int = 0) -> Optional[float]:
"""Calculates delta using the history deque for the given item_index."""
history_deque = self._get_or_create_history(item_index)
window_size = max(1, min(window_size, MAX_POSITION_HISTORY - 1))
current_pos = (x, y, z)
history_deque.append(current_pos)
if len(history_deque) > window_size:
pos_now = history_deque[-1]
pos_past = history_deque[-(window_size + 1)] # Correct indexing
if pos_now is None or pos_past is None:
logger.warning(f"{self.__class__.__name__}[{item_index}]: Missing position data in history window.")
return None
delta = math.sqrt((pos_now[0] - pos_past[0])**2 +
(pos_now[1] - pos_past[1])**2 +
(pos_now[2] - pos_past[2])**2)
return delta
else:
# logger.debug(f"{self.__class__.__name__}[{item_index}]: Not enough history ({len(history_deque)}/{window_size+1}).")
return None # Not enough history yet
# Remove _handle_input_coords, logic moved to execute
# --- Position Delta Control Nodes (Int/Float) (Handling Lists) ---
class PositionDeltaControlNode(BasePositionDeltaNode):
"""Base class for position-based delta control nodes (Int/Float).
Handles single floats or lists of coordinates.
Outputs FLOAT type (actual value is float or list[float]).
"""
# Output semantic type is FLOAT (or INT for subclass)
RETURN_TYPES = ("FLOAT",)
RETURN_NAMES = ("value",) # Singular name
IS_INT_OUTPUT = False
DEFAULT_DELTA_MIN = 0.0
DEFAULT_DELTA_MAX = 0.1
DEFAULT_OUTPUT_MIN = 0.0
DEFAULT_OUTPUT_MAX = 1.0
DEFAULT_CLAMP = True
CATEGORY = "Realtime Nodes/MediaPipe Vision/Common/Delta"
@classmethod
def INPUT_TYPES(cls):
# Input types remain the same
output_type_str = "INT" if cls.IS_INT_OUTPUT else "FLOAT"
output_min_key = f"output_min_{output_type_str.lower()}"
output_max_key = f"output_max_{output_type_str.lower()}"
return {
"required": {
"x": ("FLOAT", {"default": 0.0, "forceInput": True}),
"y": ("FLOAT", {"default": 0.0, "forceInput": True}),
"z": ("FLOAT", {"default": 0.0, "forceInput": True}),
"window_size": ("INT", { "default": 1, "min": 1, "max": MAX_POSITION_HISTORY - 1, "tooltip": "Number of frames back for delta"}),
"delta_min": ("FLOAT", {"default": cls.DEFAULT_DELTA_MIN, "step": 0.01}),
"delta_max": ("FLOAT", {"default": cls.DEFAULT_DELTA_MAX, "step": 0.01}),
output_min_key: (output_type_str, {"default": cls.DEFAULT_OUTPUT_MIN}),
output_max_key: (output_type_str, {"default": cls.DEFAULT_OUTPUT_MAX}),
"clamp": ("BOOLEAN", {"default": cls.DEFAULT_CLAMP}),
}
}
def execute(self, **kwargs):
output_type_str = "INT" if self.IS_INT_OUTPUT else "FLOAT"
output_min_key = f"output_min_{output_type_str.lower()}"
output_max_key = f"output_max_{output_type_str.lower()}"
x_in, y_in, z_in = kwargs.get('x'), kwargs.get('y'), kwargs.get('z')
window_size = kwargs.get('window_size', 1)
delta_min = kwargs.get('delta_min')
delta_max = kwargs.get('delta_max')
output_min = kwargs.get(output_min_key)
output_max = kwargs.get(output_max_key)
clamp = kwargs.get('clamp')
fallback_value = int(output_min) if self.IS_INT_OUTPUT else float(output_min)
if None in [x_in, y_in, z_in, delta_min, delta_max, output_min, output_max, clamp]:
logger.error(f"{self.__class__.__name__}: Missing required arguments.")
return (fallback_value,)
# --- Input Type Handling & History Cleanup ---
is_list = isinstance(x_in, list)
if is_list != isinstance(y_in, list) or is_list != isinstance(z_in, list):
raise ValueError(f"{self.__class__.__name__}: Inputs x, y, z must be all floats or all lists.")
if is_list:
current_batch_size = len(x_in)
# Remove histories for indices no longer present
keys_to_remove = [k for k in self._histories if k >= current_batch_size]
if keys_to_remove:
# logger.debug(f"{self.__class__.__name__}: Removing histories for indices {keys_to_remove}")
for k in keys_to_remove:
del self._histories[k]
if current_batch_size == 0:
logger.warning(f"{self.__class__.__name__}: Received empty lists for coordinates.")
return ([],)
# Process List Input
results = []
for i in range(current_batch_size):
xi, yi, zi = x_in[i], y_in[i], z_in[i]
if not all(isinstance(c, (float, int)) for c in [xi, yi, zi]):
logger.warning(f"{self.__class__.__name__}: Invalid coordinate type at index {i}. Using fallback value.")
results.append(fallback_value)
continue
raw_delta = self._get_position_delta(float(xi), float(yi), float(zi), window_size, item_index=i)
scaled_value = scale_value(raw_delta, delta_min, delta_max, output_min, output_max, clamp)
if self.IS_INT_OUTPUT:
output_val = int(round(scaled_value)) if scaled_value is not None else int(output_min)
else:
output_val = scaled_value if scaled_value is not None else float(output_min)
results.append(output_val)
return (results,)
else:
# Single Float Input: Remove all histories except for index 0
keys_to_remove = [k for k in self._histories if k != 0]
if keys_to_remove:
# logger.debug(f"{self.__class__.__name__}: Switching to single input, removing histories for indices {keys_to_remove}")
for k in keys_to_remove:
del self._histories[k]
# Process Single Float Input
if not all(isinstance(c, (float, int)) for c in [x_in, y_in, z_in]):
logger.error(f"{self.__class__.__name__}: Invalid coordinate types for single float input.")
return (fallback_value,)
raw_delta = self._get_position_delta(float(x_in), float(y_in), float(z_in), window_size, item_index=0)
scaled_value = scale_value(raw_delta, delta_min, delta_max, output_min, output_max, clamp)
if self.IS_INT_OUTPUT:
output_val = int(round(scaled_value)) if scaled_value is not None else int(output_min)
else:
output_val = scaled_value if scaled_value is not None else float(output_min)
return (output_val,)
class PositionDeltaIntControlNode(PositionDeltaControlNode):
RETURN_TYPES = ("INT",) # Semantic type is INT
RETURN_NAMES = ("value",)
IS_INT_OUTPUT = True
DEFAULT_OUTPUT_MIN = 0
DEFAULT_OUTPUT_MAX = 100
class PositionDeltaFloatControlNode(PositionDeltaControlNode):
RETURN_TYPES = ("FLOAT",) # Semantic type is FLOAT
RETURN_NAMES = ("value",)
IS_INT_OUTPUT = False
DEFAULT_OUTPUT_MIN = 0.0
DEFAULT_OUTPUT_MAX = 1.0
# --- Position Delta Trigger Node (Handling Lists) ---
class PositionDeltaTriggerNode(BasePositionDeltaNode):
"""Outputs a BOOLEAN trigger if position delta crosses a threshold.
Handles single floats or lists of coordinates.
Outputs BOOLEAN type (actual value is boolean or list[boolean]).
"""
RETURN_TYPES = ("BOOLEAN",) # Semantic type
RETURN_NAMES = ("trigger",) # Singular name
DEFAULT_THRESHOLD = 0.05
CATEGORY = "Realtime Nodes/MediaPipe Vision/Common/Delta"
@classmethod
def INPUT_TYPES(cls):
# Inputs remain the same
return {
"required": {
"x": ("FLOAT", {"default": 0.0, "forceInput": True}),
"y": ("FLOAT", {"default": 0.0, "forceInput": True}),
"z": ("FLOAT", {"default": 0.0, "forceInput": True}),
"window_size": ("INT", {"default": 1, "min": 1, "max": MAX_POSITION_HISTORY - 1}),
"threshold": ("FLOAT", {"default": cls.DEFAULT_THRESHOLD, "step": 0.01}),
"condition": (["Above", "Below", "Equals", "Not Equals"],),
}
}
def execute(self, **kwargs):
x_in, y_in, z_in = kwargs.get('x'), kwargs.get('y'), kwargs.get('z')
window_size = kwargs.get('window_size', 1)
threshold = kwargs.get('threshold')
condition = kwargs.get('condition')
fallback_value = False
if None in [x_in, y_in, z_in, threshold, condition]:
logger.error(f"{self.__class__.__name__}: Missing required arguments.")
return (fallback_value,)
# --- Input Type Handling & History Cleanup ---
is_list = isinstance(x_in, list)
if is_list != isinstance(y_in, list) or is_list != isinstance(z_in, list):
raise ValueError(f"{self.__class__.__name__}: Inputs x, y, z must be all floats or all lists.")
if is_list:
current_batch_size = len(x_in)
# Remove histories for indices no longer present
keys_to_remove = [k for k in self._histories if k >= current_batch_size]
if keys_to_remove:
# logger.debug(f"{self.__class__.__name__}: Removing histories for indices {keys_to_remove}")
for k in keys_to_remove:
del self._histories[k]
if current_batch_size == 0:
logger.warning(f"{self.__class__.__name__}: Received empty lists for coordinates.")
return ([],)
# Process List Input
results = []
for i in range(current_batch_size):
xi, yi, zi = x_in[i], y_in[i], z_in[i]
if not all(isinstance(c, (float, int)) for c in [xi, yi, zi]):
logger.warning(f"{self.__class__.__name__}: Invalid coordinate type at index {i}. Using fallback value.")
results.append(fallback_value)
continue
raw_delta = self._get_position_delta(float(xi), float(yi), float(zi), window_size, item_index=i)
triggered = fallback_value
if raw_delta is not None:
if condition == "Above": triggered = raw_delta > threshold
elif condition == "Below": triggered = raw_delta < threshold
elif condition == "Equals": triggered = abs(raw_delta - threshold) < FLOAT_EQUALITY_TOLERANCE
elif condition == "Not Equals": triggered = abs(raw_delta - threshold) >= FLOAT_EQUALITY_TOLERANCE
results.append(triggered)
return (results,)
else:
# Single Float Input: Remove all histories except for index 0
keys_to_remove = [k for k in self._histories if k != 0]
if keys_to_remove:
# logger.debug(f"{self.__class__.__name__}: Switching to single input, removing histories for indices {keys_to_remove}")
for k in keys_to_remove:
del self._histories[k]
# Process Single Float Input
if not all(isinstance(c, (float, int)) for c in [x_in, y_in, z_in]):
logger.error(f"{self.__class__.__name__}: Invalid coordinate types for single float input.")
return (fallback_value,)
raw_delta = self._get_position_delta(float(x_in), float(y_in), float(z_in), window_size, item_index=0)
triggered = fallback_value
if raw_delta is not None:
if condition == "Above": triggered = raw_delta > threshold
elif condition == "Below": triggered = raw_delta < threshold
elif condition == "Equals": triggered = abs(raw_delta - threshold) < FLOAT_EQUALITY_TOLERANCE
elif condition == "Not Equals": triggered = abs(raw_delta - threshold) >= FLOAT_EQUALITY_TOLERANCE
return (triggered,)
# --- Landmark Proximity Node (Handles Float or List Inputs) ---
class LandmarkProximityNode:
"""
Calculates the Euclidean distance between landmark positions.
Handles single coordinate sets or lists thereof.
Outputs FLOAT type (actual value is float or list[float]).
"""
CATEGORY = "Realtime Nodes/MediaPipe Vision/Common/Utils"
FUNCTION = "calculate_distances"
# Output semantic type is FLOAT
RETURN_TYPES = ("FLOAT",)
RETURN_NAMES = ("distance",) # Singular name
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
# Keep FLOAT type
"x1": ("FLOAT", {"default": 0.0, "forceInput": True}),
"y1": ("FLOAT", {"default": 0.0, "forceInput": True}),
"z1": ("FLOAT", {"default": 0.0, "forceInput": True}),
"x2": ("FLOAT", {"default": 0.0, "forceInput": True}),
"y2": ("FLOAT", {"default": 0.0, "forceInput": True}),
"z2": ("FLOAT", {"default": 0.0, "forceInput": True}),
}
}
def calculate_distances(self, x1: Union[float, List[float]], y1: Union[float, List[float]], z1: Union[float, List[float]],
x2: Union[float, List[float]], y2: Union[float, List[float]], z2: Union[float, List[float]]) -> Union[float, List[float]]: # Return type hint updated
# --- Input Type Handling ---
inputs = [x1, y1, z1, x2, y2, z2]
is_list_input = isinstance(x1, list)
fallback_value = 0.0
if any(isinstance(inp, list) != is_list_input for inp in inputs):
raise ValueError("All coordinate inputs (x1..z2) must be of the same type (all floats or all lists).")
if is_list_input:
# Process List Input
list_len = len(x1)
if any(len(inp) != list_len for inp in inputs if isinstance(inp, list)):
raise ValueError("If inputs are lists, all coordinate lists (x1..z2) must have the same length.")
if list_len == 0:
logger.warning("Input coordinate lists are empty. Returning empty distance list.")
return ([],) # Return empty list for FLOAT type
num_calcs = list_len
x1_list, y1_list, z1_list = x1, y1, z1
x2_list, y2_list, z2_list = x2, y2, z2
distances = []
for i in range(num_calcs):
try:
_x1, _y1, _z1 = x1_list[i], y1_list[i], z1_list[i]
_x2, _y2, _z2 = x2_list[i], y2_list[i], z2_list[i]
if not all(isinstance(c, (float, int)) for c in [_x1, _y1, _z1, _x2, _y2, _z2]):
logger.warning(f"Invalid coordinate type found at index {i}. Using fallback distance {fallback_value}.")
distances.append(fallback_value)
continue
dist = math.sqrt((_x2 - _x1)**2 + (_y2 - _y1)**2 + (_z2 - _z1)**2)
distances.append(dist)
except Exception as e:
logger.error(f"Error calculating distance at index {i}: {e}", exc_info=True)
distances.append(fallback_value)
return (distances,) # Return list of results
else:
# Process Single Float Input
_x1, _y1, _z1 = x1, y1, z1
_x2, _y2, _z2 = x2, y2, z2
if not all(isinstance(c, (float, int)) for c in [_x1, _y1, _z1, _x2, _y2, _z2]):
logger.error(f"Invalid coordinate types for single float input.")
return (fallback_value,) # Return single fallback
try:
dist = math.sqrt((_x2 - _x1)**2 + (_y2 - _y1)**2 + (_z2 - _z1)**2)
return (dist,) # Return single float value
except Exception as e:
logger.error(f"Error calculating distance for single input: {e}", exc_info=True)
return (fallback_value,) # Return single fallback
# --- Universal Landmark Mask Node (Handles Float or List Inputs) ---
class UniversalLandmarkMaskNode:
"""
Creates geometric masks (circle or square) centered at landmark coordinates.
Handles single (x, y) floats or lists (x_list, y_list) as input.
Outputs a batch of masks (N, 1, H, W).
"""
CATEGORY = "Realtime Nodes/MediaPipe Vision/Common/Masking"
FUNCTION = "create_masks" # Keep plural name for clarity
RETURN_TYPES = ("MASK",)
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"x": ("FLOAT", {"default": 0.5, "forceInput": True, "tooltip": "Normalized X coordinate (0.0-1.0) or List thereof"}),
"y": ("FLOAT", {"default": 0.5, "forceInput": True, "tooltip": "Normalized Y coordinate (0.0-1.0) or List thereof"}),
"image_for_dimensions": ("IMAGE", {"tooltip": "Reference image to determine the output mask size"}),
"radius": ("INT", {"default": 25, "min": 1, "max": 1024, "tooltip": "Radius of the mask shape in pixels"}),
"shape": (["Circle", "Square"], {"default": "Circle"}),
"invert": ("BOOLEAN", {"default": False}),
}
}
def create_masks(self, x: Union[float, List[float]], y: Union[float, List[float]], image_for_dimensions: torch.Tensor, radius: int, shape: str, invert: bool) -> Tuple[torch.Tensor]:
# --- Input Type Handling ---
is_x_list = isinstance(x, list)
is_y_list = isinstance(y, list)
if is_x_list != is_y_list:
raise ValueError("Inputs x and y must be both floats or both lists.")
if is_x_list:
# Input is List[float]
x_list = x
y_list = y
if len(x_list) != len(y_list):
raise ValueError(f"x_list (len {len(x_list)}) and y_list (len {len(y_list)}) must have the same length.")
if not x_list: # Handle empty list
logger.warning("Input coordinate lists are empty. Returning empty mask batch.")
# Determine shape for empty tensor
if image_for_dimensions.dim() != 4:
logger.error("Input image must be BHWC format to determine empty mask size.")
return (torch.zeros((0, 1, 64, 64), dtype=torch.float32),)
_, height, width, _ = image_for_dimensions.shape
return (torch.zeros((0, 1, height, width), dtype=torch.float32, device=image_for_dimensions.device),)
num_coords = len(x_list)
else:
# Input is float - wrap in lists for unified processing
x_list = [x]
y_list = [y]
num_coords = 1
# --- Get Dimensions ---
if image_for_dimensions.dim() != 4:
logger.error("Input image must be BHWC format.")
# Return empty batch matching input size
return (torch.zeros((num_coords, 1, 64, 64), dtype=torch.float32),)
_, height, width, _ = image_for_dimensions.shape
device = image_for_dimensions.device
# --- Generate Masks ---
output_masks = []
for i in range(num_coords):
xi = x_list[i]
yi = y_list[i]
# Validate individual coordinates
if not isinstance(xi, (float, int)) or not isinstance(yi, (float, int)):
logger.warning(f"Invalid coordinate type at index {i} (x={type(xi)}, y={type(yi)}). Skipping mask generation for this index.")
mask_tensor = torch.zeros((1, height, width), dtype=torch.float32, device=device)
output_masks.append(mask_tensor)
continue
# Calculate pixel coordinates
px = max(0, min(width - 1, int(xi * width)))
py = max(0, min(height - 1, int(yi * height)))
mask_np = np.zeros((height, width), dtype=np.uint8)
try:
center = (px, py)
if shape == "Circle":
cv2.circle(mask_np, center, radius, 1, -1)
elif shape == "Square":
half_size = radius
top_left = (max(0, px - half_size), max(0, py - half_size))
bottom_right = (min(width - 1, px + half_size), min(height - 1, py + half_size))
cv2.rectangle(mask_np, top_left, bottom_right, 1, -1)
else:
cv2.circle(mask_np, center, radius, 1, -1) # Default circle
if invert:
mask_np = 1 - mask_np
except Exception as e:
logger.error(f"Error drawing mask shape at index {i}: {e}", exc_info=True)
# mask_np remains zeros
# Convert to tensor (1, H, W)
mask_tensor = torch.from_numpy(mask_np).float().unsqueeze(0).to(device)
output_masks.append(mask_tensor)
# Stack into batch (N, 1, H, W)
output_batch = torch.stack(output_masks, dim=0)
return (output_batch,)
# --- Node Mappings (Revert Delta node display names) ---
NODE_CLASS_MAPPINGS = {
"LandmarkProximity": LandmarkProximityNode,
"UniversalLandmarkMask": UniversalLandmarkMaskNode,
"PositionDeltaIntControl": PositionDeltaIntControlNode,
"PositionDeltaFloatControl": PositionDeltaFloatControlNode,
"PositionDeltaTrigger": PositionDeltaTriggerNode,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"LandmarkProximity": "Landmark Proximity Calculator",
"UniversalLandmarkMask": "Universal Landmark Mask Generator (Batch)",
"PositionDeltaIntControl": "Position Delta Int Control",
"PositionDeltaFloatControl": "Position Delta Float Control",
"PositionDeltaTrigger": "Position Delta Trigger",
}
@@ -0,0 +1 @@
@@ -0,0 +1,75 @@
import torch
from .....src.mediapipe_vision.face_detection.detector import FaceDetector
from ...common.model_loader import MediaPipeModelLoaderBaseNode
from ...common.base_detector_node import BaseMediaPipeDetectorNode
import logging
logger = logging.getLogger(__name__)
_category = "Realtime Nodes/MediaPipe Vision/Face/FaceDetection"
# Inherit from Base Loader
class MediaPipeFaceDetectorModelLoaderNode(MediaPipeModelLoaderBaseNode):
"""ComfyUI node for loading MediaPipe Face Detector models."""
TASK_TYPE = "face_detector"
RETURN_TYPES = ("FACE_DETECTOR_MODEL_INFO",) # Specific return type
RETURN_NAMES = ("model_info",)
CATEGORY = _category
DESCRIPTION = "Loads a MediaPipe Face Detector model for finding faces in images. Required before using the Face Detector node."
# INPUT_TYPES and FUNCTION inherited
# Detection/Landmarking Nodes
class MediaPipeFaceDetectorNode(BaseMediaPipeDetectorNode):
"""ComfyUI node for MediaPipe Face Detection."""
# Define class variables required by the base class
DETECTOR_CLASS = FaceDetector
MODEL_INFO_TYPE = "FACE_DETECTOR_MODEL_INFO"
EXPECTED_TASK_TYPE = "face_detector"
RETURN_TYPES = ("FACE_DETECTIONS",)
RETURN_NAMES = ("face_detections",)
FUNCTION = "detect"
CATEGORY = _category
DESCRIPTION = "Detects faces in images using MediaPipe's face detection technology. Provides coordinates and confidence scores for all faces found in the image."
@classmethod
def INPUT_TYPES(cls):
# Start with the base inputs from the parent class
inputs = super().INPUT_TYPES()
# Add face detection specific parameters
inputs["required"].update({
"min_confidence": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.05,
"tooltip": "Minimum confidence threshold for face detection - lower values find more faces but may include false positives"}),
})
return inputs
def detect(self, image: torch.Tensor, model_info: dict, min_confidence: float,
running_mode: str, delegate: str):
"""Performs face detection with the configured parameters."""
# Validate model_info and get model path
model_path = self.validate_model_info(model_info)
# Initialize or update detector
detector = self.initialize_or_update_detector(model_path)
# Perform detection with all parameters
batch_results = detector.detect(
image,
min_detection_confidence=min_confidence,
running_mode=running_mode,
delegate=delegate
)
return (batch_results,)
# Define mappings for ComfyUI
NODE_CLASS_MAPPINGS = {
"MediaPipeFaceDetectorModelLoader": MediaPipeFaceDetectorModelLoaderNode,
"MediaPipeFaceDetector": MediaPipeFaceDetectorNode,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"MediaPipeFaceDetectorModelLoader": "Load Face Detector Model (MediaPipe)",
"MediaPipeFaceDetector": "Face Detector (MediaPipe)",
}
@@ -0,0 +1,155 @@
# face_detection_mask.py
import torch
import numpy as np
import logging
import scipy.ndimage
logger = logging.getLogger(__name__)
_category = "Realtime Nodes/MediaPipe Vision/Face/FaceDetection"
class FaceDetectionToMask:
"""Converts face detection bounding boxes to masks that can be used in ComfyUI."""
CATEGORY = _category
DESCRIPTION = "Creates masks from face detections that can be used for selective processing, inpainting, or compositing. Generate masks based on face bounding boxes or facial keypoints."
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"face_detections": ("FACE_DETECTIONS", {"tooltip": "Face detection results from the Face Detector node"}),
"width": ("INT", {"default": 512, "min": 1, "max": 8192, "step": 1,
"tooltip": "Width of the output mask in pixels"}),
"height": ("INT", {"default": 512, "min": 1, "max": 8192, "step": 1,
"tooltip": "Height of the output mask in pixels"}),
"mask_type": (["Bounding Box", "Keypoints", "Both"], {"default": "Bounding Box",
"tooltip": "Type of mask to create - Bounding Box (rectangle around face), Keypoints (dots at facial points), or Both"}),
"face_selection": (["All Faces", "First Face", "Largest Face", "Highest Confidence Face"],
{"default": "All Faces",
"tooltip": "Which faces to include in the mask - all detected faces, only the first face, largest face, or most confident detection"}),
"feather_amount": ("INT", {"default": 0, "min": 0, "max": 100, "step": 1,
"tooltip": "How much to soften the edges of the mask in pixels - higher values create smoother transitions"}),
"keypoint_radius": ("INT", {"default": 5, "min": 1, "max": 50, "step": 1,
"tooltip": "Size of each keypoint in the mask when using Keypoints or Both mask types"}),
"invert": ("BOOLEAN", {"default": False,
"tooltip": "When enabled, inverts the mask (white becomes black, black becomes white)"}),
"score_threshold": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01,
"tooltip": "Minimum confidence score for including face detections - higher values filter out less confident detections"})
},
"optional": {
"reference_image": ("IMAGE", {"tooltip": "Optional image to match dimensions - mask will be created with the same size as this image"})
}
}
RETURN_TYPES = ("MASK",)
FUNCTION = "create_mask"
def create_mask(self, face_detections, width, height, mask_type, face_selection,
feather_amount, keypoint_radius, invert, score_threshold, reference_image=None):
# If reference image provided, use its dimensions
if reference_image is not None:
height, width = reference_image.shape[1], reference_image.shape[2]
batch_size = len(face_detections)
masks = []
for i in range(batch_size):
# Create empty mask
mask = torch.zeros((height, width), dtype=torch.float32, device="cpu")
detections_for_image = face_detections[i]
if not detections_for_image:
masks.append(mask)
continue
# Filter detections by score threshold
filtered_detections = []
for detection in detections_for_image:
if detection.score is None or detection.score >= score_threshold:
filtered_detections.append(detection)
# Select which faces to include based on face_selection parameter
selected_detections = []
if face_selection == "All Faces":
selected_detections = filtered_detections
elif face_selection == "First Face" and filtered_detections:
selected_detections = [filtered_detections[0]]
elif face_selection == "Largest Face" and filtered_detections:
# Find the detection with the largest area
largest_detection = max(filtered_detections,
key=lambda d: d.bounding_box.width * d.bounding_box.height)
selected_detections = [largest_detection]
elif face_selection == "Highest Confidence Face" and filtered_detections:
# Find the detection with the highest confidence score
# Filter out detections with no score first
scored_detections = [d for d in filtered_detections if d.score is not None]
if scored_detections:
highest_conf_detection = max(scored_detections, key=lambda d: d.score)
selected_detections = [highest_conf_detection]
else:
# If no detections have scores, fall back to first detection
selected_detections = [filtered_detections[0]]
for detection in selected_detections:
# Get bounding box
bbox = detection.bounding_box
x1, y1 = max(0, bbox.origin_x), max(0, bbox.origin_y)
x2, y2 = min(width-1, bbox.origin_x + bbox.width), min(height-1, bbox.origin_y + bbox.height)
if x2 <= x1 or y2 <= y1:
continue # Skip invalid boxes
# Fill bounding box region if requested
if mask_type in ["Bounding Box", "Both"]:
mask[y1:y2+1, x1:x2+1] = 1.0
# Add keypoints if requested
if mask_type in ["Keypoints", "Both"] and detection.keypoints:
for kp in detection.keypoints:
# Keypoints are in normalized coordinates, convert to pixel coordinates
cx, cy = int(kp.x * width), int(kp.y * height)
# Draw circle for each keypoint
for y in range(max(0, cy - keypoint_radius), min(height, cy + keypoint_radius + 1)):
for x in range(max(0, cx - keypoint_radius), min(width, cx + keypoint_radius + 1)):
# Check if point is within radius
if (x - cx) ** 2 + (y - cy) ** 2 <= keypoint_radius ** 2:
mask[y, x] = 1.0
# Feather the edges if requested
if feather_amount > 0:
# Convert mask to numpy for scipy operations
mask_np = mask.numpy()
# Calculate distance transform
if invert:
distance = scipy.ndimage.distance_transform_edt(mask_np == 0)
else:
distance = scipy.ndimage.distance_transform_edt(mask_np > 0)
# Create feathered edge
feather_mask = np.clip(distance / feather_amount, 0, 1)
if invert:
feather_mask = 1.0 - feather_mask
# Convert back to tensor
mask = torch.from_numpy(feather_mask).float()
elif invert:
# Simple inversion without feathering
mask = 1.0 - mask
masks.append(mask)
# Stack all masks into a batch
mask_batch = torch.stack(masks)
return (mask_batch,)
# Node registration
NODE_CLASS_MAPPINGS = {
"FaceDetectionToMask": FaceDetectionToMask
}
NODE_DISPLAY_NAME_MAPPINGS = {
"FaceDetectionToMask": "Face Detection to Mask (MediaPipe)"
}
@@ -0,0 +1,107 @@
# visualize_face_detection.py
import torch
import numpy as np
import cv2
from ...common.base_visualization_nodes import BaseDetectionVisualizationNode
import logging
logger = logging.getLogger(__name__)
class VisualizeFaceDetections(BaseDetectionVisualizationNode):
"""Draws face detection bounding boxes and keypoints onto images."""
CATEGORY = "Realtime Nodes/MediaPipe Vision/Face/FaceDetection/Visualization"
DESCRIPTION = "Visualizes face detections by drawing boxes and keypoints onto images. Allows customization of colors and sizes to highlight detected faces."
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"face_detections": ("FACE_DETECTIONS", {"tooltip": "Face detection results from the Face Detector node"}),
"image": ("IMAGE", {"tooltip": "The image on which to draw face detection visualizations"}),
},
"optional": {
"box_color": ("STRING", {"default": "#00FF00", "tooltip": "Hex color code for the face bounding boxes (default: green)"}),
"box_thickness": ("INT", {"default": 2, "min": 1, "max": 10, "tooltip": "Thickness of the face bounding box lines"}),
"keypoint_color": ("STRING", {"default": "#FF0000", "tooltip": "Hex color code for facial keypoints (default: red)"}),
"keypoint_radius": ("INT", {"default": 2, "min": 1, "max": 10, "tooltip": "Size of the dots representing facial keypoints"}),
"label_color": ("STRING", {"default": "#FFFFFF", "tooltip": "Hex color code for the text labels (default: white)"}),
"label_bg_color": ("STRING", {"default": "#00FF00", "tooltip": "Hex color code for the label background (default: green)"}),
}
}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "draw_visualizations"
def draw_visualizations(self, face_detections, image,
box_color="#00FF00", box_thickness=2,
keypoint_color="#FF0000", keypoint_radius=2,
label_color="#FFFFFF", label_bg_color="#00FF00"):
"""Main function that processes the batch of images and detections."""
return self.process_batch(
face_detections,
image,
self.draw_faces,
box_color=box_color,
box_thickness=box_thickness,
keypoint_color=keypoint_color,
keypoint_radius=keypoint_radius,
label_color=label_color,
label_bg_color=label_bg_color
)
def draw_faces(self, image, width, height, detections,
box_color="#00FF00", box_thickness=2,
keypoint_color="#FF0000", keypoint_radius=2,
label_color="#FFFFFF", label_bg_color="#00FF00"):
"""Callback function to draw faces on a single image."""
# Convert hex colors to BGR tuples for cv2
box_color_bgr = self._hex_to_bgr(box_color)
keypoint_color_bgr = self._hex_to_bgr(keypoint_color)
for detection in detections:
bbox = detection.bounding_box
x1, y1 = bbox.origin_x, bbox.origin_y
x2, y2 = bbox.origin_x + bbox.width, bbox.origin_y + bbox.height
# Draw bounding box
cv2.rectangle(image, (int(x1), int(y1)), (int(x2), int(y2)), box_color_bgr, box_thickness)
# Draw Keypoints
if detection.keypoints:
for kp in detection.keypoints:
cx, cy = int(kp.x * width), int(kp.y * height)
cv2.circle(image, (cx, cy), keypoint_radius, keypoint_color_bgr, -1)
# Prepare label text (confidence score)
label_text = f"Face: {detection.score:.2f}" if detection.score is not None else "Face"
# Draw label with background
try:
self.draw_text_with_background(
image,
label_text,
(int(x1), int(y1)),
text_color=label_color,
bg_color=label_bg_color
)
except Exception as e:
logger.error(f"Error drawing text label: {e}")
# Fallback using basic cv2 text
cv2.putText(
image,
label_text,
(int(x1) + 2, int(y1) + 15),
cv2.FONT_HERSHEY_SIMPLEX,
0.5,
self._hex_to_bgr(label_color),
1
)
# Node registration
NODE_CLASS_MAPPINGS = {
"VisualizeFaceDetections": VisualizeFaceDetections
}
NODE_DISPLAY_NAME_MAPPINGS = {
"VisualizeFaceDetections": "Visualize Face Detections (MediaPipe)"
}
@@ -0,0 +1,179 @@
# node_wrappers/face/blendshape.py
import logging
from typing import Optional, List, Any, Dict, Tuple
from .....src.mediapipe_vision.types import BLENDSHAPES_LIST # Updated type alias
from .....src.mediapipe_vision.utils.misc_utils import get_blendshape_categories # Utility to get all category names
# Import base trigger/control logic if we adapt base classes later, or implement directly for now
from ...common.universal_landmark_nodes import FLOAT_EQUALITY_TOLERANCE # For trigger node comparisons
from .....src.mediapipe_vision.utils.delta_utils import scale_value
logger = logging.getLogger(__name__)
# --- Constants ---
# Get the list of available blendshape categories from MediaPipe
ALL_BLENDSHAPE_CATEGORIES = get_blendshape_categories()
if not ALL_BLENDSHAPE_CATEGORIES:
logger.warning("Could not retrieve blendshape categories. Blendshape nodes might not function correctly.")
ALL_BLENDSHAPE_CATEGORIES = ["_neutral"] # Provide a fallback
_category = "Realtime Nodes/MediaPipe Vision/Face/Blendshape"
# --- Helper Functions ---
def get_blendshape_score(blendshapes_result: BLENDSHAPES_LIST, category_name: str) -> Optional[float]:
"""Extracts the score for a specific blendshape category from the result."""
if not blendshapes_result or not blendshapes_result[0] or not blendshapes_result[0][0]:
#logger.debug(f"No blendshapes found in result.")
return None # No blendshape data
# BLENDSHAPES structure: List[List[Dict[str, float]]]
# We typically care about the first face in the first batch item
first_face_blendshapes = blendshapes_result[0][0]
# In the MediaPipe output, blendshapes are represented as a dictionary
# where keys are category names and values are scores
if isinstance(first_face_blendshapes, dict):
# Dictionary format: {category_name: score, ...}
if category_name in first_face_blendshapes:
return first_face_blendshapes[category_name]
elif isinstance(first_face_blendshapes, list):
# Try to handle the list format with objects that have category_name attribute
try:
target_blendshape = next((bs for bs in first_face_blendshapes if hasattr(bs, 'category_name') and bs.category_name == category_name), None)
if target_blendshape:
return target_blendshape.score
except Exception as e:
logger.error(f"Error processing blendshape list: {e}")
logger.warning(f"Blendshape category '{category_name}' not found in results. Type: {type(first_face_blendshapes)}")
return None
# --- Control Nodes ---
class BlendshapeControlFloatNode:
"""Outputs a scaled FLOAT based on the score of a selected blendshape."""
CATEGORY = _category
RETURN_TYPES = ("FLOAT",)
FUNCTION = "execute"
DESCRIPTION = "Converts facial expressions into floating point values that can control other nodes. For example, use a smile to control image parameters or animation intensity."
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"blendshapes": ("BLENDSHAPES_LIST", {"tooltip": "Facial expression data from the Face Landmarker node"}),
"blendshape_name": (ALL_BLENDSHAPE_CATEGORIES, {"default": "mouthSmileLeft",
"tooltip": "The specific facial expression to monitor (e.g., smile, frown, raised eyebrows)"}),
# Blendshape scores are typically 0.0 to 1.0, but allow flexibility
"score_min": ("FLOAT", {"default": 0.0, "step": 0.01,
"tooltip": "The minimum facial expression intensity to consider (usually 0.0)"}),
"score_max": ("FLOAT", {"default": 1.0, "step": 0.01,
"tooltip": "The maximum facial expression intensity to consider (usually 1.0)"}),
"output_min_float": ("FLOAT", {"default": 0.0, "step": 0.01,
"tooltip": "The minimum output value when expression is at or below score_min"}),
"output_max_float": ("FLOAT", {"default": 1.0, "step": 0.01,
"tooltip": "The maximum output value when expression is at or above score_max"}),
"clamp": ("BOOLEAN", {"default": True,
"tooltip": "When enabled, restricts output values to stay within min and max range"}),
}
}
def execute(self, blendshapes, blendshape_name, score_min, score_max, output_min_float, output_max_float, clamp):
raw_score = get_blendshape_score(blendshapes, blendshape_name)
scaled_value = scale_value(raw_score, score_min, score_max, output_min_float, output_max_float, clamp)
# Default to min value on error or if blendshape not found
output_val = scaled_value if scaled_value is not None else float(output_min_float)
return (output_val,)
class BlendshapeControlIntNode:
"""Outputs a scaled INT based on the score of a selected blendshape."""
CATEGORY = _category
RETURN_TYPES = ("INT",)
FUNCTION = "execute"
DESCRIPTION = "Converts facial expressions into integer values that can control other nodes. Useful for selecting steps, seeds, or indexes based on facial expressions."
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"blendshapes": ("BLENDSHAPES_LIST", {"tooltip": "Facial expression data from the Face Landmarker node"}),
"blendshape_name": (ALL_BLENDSHAPE_CATEGORIES, {"default": "mouthSmileLeft",
"tooltip": "The specific facial expression to monitor (e.g., smile, frown, raised eyebrows)"}),
"score_min": ("FLOAT", {"default": 0.0, "step": 0.01,
"tooltip": "The minimum facial expression intensity to consider (usually 0.0)"}),
"score_max": ("FLOAT", {"default": 1.0, "step": 0.01,
"tooltip": "The maximum facial expression intensity to consider (usually 1.0)"}),
"output_min_int": ("INT", {"default": 0,
"tooltip": "The minimum integer output value when expression is at or below score_min"}),
"output_max_int": ("INT", {"default": 100,
"tooltip": "The maximum integer output value when expression is at or above score_max"}),
"clamp": ("BOOLEAN", {"default": True,
"tooltip": "When enabled, restricts output values to stay within min and max range"}),
}
}
def execute(self, blendshapes, blendshape_name, score_min, score_max, output_min_int, output_max_int, clamp):
raw_score = get_blendshape_score(blendshapes, blendshape_name)
scaled_value = scale_value(raw_score, score_min, score_max, output_min_int, output_max_int, clamp)
# Default to min value on error or if blendshape not found
output_val = int(round(scaled_value)) if scaled_value is not None else int(output_min_int)
return (output_val,)
# --- Trigger Node ---
class BlendshapeTriggerNode:
"""Outputs a BOOLEAN trigger if a blendshape score crosses a threshold."""
CATEGORY = _category
RETURN_TYPES = ("BOOLEAN",)
FUNCTION = "execute"
DESCRIPTION = "Creates a true/false trigger based on facial expressions exceeding thresholds. Perfect for activating workflows when someone smiles, frowns, or makes specific expressions."
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"blendshapes": ("BLENDSHAPES_LIST", {"tooltip": "Facial expression data from the Face Landmarker node"}),
"blendshape_name": (ALL_BLENDSHAPE_CATEGORIES, {"default": "jawOpen",
"tooltip": "The specific facial expression to monitor (e.g., smile, frown, open mouth)"}),
"threshold": ("FLOAT", {"default": 0.5, "step": 0.01,
"tooltip": "The expression intensity threshold value that triggers the condition"}),
"condition": (["Above", "Below", "Equals", "Not Equals"],
{"tooltip": "The condition that determines when to trigger - Above (expression > threshold), Below (expression < threshold), etc."}),
}
}
def execute(self, blendshapes, blendshape_name, threshold, condition):
raw_score = get_blendshape_score(blendshapes, blendshape_name)
triggered = False
if raw_score is not None:
if condition == "Above" and raw_score > threshold:
triggered = True
elif condition == "Below" and raw_score < threshold:
triggered = True
elif condition == "Equals" and abs(raw_score - threshold) < FLOAT_EQUALITY_TOLERANCE:
triggered = True
elif condition == "Not Equals" and abs(raw_score - threshold) >= FLOAT_EQUALITY_TOLERANCE:
triggered = True
return (triggered,)
# --- Mappings ---
NODE_CLASS_MAPPINGS = {
"BlendshapeControlFloat": BlendshapeControlFloatNode,
"BlendshapeControlInt": BlendshapeControlIntNode,
"BlendshapeTrigger": BlendshapeTriggerNode,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"BlendshapeControlFloat": "Blendshape Control (Float)",
"BlendshapeControlInt": "Blendshape Control (Int)",
"BlendshapeTrigger": "Blendshape Trigger",
}
@@ -0,0 +1,101 @@
"""Node wrapper for face detection."""
import torch
from .....src.mediapipe_vision.face_landmark.detector import FaceLandmarkDetector
# Import Base Classes
from ...common.model_loader import MediaPipeModelLoaderBaseNode
from ...common.base_detector_node import BaseMediaPipeDetectorNode
import logging
logger = logging.getLogger(__name__)
_category = "Realtime Nodes/MediaPipe Vision/Face/FaceLandmark"
# Inherit from Base Loader
class MediaPipeFaceLandmarkerModelLoaderNode(MediaPipeModelLoaderBaseNode):
"""ComfyUI node for loading MediaPipe Face Landmarker models."""
TASK_TYPE = "face_landmarker"
RETURN_TYPES = ("FACE_LANDMARKER_MODEL_INFO",) # Specific return type
RETURN_NAMES = ("model_info",)
CATEGORY = _category
DESCRIPTION = "Loads a MediaPipe Face Landmarker model which detects facial landmarks, expressions and face orientation. Required before using the Face Landmarker node."
# INPUT_TYPES and FUNCTION inherited
class MediaPipeFaceLandmarkerNode(BaseMediaPipeDetectorNode):
"""ComfyUI node for MediaPipe Face Landmark Detection."""
# Define class variables required by the base class
DETECTOR_CLASS = FaceLandmarkDetector
MODEL_INFO_TYPE = "FACE_LANDMARKER_MODEL_INFO"
EXPECTED_TASK_TYPE = "face_landmarker"
RETURN_TYPES = ("FACE_LANDMARKS", "BLENDSHAPES_LIST", "TRANSFORM_MATRIX_LIST")
RETURN_NAMES = ("landmarks", "blendshapes", "transform_matrices")
FUNCTION = "detect"
CATEGORY = _category
DESCRIPTION = "Detects precise facial landmarks (478 points), expressions, and head orientation in images. Provides detailed data for face tracking, animation, or specialized face image processing."
@classmethod
def INPUT_TYPES(cls):
# Start with the base inputs from the parent class
inputs = super().INPUT_TYPES()
# Add face-specific parameters
inputs["required"].update({
"max_results": ("INT", {"default": 1, "min": 1, "max": 10, "step": 1,
"tooltip": "Maximum number of faces to analyze - higher values detect landmarks on more faces but use more processing power"}),
"min_confidence": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01,
"tooltip": "Minimum confidence threshold for detecting faces - lower finds more faces but may include false positives"}),
"min_presence_confidence": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01,
"tooltip": "(VIDEO mode) Minimum confidence that a face is present - lower values may detect faces that are less clear"}),
"min_tracking_confidence": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01,
"tooltip": "(VIDEO mode) Minimum confidence for tracking a face between frames - lower values maintain tracking longer but may drift"}),
"output_blendshapes": ("BOOLEAN", {"default": False,
"tooltip": "When enabled, outputs facial expression data (smile, frown, etc.) useful for animation or emotion analysis"}),
"output_transform_matrix": ("BOOLEAN", {"default": False,
"tooltip": "When enabled, outputs 3D face orientation data (head pose) useful for AR/VR applications"}),
})
return inputs
def detect(self, image: torch.Tensor, model_info: dict, max_results: int,
min_confidence: float, min_presence_confidence: float,
min_tracking_confidence: float, output_blendshapes: bool,
output_transform_matrix: bool, running_mode: str, delegate: str):
"""Performs face landmark detection with the configured parameters."""
# Validate model_info and get model path
model_path = self.validate_model_info(model_info)
# Initialize or update detector
detector = self.initialize_or_update_detector(model_path)
# Perform detection with all parameters
face_landmarks_batch, blendshapes_batch, matrices_batch = detector.detect(
image,
num_faces=max_results,
min_detection_confidence=min_confidence,
min_presence_confidence=min_presence_confidence,
min_tracking_confidence=min_tracking_confidence,
output_blendshapes=output_blendshapes,
output_transform_matrix=output_transform_matrix,
running_mode=running_mode,
delegate=delegate
)
# Return the results with optional empty lists for disabled outputs
return (
face_landmarks_batch,
blendshapes_batch if output_blendshapes else [],
matrices_batch if output_transform_matrix else []
)
# Define mappings for ComfyUI
NODE_CLASS_MAPPINGS = {
"MediaPipeFaceLandmarkerModelLoader": MediaPipeFaceLandmarkerModelLoaderNode,
"MediaPipeFaceLandmarker": MediaPipeFaceLandmarkerNode,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"MediaPipeFaceLandmarkerModelLoader": "Load Face Landmarker Model (MediaPipe)",
"MediaPipeFaceLandmarker": "Face Landmarker (MediaPipe)",
}
@@ -0,0 +1,35 @@
# visualize_face.py
import mediapipe as mp
from typing import Optional, List, Any
from ...common.base_visualization_nodes import BaseLandmarkVisualizationNode
import logging
logger = logging.getLogger(__name__)
# Define standard face connections
FACE_CONNECTIONS = mp.solutions.face_mesh.FACEMESH_TESSELATION
class VisualizeFaceLandmarks(BaseLandmarkVisualizationNode):
LANDMARKS_TYPE = "FACE_LANDMARKS"
CONNECTIONS = FACE_CONNECTIONS
DEFAULT_DRAW_POINTS = False
DEFAULT_POINT_COLOR = "#FFFF00"
DEFAULT_POINT_RADIUS = 1
DEFAULT_CONNECTION_COLOR = "#00FF00"
DEFAULT_CONNECTION_THICKNESS = 1
SUPPORTS_VISIBILITY = False
CATEGORY = "Realtime Nodes/MediaPipe Vision/Face/FaceLandmark/Visualization"
DESCRIPTION = "Visualizes facial landmarks by drawing a mesh of points and connections on the face. Creates a detailed facial wireframe that shows the precise contours and features of detected faces."
def get_landmark_list(self, detection_result: Any) -> Optional[List[Any]]:
return detection_result
# Node registration
NODE_CLASS_MAPPINGS = {
"VisualizeFaceLandmarks": VisualizeFaceLandmarks
}
NODE_DISPLAY_NAME_MAPPINGS = {
"VisualizeFaceLandmarks": "Visualize Face Landmarks (MediaPipe)"
}
@@ -0,0 +1,109 @@
# node_wrappers/face/masking.py
import torch
import numpy as np
import cv2
import logging
from typing import List, Dict, Tuple
from .....src.mediapipe_vision.types import LandmarkPoint
from .....src.mediapipe_vision.utils.mask_utils import create_mask_from_points
logger = logging.getLogger(__name__)
_category = "Realtime Nodes/MediaPipe Vision/Face/FaceLandmark/Masking"
# --- Face Part Definitions ---
FACE_PART_INDICES = { # NOTE: These indices are EXAMPLES and NEED VERIFICATION!
"FACE_OVAL": [10, 338, 297, 332, 284, 251, 389, 356, 454, 323, 361, 288, 397, 365, 379, 378, 400, 377, 152, 148, 176, 149, 150, 136, 172, 58, 132, 93, 234, 127, 162, 21, 54, 103, 67, 109],
"LIPS_OUTER": [61, 146, 91, 181, 84, 17, 314, 405, 321, 375, 291, 62, 185, 40, 39, 37, 0, 267, 269, 270, 409, 292],
"LIPS_INNER": [78, 95, 88, 178, 87, 14, 317, 402, 318, 324, 308, 76, 191, 81, 82, 13, 312, 311, 310, 415, 306],
"LEFT_EYE": [33, 7, 163, 144, 145, 153, 154, 155, 133, 246, 161, 160, 159, 158, 157, 173],
"RIGHT_EYE": [263, 249, 390, 373, 374, 380, 381, 382, 362, 466, 388, 387, 386, 385, 384, 398],
"LEFT_EYEBROW": [70, 63, 105, 66, 107, 55, 65, 52, 53, 46],
"RIGHT_EYEBROW": [300, 293, 334, 296, 336, 285, 295, 282, 283, 276],
"NOSE_BRIDGE": [168, 6, 197, 195, 5],
"NOSE_TIP_LOWER": [5, 4, 1, 19, 94, 2],
}
# --- Node Definition ---
class MaskFromFaceLandmarksNode:
"""Creates a mask for a specific facial part from Face Landmarks."""
CATEGORY = _category
DESCRIPTION = "Creates precise masks for specific facial features like eyes, lips, nose, or the entire face. Perfect for targeted editing, inpainting, or creating effects on specific facial regions."
@classmethod
def INPUT_TYPES(cls):
part_names = list(FACE_PART_INDICES.keys())
return {
"required": {
"face_landmarks": ("FACE_LANDMARKS", {"tooltip": "Facial landmark data from the Face Landmarker node containing the precise locations of facial features"}),
"image_for_dimensions": ("IMAGE", {"tooltip": "Reference image to determine the size of the output mask - the mask will match these dimensions"}),
"part_name": (part_names, {"default": "FACE_OVAL",
"tooltip": "Choose which facial feature to create a mask for - FACE_OVAL (whole face), LIPS_OUTER (mouth), LEFT_EYE, etc."}),
}
}
RETURN_TYPES = ("MASK",)
FUNCTION = "generate_mask"
def generate_mask(self, face_landmarks: List[List[List[LandmarkPoint]]], image_for_dimensions: torch.Tensor, part_name: str):
if part_name not in FACE_PART_INDICES:
raise ValueError(f"Unknown face part name: '{part_name}'. Available: {list(FACE_PART_INDICES.keys())}")
# Get dimensions from image
if image_for_dimensions.dim() != 4:
raise ValueError("image_for_dimensions must be BHWC format.")
height = image_for_dimensions.shape[1]
width = image_for_dimensions.shape[2]
input_batch_size = image_for_dimensions.shape[0]
device = image_for_dimensions.device
# Validate batch sizes match
if len(face_landmarks) != input_batch_size:
logger.warning(f"Batch size mismatch between face_landmarks ({len(face_landmarks)}) and image_for_dimensions ({input_batch_size}). Using landmark batch size.")
# Adjust effective batch size based on landmarks if mismatched, though ideally they match
batch_size = len(face_landmarks)
else:
batch_size = input_batch_size
required_indices = FACE_PART_INDICES[part_name]
output_masks = []
for i in range(batch_size):
landmarks_for_image = face_landmarks[i] # List[List[LandmarkPoint]] (Faces in image)
combined_points_for_image = []
# Combine points from all detected faces for the selected part
for face_result_landmarks in landmarks_for_image: # face_result_landmarks is List[LandmarkPoint]
points_dict = {lm.index: lm for lm in face_result_landmarks} # Iterate directly over the list
part_points = []
for idx in required_indices:
if idx in points_dict:
lm = points_dict[idx]
px = int(lm.x * width)
py = int(lm.y * height)
part_points.append((px, py))
# Removed warning here to avoid spam if landmarks are just missing
combined_points_for_image.extend(part_points)
# Create mask from all collected points for this image
mask_tensor = create_mask_from_points(height, width, combined_points_for_image, device=device)
output_masks.append(mask_tensor)
if not output_masks:
logger.warning("No landmarks found to generate any masks. Returning zero mask batch.")
output_batch = torch.zeros((input_batch_size, 1, height, width), dtype=torch.float32, device=device)
else:
# Stack along batch dimension, ensure channel dim is present
output_batch = torch.stack(output_masks, dim=0) # Shape becomes (B, 1, H, W)
return (output_batch,)
# Remove GetImageDimensionsNode class definition
# --- Mappings ---
NODE_CLASS_MAPPINGS = {
"MaskFromFaceLandmarks": MaskFromFaceLandmarksNode,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"MaskFromFaceLandmarks": "Mask From Face Landmarks (MediaPipe)",
}
@@ -0,0 +1,219 @@
# node_wrappers/face/head_pose.py
import logging
import numpy as np
from scipy.spatial.transform import Rotation as R
from typing import Optional, List, Any, Tuple, Dict
from .....src.mediapipe_vision.types import TRANSFORM_MATRIX_LIST
from .....src.mediapipe_vision.utils.delta_utils import scale_value
from ...common.universal_landmark_nodes import FLOAT_EQUALITY_TOLERANCE # For trigger node comparisons
logger = logging.getLogger(__name__)
_category = "Realtime Nodes/MediaPipe Vision/Face/HeadPose"
# --- Helper Functions ---
def get_first_transform_matrix(matrix_result: TRANSFORM_MATRIX_LIST) -> Optional[np.ndarray]:
"""Extracts the first 4x4 transformation matrix from the result."""
if matrix_result is None or len(matrix_result) == 0:
return None
if len(matrix_result[0]) == 0:
return None
if matrix_result[0][0] is None:
return None
# FACE_TRANSFORM_MATRICES structure: List[List[np.ndarray]]
return matrix_result[0][0]
def decompose_transform_matrix(matrix: np.ndarray) -> Optional[Dict[str, float]]:
"""Decomposes a 4x4 matrix into translation (x, y, z) and Euler angles (pitch, yaw, roll)."""
if matrix is None or matrix.shape != (4, 4):
return None
try:
# Extract translation
translation = matrix[:3, 3]
# Extract rotation matrix part
rotation_matrix = matrix[:3, :3]
# Convert rotation matrix to Euler angles (ZYX convention common for head pose)
# Note: Scipy uses intrinsic rotations. ZYX order corresponds to yaw, pitch, roll.
rotation = R.from_matrix(rotation_matrix)
euler_angles_rad = rotation.as_euler('zyx', degrees=False)
# Convert radians to degrees for easier interpretation
euler_angles_deg = np.degrees(euler_angles_rad)
return {
"x": translation[0],
"y": translation[1],
"z": translation[2],
"yaw": euler_angles_deg[0], # Rotation around Y (vertical) axis
"pitch": euler_angles_deg[1], # Rotation around X (side-to-side) axis
"roll": euler_angles_deg[2], # Rotation around Z (front-to-back) axis
}
except Exception as e:
logger.error(f"Error decomposing matrix: {e}")
return None
# --- Control Node ---
HEAD_POSE_COMPONENTS = ["x", "y", "z", "pitch", "yaw", "roll"]
class HeadPoseControlFloatNode:
"""Outputs a scaled FLOAT based on a selected component (translation/rotation) of the head pose."""
CATEGORY = _category
RETURN_TYPES = ("FLOAT",)
FUNCTION = "execute"
DESCRIPTION = "Converts head position and orientation into floating point values. Use this to control parameters based on head movements like nodding, turning left/right, or tilting."
@classmethod
def INPUT_TYPES(cls):
# Example ranges - these might need tuning based on expected matrix values
default_ranges = {
"x": (-0.5, 0.5), "y": (-0.5, 0.5), "z": (-1.0, 0.5), # Example translation ranges
"pitch": (-45.0, 45.0), "yaw": (-60.0, 60.0), "roll": (-45.0, 45.0) # Example angle ranges (degrees)
}
default_component = "yaw"
default_min, default_max = default_ranges[default_component]
return {
"required": {
"face_transform_matrices": ("TRANSFORM_MATRIX_LIST", {"tooltip": "3D head orientation data from the Face Landmarker node"}),
"component": (HEAD_POSE_COMPONENTS, {"default": default_component,
"tooltip": "Which aspect of head movement to track - yaw (left/right), pitch (up/down), roll (tilt), or position (x/y/z)"}),
"value_min": ("FLOAT", {"default": default_min, "step": 0.1,
"tooltip": "The minimum value of the selected movement component to consider"}),
"value_max": ("FLOAT", {"default": default_max, "step": 0.1,
"tooltip": "The maximum value of the selected movement component to consider"}),
"output_min_float": ("FLOAT", {"default": 0.0, "step": 0.01,
"tooltip": "The minimum output value when head position is at or below value_min"}),
"output_max_float": ("FLOAT", {"default": 1.0, "step": 0.01,
"tooltip": "The maximum output value when head position is at or above value_max"}),
"clamp": ("BOOLEAN", {"default": True,
"tooltip": "When enabled, restricts output values to stay within min and max range"}),
}
}
def execute(self, face_transform_matrices, component, value_min, value_max, output_min_float, output_max_float, clamp):
matrix = get_first_transform_matrix(face_transform_matrices)
decomposed = decompose_transform_matrix(matrix)
raw_value = None
if decomposed and component in decomposed:
raw_value = decomposed[component]
else:
logger.debug(f"Could not get head pose component '{component}'")
scaled_value = scale_value(raw_value, value_min, value_max, output_min_float, output_max_float, clamp)
# Default to min output value on error
output_val = scaled_value if scaled_value is not None else float(output_min_float)
return (output_val,)
class HeadPoseControlIntNode:
"""Outputs a scaled INT based on a selected component (translation/rotation) of the head pose."""
CATEGORY = _category
RETURN_TYPES = ("INT",)
FUNCTION = "execute"
DESCRIPTION = "Converts head position and orientation into integer values. Perfect for selecting options, changing steps/seeds, or triggering numbered presets based on where someone is looking."
@classmethod
def INPUT_TYPES(cls):
# Share the same default range logic as Float node
default_ranges = {
"x": (-0.5, 0.5), "y": (-0.5, 0.5), "z": (-1.0, 0.5),
"pitch": (-45.0, 45.0), "yaw": (-60.0, 60.0), "roll": (-45.0, 45.0)
}
default_component = "yaw"
default_min, default_max = default_ranges[default_component]
return {
"required": {
"face_transform_matrices": ("TRANSFORM_MATRIX_LIST", {"tooltip": "3D head orientation data from the Face Landmarker node"}),
"component": (HEAD_POSE_COMPONENTS, {"default": default_component,
"tooltip": "Which aspect of head movement to track - yaw (left/right), pitch (up/down), roll (tilt), or position (x/y/z)"}),
"value_min": ("FLOAT", {"default": default_min, "step": 0.1,
"tooltip": "The minimum value of the selected movement component to consider"}),
"value_max": ("FLOAT", {"default": default_max, "step": 0.1,
"tooltip": "The maximum value of the selected movement component to consider"}),
"output_min_int": ("INT", {"default": 0,
"tooltip": "The minimum integer output value when head position is at or below value_min"}),
"output_max_int": ("INT", {"default": 100,
"tooltip": "The maximum integer output value when head position is at or above value_max"}),
"clamp": ("BOOLEAN", {"default": True,
"tooltip": "When enabled, restricts output values to stay within min and max range"}),
}
}
def execute(self, face_transform_matrices, component, value_min, value_max, output_min_int, output_max_int, clamp):
matrix = get_first_transform_matrix(face_transform_matrices)
decomposed = decompose_transform_matrix(matrix)
raw_value = None
if decomposed and component in decomposed:
raw_value = decomposed[component]
else:
logger.debug(f"Could not get head pose component '{component}'")
scaled_value = scale_value(raw_value, value_min, value_max, output_min_int, output_max_int, clamp)
# Default to min output value on error
output_val = int(round(scaled_value)) if scaled_value is not None else int(output_min_int)
return (output_val,)
# --- Trigger Node ---
class HeadPoseTriggerNode:
"""Outputs a BOOLEAN trigger if a head pose component crosses a threshold."""
CATEGORY = _category
RETURN_TYPES = ("BOOLEAN",)
FUNCTION = "execute"
DESCRIPTION = "Creates a true/false trigger based on head movements exceeding thresholds. Use it to activate workflows when someone nods, looks left/right, or tilts their head a certain amount."
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"face_transform_matrices": ("TRANSFORM_MATRIX_LIST", {"tooltip": "3D head orientation data from the Face Landmarker node"}),
"component": (HEAD_POSE_COMPONENTS, {"default": "yaw",
"tooltip": "Which aspect of head movement to track - yaw (left/right), pitch (up/down), roll (tilt), or position (x/y/z)"}),
"threshold": ("FLOAT", {"default": 30.0, "step": 0.5,
"tooltip": "The angle or position value that triggers the condition (example: 30 degrees of head turn)"}),
"condition": (["Above", "Below", "Equals", "Not Equals"],
{"tooltip": "The condition that determines when to trigger - Above (movement > threshold), Below (movement < threshold), etc."}),
}
}
def execute(self, face_transform_matrices, component, threshold, condition):
matrix = get_first_transform_matrix(face_transform_matrices)
decomposed = decompose_transform_matrix(matrix)
raw_value = None
if decomposed and component in decomposed:
raw_value = decomposed[component]
else:
logger.debug(f"Could not get head pose component '{component}' for trigger")
triggered = False
if raw_value is not None:
if condition == "Above" and raw_value > threshold:
triggered = True
elif condition == "Below" and raw_value < threshold:
triggered = True
elif condition == "Equals" and abs(raw_value - threshold) < FLOAT_EQUALITY_TOLERANCE:
triggered = True
elif condition == "Not Equals" and abs(raw_value - threshold) >= FLOAT_EQUALITY_TOLERANCE:
triggered = True
return (triggered,)
# --- Mappings ---
NODE_CLASS_MAPPINGS = {
"HeadPoseControlFloat": HeadPoseControlFloatNode,
"HeadPoseControlInt": HeadPoseControlIntNode,
"HeadPoseTrigger": HeadPoseTriggerNode,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"HeadPoseControlFloat": "Head Pose Control (Float)",
"HeadPoseControlInt": "Head Pose Control (Int)",
"HeadPoseTrigger": "Head Pose Trigger",
}
@@ -0,0 +1,33 @@
import logging
from ...common.universal_landmark_nodes import LandmarkPositionBaseNode
logger = logging.getLogger(__name__)
_category = "Realtime Nodes/MediaPipe Vision/Face/FaceLandmark/Position"
class FaceLandmarkPositionNode(LandmarkPositionBaseNode):
"""Extracts position lists (x, y, z, vis, pres) for a specific landmark index
from FACE_LANDMARKS across the batch."""
CATEGORY = _category
# Define the specific input type string this node expects
LANDMARKS_TYPE = "FACE_LANDMARKS"
DESCRIPTION = "Extracts lists of coordinates (x, y, z) and properties (visibility, presence) for a specific Face Landmark index across an entire batch. Outputs separate lists for each property."
@classmethod
def INPUT_TYPES(cls):
# Define the specific input type for this node
return {
"required": {
"landmarks": (cls.LANDMARKS_TYPE, {"forceInput": True}), # Use the class variable
"landmark_index": ("INT", {"default": 0, "min": 0}),
}
}
# --- Node Mappings ---
NODE_CLASS_MAPPINGS = {
"FaceLandmarkPosition": FaceLandmarkPositionNode,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"FaceLandmarkPosition": "Face Landmark Position Extractor (Batch)",
}
@@ -0,0 +1,60 @@
# node_wrappers/face/texture_warp.py
import logging
import torch
import numpy as np
logger = logging.getLogger(__name__)
_category = "Realtime Nodes/MediaPipe Vision/Face/Experimental"
class FaceTextureWarpNode:
"""(Experimental) Attempts to warp a texture based on the face transform matrix.
NOTE: This is a conceptual placeholder. Actual warping requires significant
integration with sampling or image processing libraries.
"""
CATEGORY = _category
RETURN_TYPES = ("IMAGE",)
FUNCTION = "execute"
DESCRIPTION = "(Experimental) Designed to warp textures to match face orientation, potentially allowing for facial texture replacement or customized face overlays. Currently a placeholder for future implementation."
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"target_image": ("IMAGE", {"tooltip": "The base image containing a face onto which the texture will be applied"}),
"texture_image": ("IMAGE", {"tooltip": "The texture image to be warped onto the face (could be another face, pattern, or design)"}),
"face_transform_matrices": ("TRANSFORM_MATRIX_LIST", {"tooltip": "3D head orientation data from the Face Landmarker node to determine warping angles"}),
# Potentially add mask input for blending?
},
"optional": {
"offset_x": ("FLOAT", {"default": 0.0, "step": 0.01, "tooltip": "Horizontal positioning adjustment for the texture overlay"}),
"offset_y": ("FLOAT", {"default": 0.0, "step": 0.01, "tooltip": "Vertical positioning adjustment for the texture overlay"}),
"scale": ("FLOAT", {"default": 1.0, "step": 0.01, "tooltip": "Size adjustment for the texture overlay (1.0 = original size)"}),
}
}
def execute(self, target_image, texture_image, face_transform_matrices, offset_x=0.0, offset_y=0.0, scale=1.0):
logger.warning(f"{self.__class__.__name__} is experimental and does not perform actual warping yet.")
#TODO: implement this
# --- Conceptual Logic ---
# 1. Get the primary face transform matrix
# 2. Potentially combine with offset/scale parameters
# 3. Define or load a canonical 3D face model mesh.
# 4. Project the canonical mesh onto the 2D target_image plane using the matrix.
# 5. Determine the UV coordinates on the texture_image corresponding to the projected vertices.
# 6. Perform perspective warping of the texture_image onto the target_image using the
# calculated transformation (e.g., using OpenCV's getPerspectiveTransform and warpPerspective,
# or more advanced 3D rendering techniques).
# 7. Blend the warped texture onto the target image (potentially using a face mask).
# For now, just return the original target image
return (target_image,)
# --- Mappings ---
NODE_CLASS_MAPPINGS = {
"FaceTextureWarp": FaceTextureWarpNode,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"FaceTextureWarp": "Face Texture Warp (Experimental)",
}
@@ -0,0 +1,81 @@
import torch
import logging
# Import Base Loader and Detector
from ...common.model_loader import MediaPipeModelLoaderBaseNode
from .....src.mediapipe_vision.face_stylization.detector import FaceStylizer
logger = logging.getLogger(__name__)
_category = "Realtime Nodes/MediaPipe Vision/Face/FaceStylization"
# --- Model Loader ---
class MediaPipeFaceStylizerModelLoaderNode(MediaPipeModelLoaderBaseNode):
"""ComfyUI node for loading MediaPipe Face Stylizer models."""
TASK_TYPE = "face_stylizer" # Need to add this task to MEDIAPIPE_MODELS
RETURN_TYPES = ("FACE_STYLIZER_MODEL_INFO",)
RETURN_NAMES = ("model_info",)
CATEGORY = _category
DESCRIPTION = "Loads a MediaPipe Face Stylizer model that can transform face images into artistic styles. Required before using the Face Stylizer node."
# --- Stylizer Node ---
class MediaPipeFaceStylizerNode:
"""ComfyUI node for MediaPipe Face Stylization."""
def __init__(self):
self._detector: FaceStylizer = None
self._model_path: str = None
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"image": ("IMAGE", {"tooltip": "Input image containing faces to be artistically stylized"}),
"model_info": ("FACE_STYLIZER_MODEL_INFO", {"tooltip": "Face stylizer model loaded from the 'Load Face Stylizer Model' node"}),
}
}
RETURN_TYPES = ("IMAGE",) # Output is a stylized image
RETURN_NAMES = ("stylized_image",)
FUNCTION = "stylize_face"
CATEGORY = _category
DESCRIPTION = "Applies artistic styles to faces in images without affecting the background. Creates cartoon-like or artistic renditions of faces while preserving identity."
def stylize_face(self, image: torch.Tensor, model_info: dict):
"""Applies face stylization."""
task_type = model_info.get('task_type')
expected_task_type = MediaPipeFaceStylizerModelLoaderNode.TASK_TYPE
if not isinstance(model_info, dict) or task_type != expected_task_type:
raise ValueError(f"Invalid model_info. Expected task_type '{expected_task_type}' but 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.")
# Manage detector instance
if self._detector is None or self._model_path != model_path:
if self._detector and hasattr(self._detector, 'close'):
try: self._detector.close()
except Exception as e: logger.warning(f"Error closing FaceStylizer: {e}")
logger.info(f"Creating new FaceStylizer instance for {model_path}")
self._detector = FaceStylizer(model_path)
self._model_path = model_path
# Call the detector's stylize method
stylized_image_batch = self._detector.stylize(image)
return (stylized_image_batch,)
def __del__(self):
if hasattr(self, '_detector') and self._detector and hasattr(self._detector, 'close'):
try: self._detector.close()
except Exception as e: logger.warning(f"Error closing FaceStylizer in __del__: {e}")
self._detector = None
# --- Mappings ---
NODE_CLASS_MAPPINGS = {
"MediaPipeFaceStylizerModelLoader": MediaPipeFaceStylizerModelLoaderNode,
"MediaPipeFaceStylizer": MediaPipeFaceStylizerNode,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"MediaPipeFaceStylizerModelLoader": "Load Face Stylizer Model (MediaPipe)",
"MediaPipeFaceStylizer": "Face Stylizer (MediaPipe)",
}
@@ -0,0 +1 @@
@@ -0,0 +1,140 @@
# visualize_gestures.py
import torch
import numpy as np
import cv2
from ...common.base_visualization_nodes import BaseDetectionVisualizationNode
import logging
logger = logging.getLogger(__name__)
_category = "Realtime Nodes/MediaPipe Vision/Hand/GestureVisualization"
class VisualizeGestureRecognitions(BaseDetectionVisualizationNode):
"""Draws recognized gesture labels onto images near the detected hand."""
CATEGORY = _category
DESCRIPTION = "Visualizes recognized hand gestures by drawing text labels next to hands in the image. Displays gesture names like 'Thumbs Up', 'Victory', 'Open Palm' etc. near each detected hand."
# Override font scale for gesture labels (larger than default)
FONT_SCALE = 0.7 # Larger than default for gesture labels
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"gesture_recognitions": ("GESTURE_RECOGNITIONS", {"tooltip": "Gesture recognition results from the Gesture Recognizer node"}),
"hand_landmarks": ("HAND_LANDMARKS", {"tooltip": "Hand landmark data from the Hand Landmarker node - needed to position the text labels"}),
"image": ("IMAGE", {"tooltip": "The image on which to draw gesture labels"}),
},
"optional": {
"label_color": ("STRING", {"default": "#FFFFFF", "tooltip": "Hex color code for the gesture label text (default: white)"}),
"label_bg_color": ("STRING", {"default": "#0000FF", "tooltip": "Hex color code for the label background (default: blue)"}),
"y_offset": ("INT", {"default": -30, "step": 5, "tooltip": "Vertical position adjustment for text labels relative to the wrist - negative moves up, positive moves down"}),
}
}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "draw_visualizations"
def draw_visualizations(self, gesture_recognitions, hand_landmarks, image,
label_color="#FFFFFF", label_bg_color="#0000FF", y_offset=-30):
"""Main validation and batch processing function."""
if not gesture_recognitions:
logger.warning("VisualizeGestureRecognitions: No gestures provided.")
return (image,)
if not hand_landmarks:
logger.warning("VisualizeGestureRecognitions: Hand landmarks required for positioning text.")
return (image,)
# Basic batch size check
batch_size = len(image)
if len(gesture_recognitions) != batch_size or len(hand_landmarks) != batch_size:
logger.error(f"VisualizeGestureRecognitions: Batch size mismatch. Images: {batch_size}, Gestures: {len(gesture_recognitions)}, Landmarks: {len(hand_landmarks)}. Cannot proceed.")
return (image,)
# Pair gestures with landmarks for convenience
paired_data = list(zip(gesture_recognitions, hand_landmarks))
# Use the base class process_batch method with our custom drawing function
return self.process_batch(
paired_data, # We're passing paired data instead of just detections
image,
self.draw_gestures,
label_color=label_color,
label_bg_color=label_bg_color,
y_offset=y_offset
)
def draw_gestures(self, image, width, height, paired_data,
label_color="#FFFFFF", label_bg_color="#0000FF", y_offset=-30):
"""Draw gestures for a single image."""
gestures_for_image, landmarks_for_image = paired_data
# Check if number of detected hands matches for gestures and landmarks
if len(gestures_for_image) != len(landmarks_for_image):
logger.warning(f"VisualizeGestureRecognitions: Mismatch between detected gestures and landmarks. Skipping labels.")
return
if not gestures_for_image:
return
# Iterate through each detected hand/gesture set for this image
for hand_idx, gesture_result in enumerate(gestures_for_image):
landmark_result = landmarks_for_image[hand_idx]
# Get top gesture category
top_gesture = "N/A"
if gesture_result.gestures:
top_gesture = gesture_result.gestures[0].display_name or gesture_result.gestures[0].category_name or "N/A"
# Use wrist landmark (index 0) for positioning text
if not landmark_result.landmarks or len(landmark_result.landmarks) == 0:
logger.warning(f"VisualizeGestureRecognitions: No landmarks found for this hand to position text.")
continue
wrist_lm = landmark_result.landmarks[0] # Assuming index 0 is wrist
cx = int(wrist_lm.x * width)
cy = int(wrist_lm.y * height) + y_offset # Apply offset
label_text = f"{top_gesture}"
# Draw label with background
try:
# Get text size
text_width, text_height = self.get_text_dimensions(label_text, self.FONT_SCALE)
# Center the text around the wrist point
text_x = cx - (text_width // 2)
text_y = cy - (text_height // 2)
# Draw the text with background
self.draw_text_with_background(
image,
label_text,
(text_x, text_y),
text_color=label_color,
bg_color=label_bg_color,
font_scale=self.FONT_SCALE
)
except Exception as e:
logger.error(f"Error drawing text label for gesture: {e}")
# Fallback: Draw text directly
cv2.putText(
image,
label_text,
(cx, cy),
cv2.FONT_HERSHEY_SIMPLEX,
self.FONT_SCALE,
self._hex_to_bgr(label_color),
1,
cv2.LINE_AA
)
# Node registration
NODE_CLASS_MAPPINGS = {
"VisualizeGestureRecognitions": VisualizeGestureRecognitions
}
NODE_DISPLAY_NAME_MAPPINGS = {
"VisualizeGestureRecognitions": "Visualize Gestures (MediaPipe)"
}
@@ -0,0 +1,85 @@
import torch
import logging
# Import Base Loader and Detector
from ...common.model_loader import MediaPipeModelLoaderBaseNode
from ...common.base_detector_node import BaseMediaPipeDetectorNode
from .....src.mediapipe_vision.gesture_recognition.detector import GestureRecognizer
logger = logging.getLogger(__name__)
_category = "Realtime Nodes/MediaPipe Vision/Hand/GestureRecognition"
# --- Model Loader ---
class MediaPipeGestureRecognizerModelLoaderNode(MediaPipeModelLoaderBaseNode):
"""ComfyUI node for loading MediaPipe Gesture Recognizer models."""
TASK_TYPE = "gesture_recognizer" # Need to add this task to MEDIAPIPE_MODELS in src/model_loader.py
RETURN_TYPES = ("GESTURE_RECOGNIZER_MODEL_INFO",)
RETURN_NAMES = ("model_info",)
CATEGORY = _category
DESCRIPTION = "Loads a MediaPipe Gesture Recognizer model for identifying hand gestures in images, such as thumbs up, open palm, pointing, etc. Required before using the Gesture Recognizer node."
# --- Recognizer Node ---
class MediaPipeGestureRecognizerNode(BaseMediaPipeDetectorNode):
"""ComfyUI node for MediaPipe Gesture Recognition."""
# Define class variables required by the base class
DETECTOR_CLASS = GestureRecognizer
MODEL_INFO_TYPE = "GESTURE_RECOGNIZER_MODEL_INFO"
EXPECTED_TASK_TYPE = "gesture_recognizer"
RETURN_TYPES = ("GESTURE_RECOGNITIONS",)
RETURN_NAMES = ("gesture_recognitions",)
FUNCTION = "recognize_gestures"
CATEGORY = _category
DESCRIPTION = "Identifies common hand gestures in images like thumbs up, victory sign, open palm, etc. Provides both gesture names and confidence scores for recognized hand poses."
@classmethod
def INPUT_TYPES(cls):
# Start with the base inputs from the parent class
inputs = super().INPUT_TYPES()
# Add gesture-specific parameters
inputs["required"].update({
"max_results": ("INT", {"default": 2, "min": 1, "max": 4, "step": 1,
"tooltip": "Maximum number of hands to analyze in the image - higher values can detect more hands but use more processing power"}),
"min_confidence": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01,
"tooltip": "Minimum confidence threshold for hand detection - lower values detect more hands but may include false positives"}),
"min_tracking_confidence": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01,
"tooltip": "(VIDEO mode) Minimum confidence threshold for tracking hands between frames - lower values maintain tracking longer but may be less accurate"}),
"min_presence_confidence": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01,
"tooltip": "(VIDEO mode) Minimum confidence that a hand is present - lower values may detect hands that are less clear"}),
})
return inputs
def recognize_gestures(self, image: torch.Tensor, model_info: dict, max_results: int,
min_confidence: float, min_tracking_confidence: float, min_presence_confidence: float,
running_mode: str, delegate: str):
"""Performs gesture recognition on the input image."""
# Validate model_info and get model path
model_path = self.validate_model_info(model_info)
# Initialize or update detector
detector = self.initialize_or_update_detector(model_path)
# Call the detector's recognize method with standardized parameters
gesture_results_batch = detector.recognize(
image,
num_hands=max_results,
min_detection_confidence=min_confidence,
min_tracking_confidence=min_tracking_confidence,
min_presence_confidence=min_presence_confidence,
running_mode=running_mode,
delegate=delegate
)
return (gesture_results_batch,)
# --- Mappings ---
NODE_CLASS_MAPPINGS = {
"MediaPipeGestureRecognizerModelLoader": MediaPipeGestureRecognizerModelLoaderNode,
"MediaPipeGestureRecognizer": MediaPipeGestureRecognizerNode,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"MediaPipeGestureRecognizerModelLoader": "Load Gesture Recognizer Model (MediaPipe)",
"MediaPipeGestureRecognizer": "Gesture Recognizer (MediaPipe)",
}
@@ -0,0 +1,89 @@
import torch
from .....src.mediapipe_vision.hand_landmark.detector import HandLandmarkDetector
from ...common.model_loader import MediaPipeModelLoaderBaseNode
from ...common.base_detector_node import BaseMediaPipeDetectorNode
import logging
logger = logging.getLogger(__name__)
_category = "Realtime Nodes/MediaPipe Vision/Hand/HandLandmark"
class MediaPipeHandLandmarkerModelLoaderNode(MediaPipeModelLoaderBaseNode):
"""ComfyUI node for loading MediaPipe Hand Landmarker models."""
TASK_TYPE = "hand_landmarker"
RETURN_TYPES = ("HAND_LANDMARKER_MODEL_INFO",)
RETURN_NAMES = ("model_info",)
CATEGORY = _category
DESCRIPTION = "Loads a MediaPipe Hand Landmarker model for detecting hand positions and gestures in images. Required before using the Hand Landmarker node."
class MediaPipeHandLandmarkerNode(BaseMediaPipeDetectorNode):
"""ComfyUI node for MediaPipe Hand Landmark Detection."""
# Define class variables required by the base class
DETECTOR_CLASS = HandLandmarkDetector
MODEL_INFO_TYPE = "HAND_LANDMARKER_MODEL_INFO"
EXPECTED_TASK_TYPE = "hand_landmarker"
RETURN_TYPES = ("HAND_LANDMARKS", "HANDEDNESS_LIST")
RETURN_NAMES = ("landmarks", "handedness")
FUNCTION = "detect"
CATEGORY = _category
DESCRIPTION = "Detects hands in images and provides precise landmark points for each finger joint and hand position. Also identifies whether each hand is left or right."
@classmethod
def INPUT_TYPES(cls):
# Start with the base inputs from the parent class
inputs = super().INPUT_TYPES()
# Add hand-specific parameters
inputs["required"].update({
"max_results": ("INT", {"default": 2, "min": 1, "max": 4, "step": 1,
"tooltip": "Maximum number of hands to detect and track in the image - higher values can detect more hands but use more processing power"}),
"min_confidence": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01,
"tooltip": "Minimum confidence threshold for hand detection - lower values detect more hands but may include false positives"}),
"min_presence_confidence": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01,
"tooltip": "(VIDEO mode) Minimum confidence that a hand is present - lower values may detect hands that are less clear"}),
"min_tracking_confidence": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01,
"tooltip": "(VIDEO mode) Minimum confidence for tracking a hand between frames - lower values maintain tracking longer but may drift"}),
})
return inputs
def detect(self, image: torch.Tensor, model_info: dict, max_results: int,
min_confidence: float, min_presence_confidence: float,
min_tracking_confidence: float, running_mode: str, delegate: str):
"""Performs hand landmark detection with the configured parameters."""
# Validate model_info and get model path
model_path = self.validate_model_info(model_info)
# Initialize or update detector
detector = self.initialize_or_update_detector(model_path)
# Perform detection with all parameters
hand_landmarks_batch = detector.detect(
image,
num_hands=max_results,
min_detection_confidence=min_confidence,
min_presence_confidence=min_presence_confidence,
min_tracking_confidence=min_tracking_confidence,
running_mode=running_mode,
delegate=delegate
)
# Extract handedness information
handedness_batch = []
if hand_landmarks_batch:
for image_hands in hand_landmarks_batch:
handedness_per_image = [hand.handedness for hand in image_hands if hand.handedness is not None]
handedness_batch.append(handedness_per_image)
return (hand_landmarks_batch, handedness_batch)
NODE_CLASS_MAPPINGS = {
"MediaPipeHandLandmarkerModelLoader": MediaPipeHandLandmarkerModelLoaderNode,
"MediaPipeHandLandmarker": MediaPipeHandLandmarkerNode,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"MediaPipeHandLandmarkerModelLoader": "Load Hand Landmarker Model (MediaPipe)",
"MediaPipeHandLandmarker": "Hand Landmarker (MediaPipe)",
}
@@ -0,0 +1,107 @@
# node_wrappers/hand/masking.py
import torch
import numpy as np
import logging
from typing import List
from .....src.mediapipe_vision.types import HandLandmarksResult
# Reusing the helper function from face masking
from .....src.mediapipe_vision.utils.mask_utils import create_mask_from_points
logger = logging.getLogger(__name__)
_category = "Realtime Nodes/MediaPipe Vision/Hand/HandLandmark/Masking"
# --- Hand Part Definitions (Based on standard MediaPipe Hand Indices) ---
# These indices need verification
HAND_PART_INDICES = {
"PALM": [0, 1, 5, 9, 13, 17], # Example
"THUMB": [1, 2, 3, 4], # Example
"INDEX_FINGER": [5, 6, 7, 8], # Example
"MIDDLE_FINGER": [9, 10, 11, 12], # Example
"RING_FINGER": [13, 14, 15, 16], # Example
"PINKY_FINGER": [17, 18, 19, 20], # Example
"WRIST": [0], # Single point, convex hull won't work
"ALL_FINGERS": list(range(1, 21)), # Example excluding wrist
"FULL_HAND": list(range(21)), # All points
}
# --- Node Definition ---
class MaskFromHandLandmarksNode:
"""Creates a mask for a specific hand part from Hand Landmarks."""
CATEGORY = _category
DESCRIPTION = "Creates precise masks for specific hand parts like individual fingers, palm, or the entire hand. Perfect for targeted editing or creating effects on hands and fingers."
@classmethod
def INPUT_TYPES(cls):
part_names = list(HAND_PART_INDICES.keys())
return {
"required": {
"hand_landmarks": ("HAND_LANDMARKS", {"tooltip": "Hand landmark data from the Hand Landmarker node containing the precise locations of hand and finger joints"}),
"image_for_dimensions": ("IMAGE", {"tooltip": "Reference image to determine the size of the output mask - the mask will match these dimensions"}),
"part_name": (part_names, {"default": "FULL_HAND",
"tooltip": "Choose which part of the hand to create a mask for - FULL_HAND (entire hand), THUMB, INDEX_FINGER, PALM, etc."}),
}
}
RETURN_TYPES = ("MASK",)
FUNCTION = "generate_mask"
def generate_mask(self, hand_landmarks: List[List[HandLandmarksResult]], image_for_dimensions: torch.Tensor, part_name: str):
if part_name not in HAND_PART_INDICES:
raise ValueError(f"Unknown hand part name: '{part_name}'. Available: {list(HAND_PART_INDICES.keys())}")
# Get dimensions from image
if image_for_dimensions.dim() != 4:
raise ValueError("image_for_dimensions must be BHWC format.")
height = image_for_dimensions.shape[1]
width = image_for_dimensions.shape[2]
input_batch_size = image_for_dimensions.shape[0]
device = image_for_dimensions.device
# Validate batch sizes match
if len(hand_landmarks) != input_batch_size:
logger.warning(f"Batch size mismatch between hand_landmarks ({len(hand_landmarks)}) and image_for_dimensions ({input_batch_size}). Using landmark batch size.")
batch_size = len(hand_landmarks)
else:
batch_size = input_batch_size
required_indices = HAND_PART_INDICES[part_name]
output_masks = []
for i in range(batch_size):
landmarks_for_image = hand_landmarks[i] # List[HandLandmarksResult]
combined_points_for_image = []
# Combine points from all detected hands for the selected part
for hand_result in landmarks_for_image:
points_dict = {lm.index: lm for lm in hand_result.landmarks}
part_points = []
for idx in required_indices:
if idx in points_dict:
lm = points_dict[idx]
px = int(lm.x * width)
py = int(lm.y * height)
part_points.append((px, py))
# No warning for missing index
combined_points_for_image.extend(part_points)
# Create mask from all collected points for this image
mask_tensor = create_mask_from_points(height, width, combined_points_for_image, device=device)
output_masks.append(mask_tensor)
if not output_masks:
logger.warning("No landmarks found to generate any masks. Returning zero mask batch.")
output_batch = torch.zeros((input_batch_size, 1, height, width), dtype=torch.float32, device=device)
else:
output_batch = torch.stack(output_masks, dim=0)
return (output_batch,)
# --- Mappings ---
NODE_CLASS_MAPPINGS = {
"MaskFromHandLandmarks": MaskFromHandLandmarksNode,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"MaskFromHandLandmarks": "Mask From Hand Landmarks (MediaPipe)",
}
@@ -0,0 +1,40 @@
# visualize_hand.py
import mediapipe as mp
from typing import Optional, List, Any
from .....src.mediapipe_vision.types import HandLandmarksResult # Import specific result type if needed for extraction
from ...common.base_visualization_nodes import BaseLandmarkVisualizationNode
import logging
logger = logging.getLogger(__name__)
_category = "Realtime Nodes/MediaPipe Vision/Hand/HandLandmark/Visualization"
# Define standard hand connections
HAND_CONNECTIONS = mp.solutions.hands.HAND_CONNECTIONS
class VisualizeHandLandmarks(BaseLandmarkVisualizationNode):
"""Draws Hand landmarks and connections onto images."""
# Override base class properties
LANDMARKS_TYPE = "HAND_LANDMARKS"
CONNECTIONS = HAND_CONNECTIONS
DEFAULT_DRAW_POINTS = True
DEFAULT_POINT_COLOR = "#FF0000" # Red
DEFAULT_POINT_RADIUS = 3
DEFAULT_CONNECTION_COLOR = "#00FF00" # Green
DEFAULT_CONNECTION_THICKNESS = 2
SUPPORTS_VISIBILITY = False
CATEGORY = _category
DESCRIPTION = "Visualizes hand landmarks by drawing points and connections on detected hands. Creates a wireframe overlay showing finger joints, palm, and hand structure."
def get_landmark_list(self, detection_result: Any) -> Optional[List[Any]]:
return detection_result.landmarks
# Node registration
NODE_CLASS_MAPPINGS = {
"VisualizeHandLandmarks": VisualizeHandLandmarks
}
NODE_DISPLAY_NAME_MAPPINGS = {
"VisualizeHandLandmarks": "Visualize Hand Landmarks (MediaPipe)"
}
@@ -0,0 +1,39 @@
# node_wrappers/hand/position.py
import logging
from typing import List, Tuple, Any, Optional
# Import the base class
from ...common.universal_landmark_nodes import LandmarkPositionBaseNode
# Import the specific landmark type for clarity and potential future use
from .....src.mediapipe_vision.types import HAND_LANDMARKS
logger = logging.getLogger(__name__)
_category = "Realtime Nodes/MediaPipe Vision/Hand/HandLandmark/Position"
class HandLandmarkPositionNode(LandmarkPositionBaseNode):
"""Extracts position lists (x, y, z, vis, pres) for a specific landmark index
from HAND_LANDMARKS across the batch."""
CATEGORY = _category
# Define the specific input type string this node expects
LANDMARKS_TYPE = "HAND_LANDMARKS"
DESCRIPTION = "Extracts lists of coordinates (x, y, z) and properties (visibility, presence) for a specific Hand Landmark index across an entire batch. Supports world coordinates. Outputs separate lists for each property."
@classmethod
def INPUT_TYPES(cls):
# Define the specific input type for this node
return {
"required": {
"landmarks": (cls.LANDMARKS_TYPE, {"forceInput": True}),
"landmark_index": ("INT", {"default": 8, "min": 0, "max": 20, "tooltip": "Index Tip=8, Thumb Tip=4"}), # Hand specific default/range
"use_world_coordinates": ("BOOLEAN", {"default": False, "tooltip": "Use world coordinates if available"}),
}
}
# --- Node Mappings ---
NODE_CLASS_MAPPINGS = {
"HandLandmarkPosition": HandLandmarkPositionNode,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"HandLandmarkPosition": "Hand Landmark Position Extractor (Batch)",
}
@@ -0,0 +1,10 @@
"""MediaPipe Holistic node wrappers.
this module contains the nodes for the holistic solution. Holistic is only available in the legacy api
and thus is stands aout among the other imeplementations in this project. When the official update
to this is released, this module will be updated to use the new api.
"""
# Import modules to ensure they're registered
from . import holistic_landmark
from . import holistic_landmark_visualization
@@ -0,0 +1,106 @@
import torch
import cv2
import mediapipe as mp
import numpy as np
from ....src.mediapipe_vision.holistic_landmark.detector import HolisticLandmarkDetector
from ..common.base_detector_node import BaseMediaPipeDetectorNode
import logging
logger = logging.getLogger(__name__)
_category = "Realtime Nodes/MediaPipe Vision/Holistic"
class MediaPipeHolisticLandmarkerNode(BaseMediaPipeDetectorNode):
"""ComfyUI node for MediaPipe Holistic Landmark Detection.
This node uses the legacy MediaPipe Holistic API that combines face, pose and hand
tracking in a single detector.
Note: This uses the legacy MediaPipe Python Solution API while we wait for an official
Holistic Task implementation, and thus it stands out among other implementations in this project.
It does not require a model file to be loaded separately.
"""
# Define class variables required by the base class
DETECTOR_CLASS = HolisticLandmarkDetector
MODEL_INFO_TYPE = None # Not used - legacy solution doesn't need model file
EXPECTED_TASK_TYPE = None # Not used - legacy solution doesn't need model info
RETURN_TYPES = ("HOLISTIC_LANDMARKS", "MASK_LIST",)
RETURN_NAMES = ("landmarks", "segmentation_masks")
FUNCTION = "detect_holistic"
CATEGORY = _category
DESCRIPTION = """
LEGACY MODEL ***USE CAUTION***:
RUNNING THIS DETECTOR WILL BREAK ALL OTHER DETECTORS THAT USE THE TASK API UNTIL COMFYUI IS RESTARTED.
Detects face, pose, and hand landmarks in a unified detection. Provides comprehensive tracking of an entire person. This uses the legacy MediaPipe API rather than the newer Task API.
"""
@classmethod
def INPUT_TYPES(cls):
"""Define inputs for holistic detection."""
# We override the parent INPUT_TYPES completely since we don't need model_info
return {
"required": {
"image": ("IMAGE", {"tooltip": "Input image for detection"}),
"min_confidence": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01,
"tooltip": "Minimum confidence threshold for detection - lower values detect more poses but may include false positives"}),
"min_tracking_confidence": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01,
"tooltip": "Minimum confidence threshold for tracking - lower values maintain tracking longer but may be less accurate"}),
"model_complexity": ("INT", {"default": 1, "min": 0, "max": 2, "step": 1,
"tooltip": "Model complexity (0=Lite, 1=Full, 2=Heavy) - higher values are more accurate but slower"}),
"enable_segmentation": ("BOOLEAN", {"default": False,
"tooltip": "When enabled, generates masks that separate people from the background"}),
"refine_face_landmarks": ("BOOLEAN", {"default": False,
"tooltip": "When enabled, provides more detailed face landmarks around eyes and lips"}),
"static_image_mode": ("BOOLEAN", {"default": False,
"tooltip": "When enabled, treats each frame as independent - more accurate but slower and doesn't benefit from tracking"}),
}
}
def detect_holistic(self, image: torch.Tensor,
min_confidence: float,
min_tracking_confidence: float,
model_complexity: int,
enable_segmentation: bool,
refine_face_landmarks: bool,
static_image_mode: bool):
"""Performs holistic landmark detection with the configured parameters.
Unlike other detectors in this codebase, this doesn't need a model_info parameter
as it uses the built-in legacy Holistic API.
"""
# Initialize detector (without model path)
detector = self.initialize_or_update_detector(None)
# Perform detection with all parameters
holistic_landmarks_batch, segmentation_masks_batch = detector.detect(
image,
min_detection_confidence=min_confidence,
min_tracking_confidence=min_tracking_confidence,
model_complexity=model_complexity,
enable_segmentation=enable_segmentation,
refine_face_landmarks=refine_face_landmarks,
static_image_mode=static_image_mode
)
return (holistic_landmarks_batch, segmentation_masks_batch or [])
def initialize_or_update_detector(self, model_path=None):
"""Override to handle the case where model_path is not needed."""
if self._detector is None:
self._detector = self.DETECTOR_CLASS()
return self._detector
def validate_model_info(self, model_info=None):
"""Override to handle the case where model_info is not needed."""
# No validation needed for legacy API
return None
# Define mappings for ComfyUI
NODE_CLASS_MAPPINGS = {
"MediaPipeHolisticLandmarker": MediaPipeHolisticLandmarkerNode,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"MediaPipeHolisticLandmarker": "**CAUTION** Holistic Landmarker (MediaPipe Legacy)",
}
@@ -0,0 +1,202 @@
import logging
import mediapipe as mp
import numpy as np
import cv2
from typing import List, Dict, Tuple, Any, Optional
import torch
logger = logging.getLogger(__name__)
_category = "Realtime Nodes/MediaPipe Vision/Holistic"
# Use MediaPipe's standard connection definitions instead of defining our own
POSE_CONNECTIONS = mp.solutions.pose.POSE_CONNECTIONS
HAND_CONNECTIONS = mp.solutions.hands.HAND_CONNECTIONS
FACE_CONNECTIONS = mp.solutions.face_mesh.FACEMESH_TESSELATION
class MediaPipeHolisticVisualizerNode:
"""ComfyUI node for visualizing holistic landmarks using direct drawing.
Since holistic landmarks combine multiple landmark types (face, pose, hands),
this doesn't use the BaseLandmarkVisualizationNode but implements a similar approach.
"""
CATEGORY = _category
RETURN_TYPES = ("IMAGE",)
RETURN_NAMES = ("image",)
FUNCTION = "visualize_holistic"
DESCRIPTION = "Visualizes holistic landmarks by drawing face mesh, pose skeleton, and hand landmarks on the input image."
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"image": ("IMAGE", {"tooltip": "Input image on which to visualize landmarks"}),
"landmarks": ("HOLISTIC_LANDMARKS", {"tooltip": "Landmarks from the holistic detector"}),
"face_color": ("STRING", {"default": "#FF3030", "tooltip": "Color for face landmarks"}),
"pose_color": ("STRING", {"default": "#30FF30", "tooltip": "Color for pose landmarks"}),
"left_hand_color": ("STRING", {"default": "#3030FF", "tooltip": "Color for left hand landmarks"}),
"right_hand_color": ("STRING", {"default": "#FF30FF", "tooltip": "Color for right hand landmarks"}),
"draw_face": ("BOOLEAN", {"default": True, "tooltip": "Toggle drawing face landmarks"}),
"draw_pose": ("BOOLEAN", {"default": True, "tooltip": "Toggle drawing pose landmarks"}),
"draw_hands": ("BOOLEAN", {"default": True, "tooltip": "Toggle drawing hand landmarks"}),
"face_point_radius": ("INT", {"default": 1, "min": 1, "max": 5, "tooltip": "Size of face landmark points"}),
"pose_point_radius": ("INT", {"default": 2, "min": 1, "max": 10, "tooltip": "Size of pose landmark points"}),
"hand_point_radius": ("INT", {"default": 2, "min": 1, "max": 10, "tooltip": "Size of hand landmark points"}),
"connection_thickness": ("INT", {"default": 1, "min": 1, "max": 5, "tooltip": "Thickness of connection lines"}),
"min_visibility": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.05,
"tooltip": "Minimum visibility score to draw a pose landmark"})
}
}
def _hex_to_rgb(self, hex_color):
"""Convert hex color string to RGB tuple for OpenCV (BGR format)."""
hex_color = hex_color.lstrip('#')
# Convert to BGR for OpenCV
r, g, b = tuple(int(hex_color[i:i+2], 16) for i in (0, 2, 4))
return (b, g, r) # Return in BGR format for OpenCV
def visualize_holistic(self,
image,
landmarks,
face_color,
pose_color,
left_hand_color,
right_hand_color,
draw_face,
draw_pose,
draw_hands,
face_point_radius,
pose_point_radius,
hand_point_radius,
connection_thickness,
min_visibility):
"""Visualize holistic landmarks using OpenCV for faster performance."""
if not landmarks:
logger.warning("No holistic landmarks provided.")
return (image,)
batch_size = image.shape[0]
result_images = []
# For holistic, each item in the batch is already a single detection
# (unlike other landmarks which have a list of detections per image)
if len(landmarks) != batch_size:
logger.warning(f"Batch size mismatch. Images: {batch_size}, Landmark sets: {len(landmarks)}. Proceeding with available data.")
# Convert color strings to BGR tuples for OpenCV
face_color_bgr = self._hex_to_rgb(face_color)
pose_color_bgr = self._hex_to_rgb(pose_color)
left_hand_color_bgr = self._hex_to_rgb(left_hand_color)
right_hand_color_bgr = self._hex_to_rgb(right_hand_color)
for i in range(batch_size):
# Clone the tensor to avoid modifying the original
img_tensor = image[i].clone()
# Convert tensor to numpy array scaled to 0-255 range for OpenCV
np_image = (img_tensor.cpu().numpy() * 255).astype(np.uint8)
# Convert to BGR for OpenCV processing
cv_image = cv2.cvtColor(np_image, cv2.COLOR_RGB2BGR)
if i < len(landmarks):
holistic_result = landmarks[i] # One holistic result per image
# Draw face landmarks if requested
if draw_face and holistic_result.face_landmarks:
self._draw_landmarks_cv2(
cv_image,
holistic_result.face_landmarks,
FACE_CONNECTIONS,
face_color_bgr, face_point_radius, connection_thickness,
min_visibility=None # Face doesn't use visibility
)
# Draw pose landmarks if requested
if draw_pose and holistic_result.pose_landmarks:
self._draw_landmarks_cv2(
cv_image,
holistic_result.pose_landmarks,
POSE_CONNECTIONS,
pose_color_bgr, pose_point_radius, connection_thickness,
min_visibility=min_visibility
)
# Draw left hand landmarks if requested
if draw_hands and holistic_result.left_hand_landmarks:
self._draw_landmarks_cv2(
cv_image,
holistic_result.left_hand_landmarks,
HAND_CONNECTIONS,
left_hand_color_bgr, hand_point_radius, connection_thickness,
min_visibility=None # Hands don't use visibility
)
# Draw right hand landmarks if requested
if draw_hands and holistic_result.right_hand_landmarks:
self._draw_landmarks_cv2(
cv_image,
holistic_result.right_hand_landmarks,
HAND_CONNECTIONS,
right_hand_color_bgr, hand_point_radius, connection_thickness,
min_visibility=None # Hands don't use visibility
)
# Convert back to RGB and tensor format
cv_image_rgb = cv2.cvtColor(cv_image, cv2.COLOR_BGR2RGB)
result_tensor = torch.from_numpy(cv_image_rgb.astype(np.float32) / 255.0)
result_images.append(result_tensor)
return (torch.stack(result_images),)
def _draw_landmarks_cv2(self,
image: np.ndarray,
landmark_list: List,
connections: List,
color: Tuple[int, int, int],
point_radius: int,
connection_thickness: int,
min_visibility: Optional[float] = None):
"""Helper method to draw landmarks with OpenCV for faster performance."""
height, width = image.shape[:2]
points = {}
# Process landmarks and draw points
for lm in landmark_list:
# Check visibility if required
is_visible = True
if min_visibility is not None and hasattr(lm, 'visibility') and lm.visibility is not None:
is_visible = lm.visibility >= min_visibility
if is_visible:
# Convert normalized coordinates to pixel coordinates
cx, cy = int(lm.x * width), int(lm.y * height)
points[lm.index] = (cx, cy)
# Draw point as a filled circle
cv2.circle(image, (cx, cy), point_radius, color, -1)
# Draw connections
if connections and points:
for connection in connections:
start_idx, end_idx = connection
if start_idx in points and end_idx in points:
cv2.line(
image,
points[start_idx],
points[end_idx],
color,
connection_thickness
)
# Define mappings for ComfyUI
NODE_CLASS_MAPPINGS = {
"MediaPipeHolisticVisualizer": MediaPipeHolisticVisualizerNode,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"MediaPipeHolisticVisualizer": "Holistic Visualizer (MediaPipe)",
}
@@ -0,0 +1 @@
@@ -0,0 +1,78 @@
import torch
import logging
# Import Base Loader and Detector
from ..common.model_loader import MediaPipeModelLoaderBaseNode
from ..common.base_detector_node import BaseMediaPipeDetectorNode
from ....src.mediapipe_vision.image_embedding.detector import ImageEmbedder
logger = logging.getLogger(__name__)
_category = "Realtime Nodes/MediaPipe Vision/ImageEmbedding"
# --- Model Loader ---
class MediaPipeImageEmbedderModelLoaderNode(MediaPipeModelLoaderBaseNode):
"""ComfyUI node for loading MediaPipe Image Embedder models."""
TASK_TYPE = "image_embedder" # Need to add this task to MEDIAPIPE_MODELS
RETURN_TYPES = ("IMAGE_EMBEDDER_MODEL_INFO",)
RETURN_NAMES = ("model_info",)
CATEGORY = _category
# --- Embedder Node ---
class MediaPipeImageEmbedderNode(BaseMediaPipeDetectorNode):
"""ComfyUI node for MediaPipe Image Embedding."""
# Define class variables required by the base class
DETECTOR_CLASS = ImageEmbedder
MODEL_INFO_TYPE = "IMAGE_EMBEDDER_MODEL_INFO"
EXPECTED_TASK_TYPE = "image_embedder"
RETURN_TYPES = ("IMAGE_EMBEDDINGS",)
RETURN_NAMES = ("image_embeddings",)
FUNCTION = "detect"
CATEGORY = _category
@classmethod
def INPUT_TYPES(cls):
# Start with the base inputs from the parent class
inputs = super().INPUT_TYPES()
# Add embedder-specific parameters
inputs["required"].update({
"l2_normalize": ("BOOLEAN", {"default": True,
"tooltip": "Apply L2 normalization to the embedding"}),
"quantize": ("BOOLEAN", {"default": False,
"tooltip": "Output quantized (uint8) embedding instead of float"}),
})
return inputs
def detect(self, image: torch.Tensor, model_info: dict, l2_normalize: bool,
quantize: bool, running_mode: str, delegate: str):
"""Generates image embeddings with the configured parameters."""
# Validate model_info and get model path
model_path = self.validate_model_info(model_info)
# Initialize or update detector
detector = self.initialize_or_update_detector(model_path)
# Generate embeddings with all parameters
embeddings_batch = detector.embed(
image,
l2_normalize=l2_normalize,
quantize=quantize,
running_mode=running_mode,
delegate=delegate
)
return (embeddings_batch,)
# --- Mappings ---
NODE_CLASS_MAPPINGS = {
"MediaPipeImageEmbedderModelLoader": MediaPipeImageEmbedderModelLoaderNode,
"MediaPipeImageEmbedder": MediaPipeImageEmbedderNode,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"MediaPipeImageEmbedderModelLoader": "Load Image Embedder Model (MediaPipe)",
"MediaPipeImageEmbedder": "Image Embedder (MediaPipe)",
}
@@ -0,0 +1,144 @@
import torch
import torch.nn.functional as F
import logging
from typing import List, Optional
from ....src.mediapipe_vision.types import ImageEmbedderResult # Assuming this is the relevant type
logger = logging.getLogger(__name__)
class ReshapeMediaPipeEmbeddingNode:
"""Reshapes MediaPipe embedding for experimental use as ComfyUI Conditioning.
WARNING: Semantically different from standard CLIP conditioning.
Pads/truncates the embedding vector and replicates it across 77 tokens.
"""
CATEGORY = "Realtime Nodes/MediaPipe Vision/ImageEmbedding/_Experimental"
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"image_embeddings": ("IMAGE_EMBEDDINGS",), # List[List[ImageEmbedderResult]]
"target_dimension": ("INT", {"default": 768, "min": 64, "max": 4096, "step": 8,
"tooltip": "Target dimension D for the reshaped (1, 77, D) tensor (e.g., 768 for SD1.5, 1024/1280 for SDXL)"}),
}
}
RETURN_TYPES = ("CONDITIONING",)
FUNCTION = "reshape_embedding"
def reshape_embedding(self, image_embeddings: List[List[ImageEmbedderResult]], target_dimension: int):
if not image_embeddings or not image_embeddings[0]:
raise ValueError("Input image_embeddings are empty.")
# --- Process First Embedding in Batch/Head ---
# Assuming we use the first embedding result from the first image in the batch
first_embedding_result = image_embeddings[0][0]
if first_embedding_result.float_embedding:
embedding_vector = first_embedding_result.float_embedding
elif first_embedding_result.quantized_embedding:
# Note: Using quantized embedding might require dequantization logic
# For now, raise error or warn, as direct reshaping is less meaningful
raise NotImplementedError("Reshaping quantized embeddings is not directly supported. Use float embeddings.")
else:
raise ValueError("No valid embedding found in the input.")
# Convert to tensor
embed_tensor = torch.tensor(embedding_vector, dtype=torch.float32)
orig_dim = embed_tensor.shape[0]
# --- Pad or Truncate ---
if orig_dim < target_dimension:
# Pad with zeros
padding = torch.zeros(target_dimension - orig_dim, dtype=torch.float32)
final_vector = torch.cat((embed_tensor, padding), dim=0)
elif orig_dim > target_dimension:
# Truncate
final_vector = embed_tensor[:target_dimension]
else:
final_vector = embed_tensor
# --- Reshape and Replicate for Conditioning ---
# Reshape to (1, 1, D)
cond_vec = final_vector.unsqueeze(0).unsqueeze(0)
# Replicate to (1, 77, D)
cond_batch = cond_vec.repeat(1, 77, 1)
# --- Create Pooled Output ---
# Use the padded/truncated vector as the pooled output (1, D)
pooled_output = final_vector.unsqueeze(0)
# --- Package as CONDITIONING ---
# Standard ComfyUI conditioning format: List[List[Union[torch.Tensor, Dict[str, torch.Tensor]]]]
conditioning = [[cond_batch, {"pooled_output": pooled_output}]]
logger.warning("Reshaped MediaPipe embedding to CONDITIONING format. Use experimentally.")
return (conditioning,)
# --- Comparison Node ---
class CompareMediaPipeEmbeddingsNode:
"""Calculates the cosine similarity between two MediaPipe image embeddings."""
CATEGORY = "Realtime Nodes/MediaPipe Vision/ImageEmbedding"
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"image_embeddings_1": ("IMAGE_EMBEDDINGS",),
"image_embeddings_2": ("IMAGE_EMBEDDINGS",),
}
}
RETURN_TYPES = ("FLOAT",)
RETURN_NAMES = ("cosine_similarity",)
FUNCTION = "compare_embeddings"
def compare_embeddings(self, image_embeddings_1: List[List[ImageEmbedderResult]],
image_embeddings_2: List[List[ImageEmbedderResult]]):
def _extract_vector(embeddings: List[List[ImageEmbedderResult]]) -> Optional[torch.Tensor]:
"""Helper to extract the first float embedding vector."""
if not embeddings or not embeddings[0]:
return None
first_result = embeddings[0][0]
if first_result.float_embedding:
return torch.tensor(first_result.float_embedding, dtype=torch.float32)
else:
logger.warning("Cannot compare non-float embeddings.")
return None
vec1 = _extract_vector(image_embeddings_1)
vec2 = _extract_vector(image_embeddings_2)
if vec1 is None or vec2 is None:
logger.error("Could not extract valid float embedding vectors for comparison.")
return (0.0,) # Return 0 similarity on error
# Ensure vectors have the same dimension (pad/truncate shorter one)
dim1, dim2 = vec1.shape[0], vec2.shape[0]
if dim1 != dim2:
logger.warning(f"Embedding dimensions differ ({dim1} vs {dim2}). Padding/truncating for comparison.")
max_dim = max(dim1, dim2)
if dim1 < max_dim:
vec1 = F.pad(vec1, (0, max_dim - dim1))
elif dim2 < max_dim:
vec2 = F.pad(vec2, (0, max_dim - dim2))
# Calculate Cosine Similarity
# Normalize vectors before dot product for cosine similarity
similarity = F.cosine_similarity(vec1.unsqueeze(0), vec2.unsqueeze(0), dim=1)
return (similarity.item(),)
# --- Mappings ---
NODE_CLASS_MAPPINGS = {
"ReshapeMediaPipeEmbedding": ReshapeMediaPipeEmbeddingNode,
"CompareMediaPipeEmbeddings": CompareMediaPipeEmbeddingsNode, # Add new node
}
NODE_DISPLAY_NAME_MAPPINGS = {
"ReshapeMediaPipeEmbedding": "Reshape MediaPipe Embedding (Experimental)",
"CompareMediaPipeEmbeddings": "Compare MediaPipe Embeddings", # Add new node name
}
@@ -0,0 +1 @@
@@ -0,0 +1,221 @@
"""Node wrapper for image segmentation"""
import torch
import numpy as np
import cv2 # For visualization
import logging
# Imports from this project
from ....src.mediapipe_vision.image_segmentation.segmenter import ImageSegmenter
# Import Base Classes
from ..common.model_loader import MediaPipeModelLoaderBaseNode
from ..common.base_detector_node import BaseMediaPipeDetectorNode
logger = logging.getLogger(__name__)
_category = "Realtime Nodes/MediaPipe Vision/ImageSegmentation"
# Define class names and their corresponding indices for multiclass models
MULTICLASS_NAMES = {
"Background": 0,
"Hair": 1,
"Body-skin": 2,
"Face-skin": 3,
"Clothes": 4,
"Accessories/Other": 5,
}
# Uses MULTICLASS_NAMES indices
CLASS_COLORS_BGR = {
MULTICLASS_NAMES["Background"]: (0, 0, 0), # Black
MULTICLASS_NAMES["Hair"]: (255, 0, 0), # Blue
MULTICLASS_NAMES["Body-skin"]: (0, 255, 255), # Yellow
MULTICLASS_NAMES["Face-skin"]: (255, 200, 100), # Light Blue
MULTICLASS_NAMES["Clothes"]: (0, 255, 0), # Green
MULTICLASS_NAMES["Accessories/Other"]: (255, 0, 255), # Magenta
}
# --- End Constants ---
# Inherit from the base loader
class MediaPipeImageSegmenterModelLoaderNode(MediaPipeModelLoaderBaseNode):
"""ComfyUI node for loading MediaPipe Image Segmenter models."""
TASK_TYPE = "image_segmenter"
RETURN_TYPES = ("IMAGE_SEGMENTER_MODEL_INFO",)
RETURN_NAMES = ("model_info",)
# INPUT_TYPES and FUNCTION inherited
CATEGORY = _category
class MediaPipeImageSegmenterNode(BaseMediaPipeDetectorNode):
"""ComfyUI node for MediaPipe Image Segmentation, adapted from Stream-Pack.
Manages the ImageSegmenter instance for potential reuse based on config.
"""
# Define class variables required by the base class
DETECTOR_CLASS = ImageSegmenter
MODEL_INFO_TYPE = "IMAGE_SEGMENTER_MODEL_INFO"
EXPECTED_TASK_TYPE = "image_segmenter"
RETURN_TYPES = ("MASK", "IMAGE", "MP_INT_MASK",)
RETURN_NAMES = ("mask", "visualization", "multiclass_segments")
FUNCTION = "detect"
CATEGORY = _category
@classmethod
def INPUT_TYPES(cls):
# Start with the base inputs from the parent class
inputs = super().INPUT_TYPES()
# Add segmentation-specific parameters
inputs["required"].update({
"output_confidence_masks": ("BOOLEAN", {"default": False,
"tooltip": "Output confidence mask (0-1) instead of category mask. Disables multiclass/visualization."}),
"threshold": ("FLOAT", {
"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01,
"tooltip": "Confidence threshold for confidence mask output mode. Pixels below threshold become 0."
}),
"generate_visualization": ("BOOLEAN", {"default": False,
"tooltip": "Generate a colored visualization image (only works in category mask mode)."
}),
})
# Rename the delegate parameter for consistency with original implementation
inputs["required"]["delegate_mode"] = inputs["required"].pop("delegate")
return inputs
# --- Helper methods from Stream-Pack ---
def is_multiclass_model(self, model_info: dict):
# Improved check based on variant or path
model_variant = model_info.get('model_variant', '').lower()
model_path = model_info.get('model_path', '').lower()
# Add known multiclass identifiers
return 'multiclass' in model_variant or 'multiclass' in model_path
def create_visualization(self, category_mask_np, is_multiclass=False):
"""Generates a colored visualization from a category mask."""
if category_mask_np is None:
return None
if not is_multiclass:
# Simple binary visualization (white on black)
vis_image = np.where(category_mask_np > 0, 255, 0).astype(np.uint8)
return cv2.cvtColor(vis_image, cv2.COLOR_GRAY2BGR)
else:
# Multiclass visualization
height, width = category_mask_np.shape
vis_image = np.zeros((height, width, 3), dtype=np.uint8)
for class_name, class_id in MULTICLASS_NAMES.items():
color_bgr = CLASS_COLORS_BGR.get(class_id, (128, 128, 128)) # Default grey
vis_image[category_mask_np == class_id] = color_bgr
return vis_image
# --- End Helper methods ---
def detect(self, image: torch.Tensor, model_info: dict, output_confidence_masks: bool,
threshold: float, generate_visualization: bool, running_mode: str, delegate_mode: str):
"""Performs image segmentation with the configured parameters."""
# Validate model_info and get model path
model_path = self.validate_model_info(model_info)
# Initialize or update detector
detector = self.initialize_or_update_detector(model_path)
# Determine which output type the detector needs
request_confidence = output_confidence_masks
# Need category mask if confidence isn't requested OR if visualization is needed
request_category = not output_confidence_masks or generate_visualization
if not request_confidence and not request_category:
# Should not happen with the logic above, but as a safeguard
raise ValueError("Internal logic error: No mask type determined for segmenter.")
# Call segmenter's segment method
# Returns lists (one element per image in batch)
batch_results_confidence, batch_results_category = detector.segment(
image,
output_confidence_masks=request_confidence,
output_category_mask=request_category,
running_mode=running_mode,
delegate_mode=delegate_mode
)
# Initialize output lists
batch_size = image.shape[0]
h, w = image.shape[1], image.shape[2]
all_primary_masks = []
all_vis_images = []
all_category_masks = []
default_vis_tensor = torch.zeros((h, w, 3), dtype=torch.float32, device=image.device) # HWC
# --- Process results based on node inputs ---
is_multiclass = self.is_multiclass_model(model_info)
# Loop through batch results
for i in range(batch_size):
# Get results for the current image
confidence_masks_hw_list = batch_results_confidence[i] if batch_results_confidence else None
category_mask_hw = batch_results_category[i] if batch_results_category else None
# Initialize outputs for this image
current_primary_mask_hw = torch.zeros((h, w), dtype=torch.float32, device=image.device)
current_vis_image_hwc = default_vis_tensor
current_category_mask_hw = torch.zeros((h, w), dtype=torch.long, device=image.device)
if output_confidence_masks:
# User wants confidence mask as primary output
if confidence_masks_hw_list:
# Take the first confidence mask (often foreground vs background)
conf_mask_hw = confidence_masks_hw_list[0]
# Apply threshold
conf_mask_hw_thresh = torch.where(conf_mask_hw >= threshold, conf_mask_hw, torch.zeros_like(conf_mask_hw))
current_primary_mask_hw = conf_mask_hw_thresh # HW
# Visualization and multiclass output are disabled/invalid in confidence mode
generate_visualization_for_this = False
else:
# User wants category mask (default or for visualization)
if category_mask_hw is not None:
# Determine primary MASK output (HW Float)
if is_multiclass:
# Treat non-background as foreground for the primary mask
primary_mask_hw = (category_mask_hw > 0).float()
else:
# Binary models: Category 1 is foreground
primary_mask_hw = (category_mask_hw > 0).float()
current_primary_mask_hw = primary_mask_hw # HW
# Prepare raw category mask output (HW Long Tensor)
current_category_mask_hw = category_mask_hw # Already Long
# Generate visualization if requested (and category mask exists)
generate_visualization_for_this = generate_visualization
if generate_visualization_for_this:
# Convert to NumPy for OpenCV, then back to tensor
category_np = category_mask_hw.cpu().numpy().astype(np.uint8)
vis_np = self.create_visualization(category_np, is_multiclass)
if vis_np is not None:
# Convert back to tensor (HWC)
vis_tensor = torch.from_numpy(vis_np).to(dtype=torch.float32, device=image.device) / 255.0
current_vis_image_hwc = vis_tensor # HWC
# Add results for this image to the batch results
all_primary_masks.append(current_primary_mask_hw) # HW
all_vis_images.append(current_vis_image_hwc) # HWC
all_category_masks.append(current_category_mask_hw) # HW
# Stack all results to return as tensors
primary_mask_tensor = torch.stack(all_primary_masks, dim=0) # BHW
vis_image_tensor = torch.stack(all_vis_images, dim=0) # BHWC
category_mask_tensor = torch.stack(all_category_masks, dim=0) # BHW
# Ensure multiclass_segments is in the right format (BHW) for downstream nodes
# Not adding an unsqueeze here, as we're handling both formats in the select_segment node
return (primary_mask_tensor, vis_image_tensor, category_mask_tensor)
# Define mappings for ComfyUI
NODE_CLASS_MAPPINGS = {
"MediaPipeImageSegmenterModelLoader": MediaPipeImageSegmenterModelLoaderNode,
"MediaPipeImageSegmenter": MediaPipeImageSegmenterNode,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"MediaPipeImageSegmenterModelLoader": "Load Image Segmenter Model (MediaPipe)",
"MediaPipeImageSegmenter": "Image Segmenter (MediaPipe)",
}
@@ -0,0 +1,87 @@
import torch
import numpy as np
import logging
logger = logging.getLogger(__name__)
_category = "Realtime Nodes/MediaPipe Vision/ImageSegmentation/Masking"
# Define class names and their corresponding indices for multiclass models
MULTICLASS_NAMES = {
"Background": 0,
"Hair": 1,
"Body-skin": 2,
"Face-skin": 3,
"Clothes": 4,
"Accessories/Other": 5,
}
class SelectMediaPipeSegmentNode:
"""
Selects a specific segmentation mask from the raw multiclass output
of the MediaPipeImageSegmenterNode.
"""
@classmethod
def INPUT_TYPES(cls):
available_segments = list(MULTICLASS_NAMES.keys())
return {
"required": {
"multiclass_segments": ("MP_INT_MASK", {
"tooltip": "Raw multiclass segmentation data (INT_MASK) from MediaPipe Image Segmenter node."
}),
"segment_name": (available_segments, {
"default": "Clothes" if "Clothes" in available_segments else available_segments[0],
"tooltip": "Select the specific segment to extract as a mask."
}),
}
}
RETURN_TYPES = ("MASK",)
RETURN_NAMES = ("mask",)
FUNCTION = "select_segment"
CATEGORY = _category
DESCRIPTION = "Extracts a single mask for a selected segment (e.g., Hair, Clothes) from the multiclass segmentation data."
def select_segment(self, multiclass_segments: torch.Tensor, segment_name: str):
if segment_name not in MULTICLASS_NAMES:
raise ValueError(f"Invalid segment_name '{segment_name}'. Available: {list(MULTICLASS_NAMES.keys())}")
if multiclass_segments is None or multiclass_segments.numel() == 0:
logger.warning("Input multiclass_segments is empty or None. Returning zero mask.")
# Need a shape to return a zero mask. How to get expected shape?
# For now, maybe raise error or return None (caller must handle)
# Returning None is safer if shape is unknown.
return (None,) # Or raise error
# Check tensor format and reshape if needed
# Input formats can be either BHW or B1HW (both are common in ComfyUI)
if multiclass_segments.dim() == 3: # BHW format
# Reshape to B1HW format
logger.info(f"Reshaping multiclass_segments from BHW format {multiclass_segments.shape} to B1HW format")
multiclass_segments = multiclass_segments.unsqueeze(1) # Insert channel dimension
logger.info(f"New shape: {multiclass_segments.shape}")
elif multiclass_segments.dim() == 4 and multiclass_segments.shape[1] != 1:
# If it's BCHW but not B1HW, raise error
raise ValueError(f"Expected multiclass_segments to be B1HW or BHW tensor. Got shape {multiclass_segments.shape}")
# Ensure input is a LongTensor
if multiclass_segments.dtype != torch.int64:
multiclass_segments = multiclass_segments.to(torch.int64)
target_id = MULTICLASS_NAMES[segment_name]
selected_mask = (multiclass_segments == target_id).float()
# If result is B1HW, squeeze back to BHW for ComfyUI mask format compatibility
if selected_mask.dim() == 4:
selected_mask = selected_mask.squeeze(1)
return (selected_mask,)
# --- End Select Segment Node ---
# --- Mappings ---
NODE_CLASS_MAPPINGS = {
"SelectMediaPipeSegment": SelectMediaPipeSegmentNode,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"SelectMediaPipeSegment": "Select MediaPipe Segment",
}
@@ -0,0 +1,116 @@
import torch
import logging
import json
# Import Base Loader and Detector
from ..common.model_loader import MediaPipeModelLoaderBaseNode
from ....src.mediapipe_vision.interactive_segmentation.detector import InteractiveSegmenterProcessor
from ....src.mediapipe_vision. types import PointOfInterest
logger = logging.getLogger(__name__)
_category = "Realtime Nodes/MediaPipe Vision/InteractiveSegmentation"
# --- Model Loader ---
class MediaPipeInteractiveSegmenterModelLoaderNode(MediaPipeModelLoaderBaseNode):
"""ComfyUI node for loading MediaPipe Interactive Segmenter models."""
# Usually uses a specific model like selfie_segmenter or a custom one
TASK_TYPE = "interactive_segmenter" # Need to define this task and models in model_loader.py
RETURN_TYPES = ("INTERACTIVE_SEGMENTER_MODEL_INFO",)
RETURN_NAMES = ("model_info",)
CATEGORY = _category
# --- Segmenter Node ---
class MediaPipeInteractiveSegmenterNode:
"""ComfyUI node for MediaPipe Interactive Image Segmentation."""
def __init__(self):
self._processor: InteractiveSegmenterProcessor = None
self._model_path: str = None
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"image": ("IMAGE",),
"model_info": ("INTERACTIVE_SEGMENTER_MODEL_INFO",),
# Input points as a JSON string list of lists/tuples: e.g., "[[0.5, 0.5], [0.2, 0.3]]"
"points_json": ("STRING", {"default": "[[0.5, 0.5]]", "multiline": True,
"tooltip": "Points of interest (normalized coords) as JSON: [[x1,y1], [x2,y2], ...]"}),
"output_category_mask": ("BOOLEAN", {"default": True,
"tooltip": "Output a category mask (0=background, 1=foreground). Priority over confidence."}),
"output_confidence_mask": ("BOOLEAN", {"default": False,
"tooltip": "Output a confidence mask (0.0-1.0). Used if category mask is False."}),
}
}
RETURN_TYPES = ("MASK",) # Output is a standard ComfyUI mask
RETURN_NAMES = ("mask",)
FUNCTION = "segment_interactive"
CATEGORY = _category
def segment_interactive(self, image: torch.Tensor, model_info: dict, points_json: str,
output_category_mask: bool, output_confidence_mask: bool):
"""Performs interactive segmentation."""
task_type = model_info.get('task_type')
expected_task_type = MediaPipeInteractiveSegmenterModelLoaderNode.TASK_TYPE
if not isinstance(model_info, dict) or task_type != expected_task_type:
raise ValueError(f"Invalid model_info. Expected task_type '{expected_task_type}' but 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.")
# Parse points from JSON
try:
points_list = json.loads(points_json)
if not isinstance(points_list, list) or not all(isinstance(p, (list, tuple)) and len(p) == 2 for p in points_list):
raise ValueError("Invalid format")
points_of_interest = [PointOfInterest(x=p[0], y=p[1]) for p in points_list]
except Exception as e:
raise ValueError(f"Invalid points_json format: {e}. Expected list of [x, y] pairs, e.g., [[0.5, 0.5]].")
# Manage detector instance
if self._processor is None or self._model_path != model_path:
if self._processor and hasattr(self._processor, 'close'):
try: self._processor.close()
except Exception as e: logger.warning(f"Error closing InteractiveSegmenterProcessor: {e}")
logger.info(f"Creating new InteractiveSegmenterProcessor instance for {model_path}")
self._processor = InteractiveSegmenterProcessor(model_path)
self._model_path = model_path
# Call the detector's segment method
mask_batch = self._processor.segment(
image,
points_of_interest=points_of_interest,
output_category_mask=output_category_mask,
output_confidence_mask=output_confidence_mask
)
# Convert HW mask list to B1HW tensor for ComfyUI MASK type
# Assuming detector returns one mask per batch item
comfy_masks = []
for mask_hw in mask_batch:
if mask_hw is not None:
comfy_masks.append(mask_hw.unsqueeze(0)) # Add channel dim
else:
# Handle None case - maybe return zero mask? Needs image shape
h, w = image.shape[1], image.shape[2]
comfy_masks.append(torch.zeros((1, h, w), dtype=torch.float32, device=image.device))
output_tensor = torch.stack(comfy_masks, dim=0) # Stack along batch dim
return (output_tensor,)
def __del__(self):
if hasattr(self, '_processor') and self._processor and hasattr(self._processor, 'close'):
try: self._processor.close()
except Exception as e: logger.warning(f"Error closing InteractiveSegmenterProcessor in __del__: {e}")
self._processor = None
# --- Mappings ---
NODE_CLASS_MAPPINGS = {
"MediaPipeInteractiveSegmenterModelLoader": MediaPipeInteractiveSegmenterModelLoaderNode,
"MediaPipeInteractiveSegmenter": MediaPipeInteractiveSegmenterNode,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"MediaPipeInteractiveSegmenterModelLoader": "Load Interactive Segmenter Model (MediaPipe)",
"MediaPipeInteractiveSegmenter": "Interactive Segmenter (MediaPipe)",
}
@@ -0,0 +1 @@
@@ -0,0 +1,74 @@
import torch
from ....src.mediapipe_vision.object_detection.detector import ObjectDetector
from ..common.model_loader import MediaPipeModelLoaderBaseNode
from ..common.base_detector_node import BaseMediaPipeDetectorNode
import logging
logger = logging.getLogger(__name__)
_category = "Realtime Nodes/MediaPipe Vision/ObjectDetection"
class MediaPipeObjectDetectorModelLoaderNode(MediaPipeModelLoaderBaseNode):
"""ComfyUI node for loading MediaPipe Object Detector models."""
TASK_TYPE = "object_detector"
RETURN_TYPES = ("OBJECT_DETECTOR_MODEL_INFO",)
RETURN_NAMES = ("model_info",)
CATEGORY = _category
class MediaPipeObjectDetectorNode(BaseMediaPipeDetectorNode):
"""ComfyUI node for MediaPipe Object Detection."""
# Define class variables required by the base class
DETECTOR_CLASS = ObjectDetector
MODEL_INFO_TYPE = "OBJECT_DETECTOR_MODEL_INFO"
EXPECTED_TASK_TYPE = "object_detector"
RETURN_TYPES = ("OBJECT_DETECTIONS",)
RETURN_NAMES = ("object_detections",)
FUNCTION = "detect"
CATEGORY = _category
@classmethod
def INPUT_TYPES(cls):
# Start with the base inputs from the parent class
inputs = super().INPUT_TYPES()
# Add object detection specific parameters
inputs["required"].update({
"min_confidence": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.05,
"tooltip": "Minimum confidence score for detected objects"}),
"max_results": ("INT", {"default": 5, "min": 1, "max": 50, "step": 1,
"tooltip": "Maximum number of objects to detect"}),
})
return inputs
def detect(self, image: torch.Tensor, model_info: dict, min_confidence: float,
max_results: int, running_mode: str, delegate: str):
"""Performs object detection with the configured parameters."""
# Validate model_info and get model path
model_path = self.validate_model_info(model_info)
# Initialize or update detector
detector = self.initialize_or_update_detector(model_path)
# Perform detection with all parameters
batch_results = detector.detect(
image,
score_threshold=min_confidence,
max_results=max_results,
running_mode=running_mode,
delegate=delegate
)
return (batch_results,)
# Define mappings for ComfyUI
NODE_CLASS_MAPPINGS = {
"MediaPipeObjectDetectorModelLoader": MediaPipeObjectDetectorModelLoaderNode,
"MediaPipeObjectDetector": MediaPipeObjectDetectorNode,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"MediaPipeObjectDetectorModelLoader": "Load Object Detector Model (MediaPipe)",
"MediaPipeObjectDetector": "Object Detector (MediaPipe)",
}
@@ -0,0 +1,106 @@
# visualize_object_detection.py
import torch
import numpy as np
import cv2
from ....src.mediapipe_vision.utils.image_utils import convert_to_cv2, convert_to_tensor
from ..common.base_visualization_nodes import BaseDetectionVisualizationNode
import logging
logger = logging.getLogger(__name__)
_category = "Realtime Nodes/MediaPipe Vision/ObjectDetection/Visualization"
class VisualizeObjectDetections(BaseDetectionVisualizationNode):
"""Draws object detection bounding boxes and labels onto images."""
CATEGORY = _category
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"object_detections": ("OBJECT_DETECTIONS",), # Expect the specific type
"image": ("IMAGE",),
},
"optional": {
"box_color": ("STRING", {"default": "#FF0000"}), # Red
"box_thickness": ("INT", {"default": 2, "min": 1, "max": 10}),
"label_color": ("STRING", {"default": "#FFFFFF"}), # White
"label_bg_color": ("STRING", {"default": "#FF0000"}), # Red BG
# Font size might be tricky to implement well, skipping for now
}
}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "draw_visualizations"
def draw_visualizations(self, object_detections, image,
box_color="#FF0000", box_thickness=2,
label_color="#FFFFFF", label_bg_color="#FF0000"):
"""Main function that processes the batch of images and detections."""
return self.process_batch(
object_detections,
image,
self.draw_objects,
box_color=box_color,
box_thickness=box_thickness,
label_color=label_color,
label_bg_color=label_bg_color
)
def draw_objects(self, image, width, height, detections,
box_color="#FF0000", box_thickness=2,
label_color="#FFFFFF", label_bg_color="#FF0000"):
"""Callback function to draw objects on a single image."""
# Convert hex colors to BGR tuples for cv2
box_color_bgr = self._hex_to_bgr(box_color)
for detection in detections:
bbox = detection.bounding_box
# Top-left corner
x1, y1 = int(bbox.origin_x), int(bbox.origin_y)
# Bottom-right corner
x2, y2 = int(bbox.origin_x + bbox.width), int(bbox.origin_y + bbox.height)
# Draw bounding box
cv2.rectangle(image, (x1, y1), (x2, y2), box_color_bgr, box_thickness)
# Prepare label text (use first category's display name if available)
label = "N/A"
score = 0.0
if detection.categories:
top_category = detection.categories[0]
label = top_category.display_name or top_category.category_name or "N/A"
score = top_category.score
label_text = f"{label}: {score:.2f}"
# Draw label background and text inside the box top-left
try:
self.draw_text_with_background(
image,
label_text,
(x1, y1),
text_color=label_color,
bg_color=label_bg_color
)
except Exception as e:
logger.error(f"Error drawing text label: {e}")
# Fallback: Draw text without background
cv2.putText(
image,
label_text,
(x1 + 2, y1 + 15),
cv2.FONT_HERSHEY_SIMPLEX,
0.5,
self._hex_to_bgr(label_color),
1
)
# Node registration
NODE_CLASS_MAPPINGS = {
"VisualizeObjectDetections": VisualizeObjectDetections
}
NODE_DISPLAY_NAME_MAPPINGS = {
"VisualizeObjectDetections": "Visualize Object Detections (MediaPipe)"
}
@@ -0,0 +1 @@
@@ -0,0 +1,94 @@
import torch
from ....src.mediapipe_vision.pose_landmark.detector import PoseLandmarkDetector
from ..common.model_loader import MediaPipeModelLoaderBaseNode
from ..common.base_detector_node import BaseMediaPipeDetectorNode
import logging
logger = logging.getLogger(__name__)
_category = "Realtime Nodes/MediaPipe Vision/Pose"
class MediaPipePoseLandmarkerModelLoaderNode(MediaPipeModelLoaderBaseNode):
"""ComfyUI node for loading MediaPipe Pose Landmarker models."""
TASK_TYPE = "pose_landmarker"
RETURN_TYPES = ("POSE_LANDMARKER_MODEL_INFO",)
RETURN_NAMES = ("model_info",)
CATEGORY = _category
DESCRIPTION = "Loads a MediaPipe Pose Landmarker model for detecting human body poses and postures in images. Required before using the Pose Landmarker node."
class MediaPipePoseLandmarkerNode(BaseMediaPipeDetectorNode):
"""ComfyUI node for MediaPipe Pose Landmark Detection."""
# Define class variables required by the base class
DETECTOR_CLASS = PoseLandmarkDetector
MODEL_INFO_TYPE = "POSE_LANDMARKER_MODEL_INFO"
EXPECTED_TASK_TYPE = "pose_landmarker"
RETURN_TYPES = ("POSE_LANDMARKS", "MASK_LIST",)
RETURN_NAMES = ("landmarks", "segmentation_masks")
FUNCTION = "detect"
CATEGORY = _category
DESCRIPTION = "Detects human body poses in images and provides precise landmark points for major body joints and posture. Can also generate segmentation masks that isolate people from the background."
@classmethod
def INPUT_TYPES(cls):
# Start with the base inputs from the parent class
inputs = super().INPUT_TYPES()
# Add pose-specific parameters
inputs["required"].update({
"max_results": ("INT", {"default": 1, "min": 1, "max": 4, "step": 1,
"tooltip": "Maximum number of people/poses to detect in the image - higher values can detect more people but use more processing power"}),
"min_confidence": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01,
"tooltip": "Minimum confidence threshold for pose detection - lower values detect more poses but may include false positives"}),
"min_tracking_confidence": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01,
"tooltip": "(VIDEO mode) Minimum confidence threshold for pose tracking between frames - lower values maintain tracking longer but may be less accurate"}),
"min_presence_confidence": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01,
"tooltip": "(VIDEO mode) Minimum confidence threshold for pose presence in the image - lower values may include false negatives"}),
"output_segmentation_masks": ("BOOLEAN", {"default": False,
"tooltip": "When enabled, generates masks that separate people from the background"}),
})
return inputs
def detect(self, image: torch.Tensor, model_info: dict, max_results: int,
min_confidence: float, min_tracking_confidence: float,
min_presence_confidence: float, output_segmentation_masks: bool,
running_mode: str, delegate: str):
"""Performs pose landmark detection with the configured parameters."""
# Validate model_info and get model path
model_path = self.validate_model_info(model_info)
# Initialize or update detector
detector = self.initialize_or_update_detector(model_path)
# Perform detection with all parameters
pose_landmarks_batch, segmentation_masks_batch = detector.detect(
image,
num_poses=max_results,
min_detection_confidence=min_confidence,
min_tracking_confidence=min_tracking_confidence,
min_presence_confidence=min_presence_confidence,
output_segmentation_masks=output_segmentation_masks,
running_mode=running_mode,
delegate=delegate
)
# Process segmentation masks if needed
flat_masks = []
if segmentation_masks_batch:
for img_masks in segmentation_masks_batch:
if img_masks:
flat_masks.extend(img_masks)
return (pose_landmarks_batch, flat_masks if output_segmentation_masks else [])
# Define mappings for ComfyUI
NODE_CLASS_MAPPINGS = {
"MediaPipePoseLandmarkerModelLoader": MediaPipePoseLandmarkerModelLoaderNode,
"MediaPipePoseLandmarker": MediaPipePoseLandmarkerNode,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"MediaPipePoseLandmarkerModelLoader": "Load Pose Landmarker Model (MediaPipe)",
"MediaPipePoseLandmarker": "Pose Landmarker (MediaPipe)",
}
@@ -0,0 +1,115 @@
# node_wrappers/pose/masking.py
import torch
import numpy as np
import cv2
import logging
from typing import List, Dict, Tuple
from ....src.mediapipe_vision.types import PoseLandmarksResult
from ....src.mediapipe_vision.utils.mask_utils import create_mask_from_points
# Removed GetImageDimensionsNode import
logger = logging.getLogger(__name__)
_category = "Realtime Nodes/MediaPipe Vision/Pose/Masking"
# --- Pose Part Definitions (Based on standard MediaPipe Pose Indices - BlazePose) ---
# These indices need verification
POSE_PART_INDICES = {
"HEAD": [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10], # Example
"TORSO": [11, 12, 23, 24], # Example
"LEFT_ARM": [11, 13, 15], # Upper arm only
"RIGHT_ARM": [12, 14, 16], # Upper arm only
"LEFT_FOREARM": [13, 15, 17, 19, 21], # Includes hand points?
"RIGHT_FOREARM": [14, 16, 18, 20, 22], # Includes hand points?
"LEFT_LEG": [23, 25, 27], # Upper leg
"RIGHT_LEG": [24, 26, 28], # Upper leg
"LEFT_LOWER_LEG": [25, 27, 29, 31], # Includes foot?
"RIGHT_LOWER_LEG": [26, 28, 30, 32], # Includes foot?
"FULL_BODY": list(range(33)),
}
# --- Node Definition ---
class MaskFromPoseLandmarksNode:
"""Creates a mask for a specific body part from Pose Landmarks."""
CATEGORY = _category
DESCRIPTION = "Creates precise masks for specific body parts like head, torso, arms, or legs. Perfect for targeted editing or creating effects on specific body regions."
@classmethod
def INPUT_TYPES(cls):
part_names = list(POSE_PART_INDICES.keys())
return {
"required": {
"pose_landmarks": ("POSE_LANDMARKS", {"tooltip": "Pose landmark data from the Pose Landmarker node containing the precise locations of body joints"}),
"image_for_dimensions": ("IMAGE", {"tooltip": "Reference image to determine the size of the output mask - the mask will match these dimensions"}),
"part_name": (part_names, {"default": "FULL_BODY",
"tooltip": "Choose which body part to create a mask for - FULL_BODY (entire person), HEAD, TORSO, LEFT_ARM, etc."}),
"min_visibility": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.05,
"tooltip": "Only include landmarks that are at least this visible (0.0 = include all points regardless of visibility, 1.0 = only fully visible points)"}),
}
}
RETURN_TYPES = ("MASK",)
FUNCTION = "generate_mask"
def generate_mask(self, pose_landmarks: List[List[PoseLandmarksResult]], image_for_dimensions: torch.Tensor, part_name: str, min_visibility: float):
if part_name not in POSE_PART_INDICES:
raise ValueError(f"Unknown pose part name: '{part_name}'. Available: {list(POSE_PART_INDICES.keys())}")
# Get dimensions from image
if image_for_dimensions.dim() != 4:
raise ValueError("image_for_dimensions must be BHWC format.")
height = image_for_dimensions.shape[1]
width = image_for_dimensions.shape[2]
input_batch_size = image_for_dimensions.shape[0]
device = image_for_dimensions.device
# Validate batch sizes match
if len(pose_landmarks) != input_batch_size:
logger.warning(f"Batch size mismatch between pose_landmarks ({len(pose_landmarks)}) and image_for_dimensions ({input_batch_size}). Using landmark batch size.")
batch_size = len(pose_landmarks)
else:
batch_size = input_batch_size
required_indices = POSE_PART_INDICES[part_name]
output_masks = []
for i in range(batch_size):
landmarks_for_image = pose_landmarks[i] # List[PoseLandmarksResult]
combined_points_for_image = []
# Combine points from all detected poses for the selected part
for pose_result in landmarks_for_image:
# Pose result already contains the .landmarks attribute
points_dict = {lm.index: lm for lm in pose_result.landmarks}
part_points = []
for idx in required_indices:
if idx in points_dict:
lm = points_dict[idx]
is_visible = lm.visibility if lm.visibility is not None else 1.0
if is_visible >= min_visibility:
px = int(lm.x * width)
py = int(lm.y * height)
part_points.append((px, py))
# No warning for missing index
combined_points_for_image.extend(part_points)
# Create mask from all collected points for this image
mask_tensor = create_mask_from_points(height, width, combined_points_for_image, device=device)
output_masks.append(mask_tensor)
if not output_masks:
logger.warning("No landmarks found to generate any masks. Returning zero mask batch.")
output_batch = torch.zeros((input_batch_size, 1, height, width), dtype=torch.float32, device=device)
else:
output_batch = torch.stack(output_masks, dim=0)
return (output_batch,)
# --- Mappings ---
NODE_CLASS_MAPPINGS = {
"MaskFromPoseLandmarks": MaskFromPoseLandmarksNode,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"MaskFromPoseLandmarks": "Mask From Pose Landmarks (MediaPipe)",
}
@@ -0,0 +1,39 @@
# visualize_pose.py
import mediapipe as mp
from typing import Optional, List, Any
from ....src.mediapipe_vision.types import PoseLandmarksResult # Import specific result type if needed for extraction
from ..common.base_visualization_nodes import BaseLandmarkVisualizationNode
import logging
logger = logging.getLogger(__name__)
_category = "Realtime Nodes/MediaPipe Vision/Pose/Visualization"
# Define standard pose connections
POSE_CONNECTIONS = mp.solutions.pose.POSE_CONNECTIONS
class VisualizePoseLandmarks(BaseLandmarkVisualizationNode):
"""Draws Pose landmarks and connections onto images."""
# Override base class properties
LANDMARKS_TYPE = "POSE_LANDMARKS"
CONNECTIONS = POSE_CONNECTIONS
DEFAULT_DRAW_POINTS = True
DEFAULT_POINT_COLOR = "#00FFFF" # Cyan
DEFAULT_POINT_RADIUS = 3
DEFAULT_CONNECTION_COLOR = "#FFFFFF" # White
DEFAULT_CONNECTION_THICKNESS = 2
SUPPORTS_VISIBILITY = True # Pose landmarks have visibility
CATEGORY = _category
def get_landmark_list(self, detection_result: Any) -> Optional[List[Any]]:
return detection_result.landmarks
# Node registration
NODE_CLASS_MAPPINGS = {
"VisualizePoseLandmarks": VisualizePoseLandmarks
}
NODE_DISPLAY_NAME_MAPPINGS = {
"VisualizePoseLandmarks": "Visualize Pose Landmarks (MediaPipe)"
}
@@ -0,0 +1,41 @@
import logging
from typing import List, Tuple, Any, Optional
# Import the base class
from ..common.universal_landmark_nodes import LandmarkPositionBaseNode
# Import the specific landmark type for clarity and potential future use
from ....src.mediapipe_vision.types import POSE_LANDMARKS
logger = logging.getLogger(__name__)
_category = "Realtime Nodes/MediaPipe Vision/Pose/PoseLandmark/Position"
class PoseLandmarkPositionNode(LandmarkPositionBaseNode):
"""Extracts position lists (x, y, z, vis, pres) for a specific landmark index
from POSE_LANDMARKS across the batch."""
CATEGORY = _category
# Define the specific input type string this node expects
LANDMARKS_TYPE = "POSE_LANDMARKS"
DESCRIPTION = "Extracts lists of coordinates (x, y, z) and properties (visibility, presence) for a specific Pose Landmark index across an entire batch. Supports world coordinates. Outputs separate lists for each property."
@classmethod
def INPUT_TYPES(cls):
# Define the specific input type for this node
return {
"required": {
"landmarks": (cls.LANDMARKS_TYPE, {"forceInput": True}),
"landmark_index": ("INT", {"default": 0, "min": 0, "max": 32, "tooltip": "Nose=0, Left Wrist=15, Right Wrist=16"}), # Pose specific default/range
"use_world_coordinates": ("BOOLEAN", {"default": False, "tooltip": "Use world coordinates if available"}),
}
}
# No need to override INPUT_TYPES or extract_positions, inheriting from base is sufficient.
# The base class handles everything using self.LANDMARKS_TYPE.
# --- Node Mappings ---
NODE_CLASS_MAPPINGS = {
"PoseLandmarkPosition": PoseLandmarkPositionNode,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"PoseLandmarkPosition": "Pose Landmark Position Extractor (Batch)",
}
@@ -48,7 +48,7 @@ class RepulsiveMaskNode(ControlNodeBase, MaskControlMixin):
RETURN_TYPES = ("MASK",)
FUNCTION = "update"
CATEGORY = "real-time/control/mask"
CATEGORY = "Realtime Nodes/control/mask"
def interact(self, mask, input_mask_np, state, strength, height, width, is_attract):
if np.max(input_mask_np) < 0.1:
@@ -168,7 +168,7 @@ class ResizeMaskNode(ControlNodeBase, MaskControlMixin):
RETURN_TYPES = ("MASK",)
FUNCTION = "update"
CATEGORY = "real-time/control/mask"
CATEGORY = "Realtime Nodes/control/mask"
def resize(self, mask, input_mask_np, state, strength, height, width):
if np.max(input_mask_np) < 0.1:
@@ -40,7 +40,7 @@ class ROINode(ControlNodeBase):
RETURN_TYPES = ("ROI",)
FUNCTION = "define_roi"
CATEGORY = "real-time/control/motion"
CATEGORY = "Realtime Nodes/control/motion"
def define_roi(self, mask, action, value, next_roi=None):
# Convert mask to numpy for bounding box calculation
@@ -100,7 +100,7 @@ class MotionController(ControlNodeBase):
RETURN_TYPES = ("FLOAT", "MASK")
FUNCTION = "process_motion"
CATEGORY = "real-time/control/motion"
CATEGORY = "Realtime Nodes/control/motion"
def process_motion(self, image, roi_chain, threshold, blur_size,
minimum_value, maximum_value, starting_value, always_execute=True):
@@ -252,7 +252,7 @@ class IntegerMotionController(MotionController):
RETURN_TYPES = ("INT", "MASK")
FUNCTION = "process_motion"
CATEGORY = "real-time/control/motion"
CATEGORY = "Realtime Nodes/control/motion"
def _apply_rounding(self, value, mode):
"""Apply the specified rounding mode to convert float to int"""
@@ -65,7 +65,7 @@ class FloatSequence(SequenceControlBase):
RETURN_TYPES = ("FLOAT",)
FUNCTION = "update"
CATEGORY = "real-time/control/sequence"
CATEGORY = "Realtime Nodes/control/sequence"
@classmethod
def INPUT_TYPES(cls):
@@ -88,7 +88,7 @@ class IntSequence(SequenceControlBase):
RETURN_TYPES = ("INT",)
FUNCTION = "update"
CATEGORY = "real-time/control/sequence"
CATEGORY = "Realtime Nodes/control/sequence"
@classmethod
def INPUT_TYPES(cls):
@@ -111,7 +111,7 @@ class StringSequence(SequenceControlBase):
RETURN_TYPES = ("STRING",)
FUNCTION = "update"
CATEGORY = "real-time/control/sequence"
CATEGORY = "Realtime Nodes/control/sequence"
@classmethod
def INPUT_TYPES(cls):
@@ -18,7 +18,7 @@ class StateResetNode(ControlNodeBase):
RETURN_TYPES = ("BOOLEAN",)
FUNCTION = "update"
CATEGORY = "real-time/control/utility"
CATEGORY = "Realtime Nodes/control/utility"
def update(self, trigger, always_execute=True):
print(f"\n=== StateResetNode UPDATE - node_id: {self.node_id} ===")
@@ -49,7 +49,7 @@ class StateTestNode(ControlNodeBase):
RETURN_TYPES = ("INT",)
FUNCTION = "update"
CATEGORY = "real-time/control/utility"
CATEGORY = "Realtime Nodes/control/utility"
def update(self, increment, always_execute=True):
print(f"\n=== StateTestNode UPDATE - node_id: {self.node_id} ===")
@@ -85,7 +85,7 @@ class FloatControl(ValueControlBase):
RETURN_TYPES = ("FLOAT",)
FUNCTION = "update_value"
CATEGORY = "real-time/control/value"
CATEGORY = "Realtime Nodes/control/value"
def update_value(self, maximum_value, minimum_value, starting_value, steps_per_cycle, movement_type, always_execute=True):
return self.update_value_base(maximum_value, minimum_value, starting_value, steps_per_cycle, movement_type, always_execute)
@@ -124,7 +124,7 @@ class IntControl(ValueControlBase):
RETURN_TYPES = ("INT",)
FUNCTION = "update_value"
CATEGORY = "real-time/control/value"
CATEGORY = "Realtime Nodes/control/value"
def update_value(self, maximum_value, minimum_value, starting_value, steps_per_cycle, movement_type, always_execute=True):
result = self.update_value_base(maximum_value, minimum_value, starting_value, steps_per_cycle, movement_type, always_execute)
@@ -159,7 +159,7 @@ class StringControl(ControlNodeBase):
RETURN_TYPES = ("STRING",)
FUNCTION = "update"
CATEGORY = "real-time/control/value"
CATEGORY = "Realtime Nodes/control/value"
def update(self, strings, steps_per_cycle, movement_type, always_execute=True):
# Split the input strings into a list
+2 -2
View File
@@ -28,7 +28,7 @@ class LazyCondition(ControlNodeBase):
RETURN_TYPES = (AlwaysEqualProxy("*"),)
FUNCTION = "update"
CATEGORY = "real-time/control/utility"
CATEGORY = "Realtime Nodes/control/utility"
def check_lazy_status(self, condition, if_true, fallback, use_fallback):
"""Only evaluate the if_true path if condition is truthy."""
@@ -84,7 +84,7 @@ class LogicOperator(ControlNodeBase):
RETURN_TYPES = ("BOOLEAN",)
FUNCTION = "update"
CATEGORY = "real-time/control/logic"
CATEGORY = "Realtime Nodes/control/logic"
def update(self, operation, input_a, always_execute=True, input_b=None):
"""Perform the selected logical operation on inputs based on their truthiness."""
@@ -1 +0,0 @@
# This file intentionally left blank
@@ -16,7 +16,7 @@ class FPSMonitor(ControlNodeBase):
RETURN_TYPES = ("IMAGE", "MASK")
FUNCTION = "update"
CATEGORY = "real-time/control/utility"
CATEGORY = "Realtime Nodes/control/utility"
@classmethod
def INPUT_TYPES(cls):
@@ -162,7 +162,7 @@ class SimilarityFilter(ControlNodeBase):
RETURN_TYPES = ("IMAGE", "BOOLEAN")
RETURN_NAMES = ("image", "should_execute")
FUNCTION = "update"
CATEGORY = "real-time/control/utility"
CATEGORY = "Realtime Nodes/control/utility"
def __init__(self):
super().__init__()
+6 -2
View File
@@ -1,3 +1,7 @@
opencv-python
Pillow
mediapipe
mediapipe>=0.10.9
opencv-python>=4.7.0
numpy>=1.20.0
scipy>=1.8.0
#TODO remove Pillow req
+4
View File
@@ -0,0 +1,4 @@
"""Source code for MediaPipe Vision tasks implementation."""
# Import utility modules
from . import utils
@@ -0,0 +1 @@
"""Face Detection implementation using MediaPipe."""
@@ -0,0 +1,131 @@
"""Face Detection implementation using MediaPipe.
This module contains the implementation of MediaPipe Face Detection functionality.
"""
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 typing import List, Tuple
# Import new types
from ..types import FaceDetectionResult, BoundingBox, FaceKeypoint
from ..utils.timestamp_provider import TimestampProvider # Import from utils
class FaceDetector:
"""Detects faces in an image using MediaPipe FaceDetector."""
def __init__(self, model_path: str):
"""Initialize the detector with the model path.
Args:
model_path: Path to the MediaPipe FaceDetector .task 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 # 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(
base_options=base_options,
running_mode=mode_enum, # Use enum mode
min_detection_confidence=min_detection_confidence
)
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
return vision.FaceDetector.create_from_options(options)
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
@@ -0,0 +1 @@
"""Face Landmark Detection implementation using MediaPipe."""
@@ -0,0 +1,133 @@
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 typing import List, Tuple, Optional, Dict
# Import new types
from ..types import LandmarkPoint, FaceLandmarksResult
from ..utils.timestamp_provider import TimestampProvider # Import TimestampProvider
class FaceLandmarkDetector:
"""Detects face landmarks and blendshapes using MediaPipe FaceLandmarker."""
def __init__(self, model_path: str):
"""Initialize the detector with the model path.
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
self._current_options = None # Store current options for comparison
self._timestamp_provider = None # Added
def _initialize_detector(self, num_faces: int, min_detection_confidence: float,
min_presence_confidence: float, min_tracking_confidence: float,
output_blendshapes: bool, output_transform_matrix: bool,
running_mode: str, delegate: str): # Added running_mode
"""Initializes the FaceLandmarker 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.FaceLandmarkerOptions(
base_options=base_options,
running_mode=mode_enum, # Use enum mode
num_faces=num_faces,
min_face_detection_confidence=min_detection_confidence,
# min_face_presence_confidence and min_tracking_confidence are for VIDEO/LIVE_STREAM
min_face_presence_confidence=min_presence_confidence,
min_tracking_confidence=min_tracking_confidence,
output_face_blendshapes=output_blendshapes,
output_facial_transformation_matrixes=output_transform_matrix
)
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
return vision.FaceLandmarker.create_from_options(options)
def detect(self, image: torch.Tensor, num_faces: int = 1, min_detection_confidence: float = 0.5,
min_presence_confidence: float = 0.5, min_tracking_confidence: float = 0.5,
output_blendshapes: bool = False, 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
"""Detects face landmarks in the input image tensor.
Returns:
A tuple containing:
- face_landmarks_batch: List (batch) of lists (faces) of LandmarkPoint objects.
- blendshapes: List of dictionaries mapping blendshape names to scores (if requested).
- transform_matrices: List of transformation matrices (NumPy arrays) (if requested).
"""
if image.dim() != 4:
raise ValueError("Input tensor must be in BHWC format.")
# Check if options changed
# Add running_mode to tuple check
new_options_tuple = (num_faces, min_detection_confidence, min_presence_confidence, min_tracking_confidence,
output_blendshapes, output_transform_matrix, running_mode, delegate.lower())
current_options_tuple = None
if self._current_options:
current_options_tuple = (
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,
"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'): # Close previous instance if options changed
self._detector_instance.close()
# Pass new params to initializer
self._detector_instance = self._initialize_detector(num_faces, min_detection_confidence,
min_presence_confidence, min_tracking_confidence,
output_blendshapes, output_transform_matrix,
running_mode, delegate) # Pass running_mode
batch_size = image.shape[0]
batch_results_landmarks = []
batch_results_blendshapes = [] if output_blendshapes else None
batch_results_matrices = [] if output_transform_matrix 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]
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_blendshapes = [] if output_blendshapes else None
current_image_matrices = [] if output_transform_matrix else None
if detection_result and detection_result.face_landmarks: # Added check for detection_result
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 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]}
current_image_blendshapes.append(blendshapes)
if output_transform_matrix and detection_result.facial_transformation_matrixes and face_idx < len(detection_result.facial_transformation_matrixes):
current_image_matrices.append(detection_result.facial_transformation_matrixes[face_idx])
batch_results_landmarks.append(current_image_landmarks)
if output_blendshapes: batch_results_blendshapes.append(current_image_blendshapes)
if output_transform_matrix: batch_results_matrices.append(current_image_matrices)
# Close detector after batch? Landmarker might need closing per batch? Check docs.
# For IMAGE mode, closing after batch seems okay.
# self._detector_instance.close() # Close is handled if re-initialized
return batch_results_landmarks, batch_results_blendshapes, batch_results_matrices
@@ -0,0 +1 @@
"""Face Stylization implementation using MediaPipe."""
@@ -0,0 +1,67 @@
import mediapipe as mp
import numpy as np
import torch
from mediapipe.tasks import python
from mediapipe.tasks.python import vision
from typing import List, Optional
from ..utils.timestamp_provider import TimestampProvider
class FaceStylizer:
"""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.current_options = None
self._timestamp_provider = None
def _initialize_detector(self):
"""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 = vision.FaceStylizerOptions(
base_options=base_options
# running_mode=vision.RunningMode.VIDEO # Not supported
)
self.current_options = ()
self._timestamp_provider = TimestampProvider()
return vision.FaceStylizer.create_from_options(options)
def stylize(self, image: torch.Tensor) -> torch.Tensor: # Output is IMAGE
"""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()
batch_size = image.shape[0]
batch_results = []
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)
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
output_image_tensor = img_tensor
batch_results.append(output_image_tensor)
output_batch_tensor = torch.stack(batch_results, dim=0)
return output_batch_tensor
@@ -0,0 +1 @@
"""Gesture Recognition implementation using MediaPipe."""
@@ -0,0 +1,102 @@
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 typing import List
# Import new types
from ..types import GestureRecognitionResult, GestureCategory
from ..utils.timestamp_provider import TimestampProvider # Import TimestampProvider
class GestureRecognizer:
"""Recognizes hand gestures using MediaPipe GestureRecognizer."""
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, 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(
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)
)
self._current_options = options
# Only create timestamp provider for VIDEO mode
self._timestamp_provider = TimestampProvider() if mode_enum == vision.RunningMode.VIDEO else None
return vision.GestureRecognizer.create_from_options(options)
def recognize(self, image: torch.Tensor, num_hands: int = 2, min_detection_confidence: float = 0.5,
min_tracking_confidence: float = 0.5, min_presence_confidence: float = 0.5,
running_mode: str = "video", delegate: str = 'cpu') -> List[List[GestureRecognitionResult]]: # Added running_mode
"""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())
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
@@ -0,0 +1 @@
"""Hand Landmark Detection implementation using MediaPipe."""
@@ -0,0 +1,131 @@
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 typing import List, Tuple, Optional
# Import new types
from ..types import LandmarkPoint, HandLandmarksResult
from ..utils.timestamp_provider import TimestampProvider # Import TimestampProvider
class HandLandmarkDetector:
"""Detects hand landmarks using MediaPipe HandLandmarker."""
def __init__(self, model_path: str):
"""Initialize the detector with the model path.
Args:
model_path: Path to the MediaPipe HandLandmarker .task file.
"""
if not model_path:
raise ValueError("A valid model_path must be provided.")
# Verify file exists? Optional, MediaPipe might handle this.
self.model_path = model_path
self._detector_instance = None
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(
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
)
self._current_options = options
# Only create timestamp provider for VIDEO mode
self._timestamp_provider = TimestampProvider() if mode_enum == vision.RunningMode.VIDEO else None
return vision.HandLandmarker.create_from_options(options)
def detect(self, image: torch.Tensor, num_hands: int = 2, min_detection_confidence: float = 0.5,
min_presence_confidence: float = 0.5, 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())
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
@@ -0,0 +1,3 @@
"""Holistic landmark detection module using MediaPipe Legacy API."""
from .detector import HolisticLandmarkDetector
@@ -0,0 +1,241 @@
"""Holistic landmark detector using MediaPipe Legacy Solutions API.
This implementation specifically handles the MediaPipe Holistic solution
which is currently only available through the legacy API, not the Task API.
"""
import mediapipe as mp
import numpy as np
import torch
import cv2
from typing import List, Tuple, Optional, Dict, Any
# Import needed types
from ..types import LandmarkPoint, HolisticLandmarksResult
from ..utils.timestamp_provider import TimestampProvider
class HolisticLandmarkDetector:
"""Detects holistic landmarks combining face, pose and hand tracking.
This detector uses the legacy MediaPipe Solutions API rather than the Task API
since Holistic is not yet available in the Task API format.
"""
def __init__(self, model_path: Optional[str] = None):
"""Initialize the detector.
Args:
model_path: Unused parameter to maintain API compatibility with other detectors.
The holistic solution doesn't use a model path as it uses the built-in models.
"""
# Model path is ignored for legacy API - it's for API compatibility only
self._detector_instance = None
self._current_options = None
self._timestamp_provider = TimestampProvider()
def _initialize_detector(self,
min_detection_confidence: float,
min_tracking_confidence: float,
model_complexity: int,
enable_segmentation: bool,
refine_face_landmarks: bool,
static_image_mode: bool) -> mp.solutions.holistic.Holistic:
"""Initializes the Holistic detector with the appropriate parameters.
Args:
min_detection_confidence: Minimum confidence for detection to be considered
min_tracking_confidence: Minimum confidence to maintain tracking
model_complexity: 0, 1, or 2 determining model complexity
enable_segmentation: Whether to generate segmentation masks
refine_face_landmarks: Whether to use refined face landmarks (more details around eyes/lips)
static_image_mode: If True, treats input as independent images (slower but more accurate)
Returns:
Initialized Holistic detector instance
"""
# Store options as a dictionary for parameter comparison
self._current_options = {
'min_detection_confidence': min_detection_confidence,
'min_tracking_confidence': min_tracking_confidence,
'model_complexity': model_complexity,
'enable_segmentation': enable_segmentation,
'refine_face_landmarks': refine_face_landmarks,
'static_image_mode': static_image_mode
}
# Create the holistic solution instance with specified options
return mp.solutions.holistic.Holistic(
static_image_mode=static_image_mode,
model_complexity=model_complexity,
enable_segmentation=enable_segmentation,
refine_face_landmarks=refine_face_landmarks,
min_detection_confidence=min_detection_confidence,
min_tracking_confidence=min_tracking_confidence
)
def detect(self, image: torch.Tensor,
min_detection_confidence: float = 0.5,
min_tracking_confidence: float = 0.5,
model_complexity: int = 1,
enable_segmentation: bool = False,
refine_face_landmarks: bool = False,
static_image_mode: bool = False) -> Tuple[List[HolisticLandmarksResult], Optional[List[torch.Tensor]]]:
"""Detects holistic landmarks in the input image tensor.
Args:
image: Input image tensor in BHWC format
min_detection_confidence: Minimum confidence for detection (0.0-1.0)
min_tracking_confidence: Minimum confidence for tracking (0.0-1.0)
model_complexity: Complexity of the pose landmark model (0, 1, or 2)
enable_segmentation: Whether to output segmentation masks
refine_face_landmarks: Whether to use refined face landmarks
static_image_mode: If True, treats input as independent images
Returns:
Tuple containing:
- holistic_landmarks_batch: List of HolisticLandmarksResult objects for each image
- segmentation_masks_batch: List of segmentation masks (if requested)
"""
if image.dim() != 4:
raise ValueError("Input tensor must be in BHWC format.")
# Create a parameters tuple for comparison
new_options = {
'min_detection_confidence': min_detection_confidence,
'min_tracking_confidence': min_tracking_confidence,
'model_complexity': model_complexity,
'enable_segmentation': enable_segmentation,
'refine_face_landmarks': refine_face_landmarks,
'static_image_mode': static_image_mode
}
# Initialize or update detector if needed
if self._detector_instance is None or self._current_options != new_options:
if self._detector_instance and hasattr(self._detector_instance, 'close'):
self._detector_instance.close()
self._detector_instance = self._initialize_detector(
min_detection_confidence,
min_tracking_confidence,
model_complexity,
enable_segmentation,
refine_face_landmarks,
static_image_mode
)
batch_size = image.shape[0]
batch_holistic_landmarks = []
batch_segmentation_masks = [] if enable_segmentation else None
for i in range(batch_size):
img_tensor = image[i] # HWC
# Convert to numpy with proper scaling and datatype
np_image = (img_tensor.cpu().numpy() * 255).astype(np.uint8)
# Process the image with MediaPipe Holistic
holistic_results = self._detector_instance.process(
cv2.cvtColor(np_image, cv2.COLOR_BGR2RGB)
)
# Convert results to our standardized format
holistic_result = self._convert_to_landmark_result(holistic_results)
batch_holistic_landmarks.append(holistic_result)
# Handle segmentation mask if requested
if enable_segmentation and holistic_results.segmentation_mask is not None:
mask_np = holistic_results.segmentation_mask
mask_tensor = torch.from_numpy(mask_np).float() # HW
if batch_segmentation_masks is not None:
batch_segmentation_masks.append(mask_tensor)
return batch_holistic_landmarks, batch_segmentation_masks
def _convert_to_landmark_result(self, holistic_results: Any) -> HolisticLandmarksResult:
"""Converts MediaPipe Holistic results to standardized HolisticLandmarksResult.
Args:
holistic_results: Raw results from mp.solutions.holistic
Returns:
Standardized HolisticLandmarksResult containing all landmarks
"""
# Convert face landmarks
face_landmarks = None
if holistic_results.face_landmarks:
face_landmarks = [
LandmarkPoint(
index=idx,
x=landmark.x,
y=landmark.y,
z=landmark.z,
visibility=getattr(landmark, 'visibility', None)
)
for idx, landmark in enumerate(holistic_results.face_landmarks.landmark)
]
# Convert pose landmarks
pose_landmarks = None
if holistic_results.pose_landmarks:
pose_landmarks = [
LandmarkPoint(
index=idx,
x=landmark.x,
y=landmark.y,
z=landmark.z,
visibility=landmark.visibility
)
for idx, landmark in enumerate(holistic_results.pose_landmarks.landmark)
]
# Convert pose world landmarks
pose_world_landmarks = None
if holistic_results.pose_world_landmarks:
pose_world_landmarks = [
LandmarkPoint(
index=idx,
x=landmark.x,
y=landmark.y,
z=landmark.z,
visibility=landmark.visibility
)
for idx, landmark in enumerate(holistic_results.pose_world_landmarks.landmark)
]
# Convert left hand landmarks
left_hand_landmarks = None
if holistic_results.left_hand_landmarks:
left_hand_landmarks = [
LandmarkPoint(
index=idx,
x=landmark.x,
y=landmark.y,
z=landmark.z
)
for idx, landmark in enumerate(holistic_results.left_hand_landmarks.landmark)
]
# Convert right hand landmarks
right_hand_landmarks = None
if holistic_results.right_hand_landmarks:
right_hand_landmarks = [
LandmarkPoint(
index=idx,
x=landmark.x,
y=landmark.y,
z=landmark.z
)
for idx, landmark in enumerate(holistic_results.right_hand_landmarks.landmark)
]
return HolisticLandmarksResult(
face_landmarks=face_landmarks,
pose_landmarks=pose_landmarks,
pose_world_landmarks=pose_world_landmarks,
left_hand_landmarks=left_hand_landmarks,
right_hand_landmarks=right_hand_landmarks
)
def close(self):
"""Releases resources used by the detector."""
if self._detector_instance and hasattr(self._detector_instance, 'close'):
self._detector_instance.close()
self._detector_instance = None
@@ -0,0 +1 @@
"""Image Classification implementation using MediaPipe."""
@@ -0,0 +1 @@
"""Image Embedding implementation using MediaPipe."""
@@ -0,0 +1,86 @@
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 typing import List, Optional
# Import new types
from ..types import ImageEmbedderResult
from ..utils.timestamp_provider import TimestampProvider # Import TimestampProvider
class ImageEmbedder:
"""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(
base_options=base_options,
running_mode=mode_enum, # Use enum mode
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
return vision.ImageEmbedder.create_from_options(options)
def embed(self, image: torch.Tensor, l2_normalize: bool = True, 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
@@ -0,0 +1 @@
"""Image Segmentation implementation using MediaPipe."""
@@ -0,0 +1,194 @@
"""Image Segmentation implementation using MediaPipe.
This module contains the implementation of MediaPipe Image Segmentation functionality.
"""
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
import platform # Needed for OS check
import logging
from ..utils.timestamp_provider import TimestampProvider # Import TimestampProvider
logger = logging.getLogger(__name__)
class ImageSegmenter:
"""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.
Args:
model_path: Path to the MediaPipe ImageSegmenter .tflite or .task file.
"""
if not model_path:
raise ValueError("A valid model_path must be provided.")
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.
"""
# Return type removed from signature as it's handled internally
new_config = (output_confidence_masks, output_category_mask, running_mode, delegate_mode)
# 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
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}")
# 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"
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)
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.")
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
# No explicit return needed, instance is stored in self._detector_instance
def segment(self, image: torch.Tensor, output_confidence_masks: bool = True, output_category_mask: bool = False,
running_mode: str = "video", delegate_mode="cpu"):
"""Segments the input image tensor.
Args:
image: Input image tensor (BHWC, RGB, 0-1 float).
NOTE: This implementation currently only supports batch size 1 for simplicity.
output_confidence_masks: Whether to output confidence masks (float, 0-1).
output_category_mask: Whether to output a category mask (uint8).
running_mode: Processing mode ("image" or "video").
delegate_mode: Computation delegate ('cpu' or 'gpu').
Returns:
A tuple containing:
- confidence_masks_batch: List (batch) of lists (mask per category) of confidence masks (Torch tensors, HW) or None.
- category_mask_batch: List (batch) of category masks (Torch tensor, HW) or None.
"""
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
def close(self):
"""Closes the segmenter instance if it exists."""
if self._detector_instance:
logger.info("Closing ImageSegmenter instance.")
try:
self._detector_instance.close()
except Exception as e:
logger.warning(f"Error closing segmenter instance: {e}")
self._detector_instance = None
self.current_config = None
self.delegate = None
def __del__(self):
# Ensure resources are released when the object is destroyed
self.close()
@@ -0,0 +1 @@
@@ -0,0 +1,113 @@
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.interactive_segmenter import (
InteractiveSegmenter,
)
from mediapipe.tasks.python.vision import interactive_segmenter
from mediapipe.tasks.python.components.containers import keypoint
from typing import List, Optional
from ..types import PointOfInterest
from ..utils.timestamp_provider import TimestampProvider
import logging
logger = logging.getLogger(__name__)
NormalizedKeypoint = keypoint.NormalizedKeypoint
class InteractiveSegmenterProcessor:
"""Performs interactive image segmentation using MediaPipe."""
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
def _initialize_processor(self, output_category_mask: bool, output_confidence_mask: bool):
base_options = python.BaseOptions(model_asset_path=self.model_path)
options = interactive_segmenter.InteractiveSegmenterOptions(
base_options=base_options,
output_category_mask=output_category_mask,
output_confidence_masks=output_confidence_mask
)
self.current_options = (output_category_mask, output_confidence_mask)
self._timestamp_provider = TimestampProvider()
return interactive_segmenter.InteractiveSegmenter.create_from_options(options)
def segment(self, image: torch.Tensor, points_of_interest: List[PointOfInterest],
output_category_mask: bool = True, output_confidence_mask: bool = False) -> List[Optional[torch.Tensor]]:
"""Performs interactive segmentation based on points of interest.
Iterates through points, generates mask for each, and combines them.
Args:
image: Input image tensor (BHWC).
points_of_interest: List of PointOfInterest objects (normalized coords).
output_category_mask: If True, outputs a category mask.
output_confidence_mask: If True, outputs confidence masks.
Returns:
A list (batch) containing a single mask tensor (HW) or None per image.
Priority is given to category mask if both outputs requested.
"""
if image.dim() != 4:
raise ValueError("Input tensor must be in BHWC format.")
if not points_of_interest:
raise ValueError("At least one point_of_interest must be provided.")
batch_size = image.shape[0]
batch_results_combined = []
new_options = (output_category_mask, output_confidence_mask)
if self.detector_instance is None or self.current_options != new_options:
if self.detector_instance and hasattr(self.detector_instance, 'close'):
self.detector_instance.close()
self.detector_instance = self._initialize_processor(output_category_mask, output_confidence_mask)
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)
individual_masks_for_image = []
for point in points_of_interest:
roi_keypoint_single = NormalizedKeypoint(x=point.x, y=point.y, label=point.label)
roi = interactive_segmenter.RegionOfInterest(
format=interactive_segmenter.RegionOfInterest.Format.KEYPOINT,
keypoint=roi_keypoint_single
)
segmentation_result = self.detector_instance.segment(mp_image, roi)
current_point_mask = None
if output_category_mask and segmentation_result.category_mask:
mask_np = segmentation_result.category_mask.numpy_view()
current_point_mask = torch.from_numpy(mask_np).float()
elif output_confidence_mask and segmentation_result.confidence_masks:
if segmentation_result.confidence_masks:
mask_np = segmentation_result.confidence_masks[0].numpy_view()
current_point_mask = torch.from_numpy(mask_np).float()
if current_point_mask is not None:
individual_masks_for_image.append(current_point_mask)
else:
logger.warning(f"Segmentation failed for point {point} on image {i}. Skipping point.")
combined_mask_for_image = None
if not individual_masks_for_image:
h, w = img_tensor.shape[0], img_tensor.shape[1]
combined_mask_for_image = torch.zeros((h, w), dtype=torch.float32, device=image.device)
logger.warning(f"No successful segmentations for any points on image {i}. Returning zero mask.")
elif len(individual_masks_for_image) == 1:
combined_mask_for_image = individual_masks_for_image[0]
else:
mask_stack = torch.stack(individual_masks_for_image, dim=0)
combined_mask_for_image = torch.max(mask_stack, dim=0).values
batch_results_combined.append(combined_mask_for_image)
return batch_results_combined
+415
View File
@@ -0,0 +1,415 @@
"""Model loader implementation for MediaPipe Vision tasks.
Handles model definition, download, and path retrieval.
"""
import os
import logging
import urllib.request
import folder_paths # ComfyUI path helper
from typing import Dict, Any, Optional, List
import traceback
# Configure logging
# logging.basicConfig(level=logging.INFO) # Avoid reconfiguring root logger
logger = logging.getLogger(__name__)
# --- Model Definitions ---
# Central dictionary mapping task types and variants to model details
# Structure: { task_type: { variant_name: { url: ..., filename: ... } } }
# Task types should match the keys used elsewhere (e.g., detector class validation)
# Variant names are displayed in the loader node UI
# Filenames should be unique within their task type directory
# URLs point to the official MediaPipe model files (.task or .tflite)
MEDIAPIPE_MODELS = {
"face_detector": {
"short_range": {
"url": "https://storage.googleapis.com/mediapipe-models/face_detector/blaze_face_short_range/float16/latest/blaze_face_short_range.tflite",
"filename": "blaze_face_short_range.tflite",
"description": "Optimized for detecting faces at close range, ideal for selfie-style images and near-field face detection."
},
"full_range_sparse": {
"url": "https://storage.googleapis.com/mediapipe-models/face_detector/blaze_face_back_camera/float16/latest/blaze_face_back_camera.tflite",
"filename": "blaze_face_full_range_sparse.tflite",
"description": "Designed for detecting faces at various distances, suitable for group photos and wider-angle shots."
}
},
"face_landmarker": {
"default": {
"url": "https://storage.googleapis.com/mediapipe-models/face_landmarker/face_landmarker/float16/latest/face_landmarker.task",
"filename": "face_landmarker.task",
"description": "Standard face landmark detection model that locates key facial points for tracking and analysis."
},
"with_blendshapes": { # If a separate model exists
"url": "https://storage.googleapis.com/mediapipe-models/face_landmarker/face_landmarker_v2_with_blendshapes/float16/latest/face_landmarker_v2_with_blendshapes.task",
"filename": "face_landmarker_v2_with_blendshapes.task",
"description": "Advanced face landmark model that includes blendshape coefficients for facial expression analysis."
}
},
"hand_landmarker": {
"default": { # Often includes world landmarks
"url": "https://storage.googleapis.com/mediapipe-models/hand_landmarker/hand_landmarker/float16/latest/hand_landmarker.task",
"filename": "hand_landmarker.task",
"description": "Detects hand landmarks for tracking finger positions and hand gestures in 2D and 3D space."
}
},
"pose_landmarker": {
"lite": {
"url": "https://storage.googleapis.com/mediapipe-models/pose_landmarker/pose_landmarker_lite/float16/latest/pose_landmarker_lite.task",
"filename": "pose_landmarker_lite.task",
"description": "Lightweight pose landmark model for faster inference with slightly reduced accuracy."
},
"full": {
"url": "https://storage.googleapis.com/mediapipe-models/pose_landmarker/pose_landmarker_full/float16/latest/pose_landmarker_full.task",
"filename": "pose_landmarker_full.task",
"description": "Balance of performance and accuracy for detecting body pose landmarks in typical use cases."
},
"heavy": {
"url": "https://storage.googleapis.com/mediapipe-models/pose_landmarker/pose_landmarker_heavy/float16/latest/pose_landmarker_heavy.task",
"filename": "pose_landmarker_heavy.task",
"description": "Higher accuracy pose detection model for more precise landmark tracking at the cost of performance."
}
},
"object_detector": {
"efficientdet_lite0": {
"url": "https://storage.googleapis.com/mediapipe-models/object_detector/efficientdet_lite0/float16/latest/efficientdet_lite0.tflite",
"filename": "efficientdet_lite0.tflite",
"description": "Lightweight object detection model that balances speed and accuracy for real-time applications."
},
"efficientdet_lite2": {
"url": "https://storage.googleapis.com/mediapipe-models/object_detector/efficientdet_lite2/float16/latest/efficientdet_lite2.tflite",
"filename": "efficientdet_lite2.tflite",
"description": "Medium-sized object detection model with improved accuracy over lite0 while maintaining reasonable performance."
},
"ssd_mobilenet_v2": {
"url": "https://storage.googleapis.com/mediapipe-models/object_detector/ssd_mobilenet_v2/float32/latest/ssd_mobilenet_v2.tflite",
"filename": "ssd_mobilenet_v2.tflite",
"description": "Single Shot Detector with MobileNetV2 backbone for efficient object detection tasks."
}
},
"image_segmenter": {
"deeplab_v3": {
"url": "https://storage.googleapis.com/mediapipe-models/image_segmenter/deeplab_v3/float32/latest/deeplab_v3.tflite",
"filename": "deeplab_v3.tflite",
"description": "General-purpose semantic segmentation model that can segment images into different object categories."
},
"selfie_segmenter_landscape": {
"url": "https://storage.googleapis.com/mediapipe-models/image_segmenter/selfie_segmenter_landscape/float16/latest/selfie_segmenter_landscape.tflite",
"filename": "selfie_segmenter_landscape.tflite",
"description": "Optimized for landscape-oriented images to accurately separate people from backgrounds."
},
"selfie_segmenter_general": {
"url": "https://storage.googleapis.com/mediapipe-models/image_segmenter/selfie_segmenter/float16/latest/selfie_segmenter.tflite",
"filename": "selfie_segmenter.tflite",
"description": "General-purpose person segmentation model that works across different orientations and scenarios."
},
"hair_segmenter": {
"url": "https://storage.googleapis.com/mediapipe-models/image_segmenter/hair_segmenter/float32/latest/hair_segmenter.tflite",
"filename": "hair_segmenter.tflite",
"description": "Specialized segmentation model for isolating hair in portrait images with high precision."
},
"selfie_multiclass": {
"url": "https://storage.googleapis.com/mediapipe-models/image_segmenter/selfie_multiclass_256x256/float32/latest/selfie_multiclass_256x256.tflite",
"filename": "selfie_multiclass_256x256.tflite",
"description": "Multi-class segmentation model that can separate people into different body parts and clothing items."
}
},
"gesture_recognizer": {
"default": {
"url": "https://storage.googleapis.com/mediapipe-models/gesture_recognizer/gesture_recognizer/float16/latest/gesture_recognizer.task",
"filename": "gesture_recognizer.task",
"description": "Detects and classifies hand gestures such as thumbs up, victory sign, and pointing gestures."
}
},
# "holistic_landmarker": {
# "lite": {
# "url": "https://storage.googleapis.com/mediapipe-models/holistic_landmarker/holistic_landmarker_lite/holistic_landmarker_lite.task",
# "filename": "holistic_landmarker_lite.task"
# },
# "full": {
# "url": "https://storage.googleapis.com/mediapipe-models/holistic_landmarker/holistic_landmarker_full/holistic_landmarker_full.task",
# "filename": "holistic_landmarker_full.task"
# }
# }
# Add other tasks like image_classification etc. here
"image_embedder": {
"mobilenet_v3_small": {
"url": "https://storage.googleapis.com/mediapipe-models/image_embedder/mobilenet_v3_small/float32/1/mobilenet_v3_small.tflite",
"filename": "mobilenet_v3_small.tflite",
"description": "Converts images into feature vectors (embeddings) for similarity comparison and classification tasks."
},
"efficientnet_lite0": {
"url": "https://storage.googleapis.com/mediapipe-models/image_embedder/efficientnet_lite0/float32/latest/efficientnet_lite0.tflite",
"filename": "efficientnet_lite0.tflite",
"description": "EfficientNet Lite0-based embedder providing robust image embeddings with a balance of accuracy and efficiency."
}
},
# Placeholder for Interactive Segmenter - Typically uses a general segmentation model
"interactive_segmenter": {
"magic_touch": {
"url": "https://storage.googleapis.com/mediapipe-models/interactive_segmenter/magic_touch/float32/latest/magic_touch.tflite",
"filename": "magic_touch.tflite",
"description": "Interactive segmentation model that segments objects in an image based on a user-provided point of interest."
},
"default": {
"url": "https://storage.googleapis.com/mediapipe-models/image_segmenter/selfie_segmenter/float16/latest/selfie_segmenter.tflite",
"filename": "selfie_segmenter_interactive.tflite",
"description": "Enables user-guided segmentation by providing points or regions to indicate foreground objects."
}
},
"face_stylizer": {
"color_sketch": {
"url": "https://storage.googleapis.com/mediapipe-models/face_stylizer/blaze_face_stylizer/float32/latest/face_stylizer_color_sketch.task",
"filename": "face_stylizer_color_sketch.task",
"description": "Transforms faces into colorful sketch-like artistic renderings with clean lines and vibrant colors."
},
"color_ink": {
"url": "https://storage.googleapis.com/mediapipe-models/face_stylizer/blaze_face_stylizer/float32/latest/face_stylizer_color_ink.task",
"filename": "face_stylizer_color_ink.task",
"description": "Applies a color ink effect to faces with flowing brush-like strokes and dynamic color blending."
},
"oil_painting": {
"url": "https://storage.googleapis.com/mediapipe-models/face_stylizer/blaze_face_stylizer/float32/latest/face_stylizer_oil_painting.task",
"filename": "face_stylizer_oil_painting.task",
"description": "Renders faces in classic oil painting style with textured brush strokes and rich color tones."
}
},
"image_classifier": {
"efficientnet_lite0": {
"url": "https://storage.googleapis.com/mediapipe-models/image_classifier/efficientnet_lite0/float32/latest/efficientnet_lite0.tflite",
"filename": "efficientnet_lite0.tflite",
"description": "EfficientNet-Lite0 model for general-purpose image classification tasks with optimized performance."
},
"efficientnet_lite2": {
"url": "https://storage.googleapis.com/mediapipe-models/image_classifier/efficientnet_lite2/float32/latest/efficientnet_lite2.tflite",
"filename": "efficientnet_lite2.tflite",
"description": "EfficientNet-Lite2 model offering higher accuracy for image classification tasks at the cost of increased computation."
}
},
"face_mesh": {
"default": {
"url": "https://storage.googleapis.com/mediapipe-models/face_mesh/face_mesh/float16/latest/face_mesh.tflite",
"filename": "face_mesh.tflite",
"description": "High-fidelity 3D face landmark model estimating 468 landmarks for detailed facial geometry analysis."
},
"attention": {
"url": "https://storage.googleapis.com/mediapipe-models/face_mesh/face_mesh_attention/float16/latest/face_mesh_attention.tflite",
"filename": "face_mesh_attention.tflite",
"description": "Enhanced face mesh model with attention mechanism for improved accuracy around eyes and lips."
}
}
}
# --- Utility Functions ---
def get_mediapipe_model_dir(task_type: str) -> str:
"""Gets the dedicated directory for a specific MediaPipe task type within ComfyUI models."""
base_path = None
try:
# Attempt to get the specific 'mediapipe' path defined in ComfyUI config
mediapipe_paths = folder_paths.get_folder_paths("mediapipe")
if mediapipe_paths and os.path.isdir(mediapipe_paths[0]):
base_path = mediapipe_paths[0]
logger.debug(f"Using specific MediaPipe model path: {base_path}")
else:
logger.debug("No valid specific 'mediapipe' path found or defined.")
except (KeyError, IndexError, TypeError) as e:
logger.debug(f"Could not retrieve specific 'mediapipe' path ({e}). Falling back to default.")
# If specific path wasn't found or valid, use the main models directory
if base_path is None:
try:
# Ensure models_dir exists
if hasattr(folder_paths, 'models_dir') and os.path.isdir(folder_paths.models_dir):
base_path = os.path.join(folder_paths.models_dir, "mediapipe")
logger.debug(f"Using fallback MediaPipe model path: {base_path}")
else:
raise RuntimeError("ComfyUI models directory ('models_dir') not found or invalid.")
except Exception as e:
logger.error(f"Error determining fallback MediaPipe directory: {e}")
raise RuntimeError("Could not determine a valid directory for MediaPipe models.") from e
# Construct the task-specific directory path
task_dir = os.path.join(base_path, task_type)
# Ensure the final task directory exists
try:
os.makedirs(task_dir, exist_ok=True)
except OSError as e:
logger.error(f"Failed to create MediaPipe task directory {task_dir}: {e}")
# Propagate error as directory creation is critical
raise RuntimeError(f"Failed to create required directory: {task_dir}") from e
return task_dir
def _download_url_to_file(url: str, save_path: str):
"""Downloads a file with progress indication."""
logger.info(f"Downloading model from {url} to {save_path}...")
try:
with urllib.request.urlopen(url) as response, open(save_path, 'wb') as out_file:
total_size_header = response.getheader('Content-Length')
total_size = int(total_size_header) if total_size_header else None
bytes_downloaded = 0
block_size = 8192
last_print_bytes = 0
print_interval = 1 * 1024 * 1024 # Print every 1MB
if total_size:
logger.info(f" File size: {total_size / (1024*1024):.2f} MB")
else:
logger.warning(" File size unknown.")
while True:
buffer = response.read(block_size)
if not buffer:
break
out_file.write(buffer)
bytes_downloaded += len(buffer)
# Progress reporting (less frequent)
if bytes_downloaded - last_print_bytes > print_interval or (total_size and bytes_downloaded == total_size):
progress_str = f" Downloaded: {bytes_downloaded / (1024*1024):.2f} MB"
if total_size:
progress = (bytes_downloaded / total_size) * 100
progress_str += f" / {total_size / (1024*1024):.2f} MB ({progress:.1f}%)"
# Use print with carriage return for inline update
print(progress_str, end=' \r')
last_print_bytes = bytes_downloaded
print() # Newline after download finishes
logger.info(f"Download complete: {save_path}")
return True
except Exception as e:
logger.error(f"Error downloading model: {e}")
logger.error(traceback.format_exc())
# Clean up potentially corrupted file
if os.path.exists(save_path):
try:
os.remove(save_path)
logger.info(f"Removed incomplete download: {save_path}")
except Exception as rm_err:
logger.error(f"Error removing incomplete download {save_path}: {rm_err}")
return False
def get_model_path(task_type: str, model_variant: str) -> Optional[str]:
"""Gets the full path to the model file, downloading if necessary."""
if task_type not in MEDIAPIPE_MODELS:
logger.error(f"Unknown MediaPipe task type: '{task_type}'")
return None
if model_variant not in MEDIAPIPE_MODELS[task_type]:
logger.error(f"Unknown model variant '{model_variant}' for task '{task_type}'")
# Log available variants for the task
available = list(MEDIAPIPE_MODELS[task_type].keys())
logger.info(f"Available variants for {task_type}: {available}")
return None
model_details = MEDIAPIPE_MODELS[task_type][model_variant]
model_filename = model_details["filename"]
model_url = model_details["url"]
# Get the specific directory for this task
try:
model_dir = get_mediapipe_model_dir(task_type)
except Exception as e:
logger.error(f"Cannot determine model directory for {task_type}: {e}")
return None
model_path = os.path.join(model_dir, model_filename)
# Check and download
if not os.path.exists(model_path):
logger.info(f"Model not found locally at {model_path}. Attempting download.")
if not _download_url_to_file(model_url, model_path):
return None # Download failed
else:
logger.debug(f"Found model locally: {model_path}")
return model_path
# --- Public Functions for Node Wrapper ---
def get_available_tasks() -> List[str]:
"""Get a list of available MediaPipe Vision task types from the definition."""
return list(MEDIAPIPE_MODELS.keys())
def get_available_models(task_type: str) -> List[str]:
"""Get a list of available model variants for a specific task."""
if task_type in MEDIAPIPE_MODELS:
return list(MEDIAPIPE_MODELS[task_type].keys())
else:
logger.warning(f"Task type '{task_type}' not found in model definitions.")
return []
def get_model_description(task_type: str, model_variant: str) -> str:
"""Get the description for a specific model variant.
Args:
task_type: The MediaPipe task type
model_variant: The specific model variant
Returns:
The description string if found, or a default message if not available
"""
if task_type in MEDIAPIPE_MODELS and model_variant in MEDIAPIPE_MODELS[task_type]:
return MEDIAPIPE_MODELS[task_type][model_variant].get(
"description", "No description available"
)
return "Model description not available"
# --- Model Loader Class (Simplified) ---
class MediaPipeModelLoader:
"""Helper class to manage loading state and info.
Used temporarily by the loader node to get model info.
"""
def __init__(self):
self.model_path = None
self.task_type = None
self.model_variant = None
self.model_loaded = False
self.model_info = {}
def load_model(self, task_type: str, model_variant: str) -> bool:
"""Sets model info and triggers path retrieval (incl. download)."""
logger.debug(f"MediaPipeModelLoader: Loading {task_type}/{model_variant}")
self.task_type = task_type
self.model_variant = model_variant
self.model_path = get_model_path(task_type, model_variant)
if self.model_path and os.path.exists(self.model_path):
self.model_loaded = True
self.model_info = {
'task_type': self.task_type,
'model_variant': self.model_variant,
'model_path': self.model_path,
'description': get_model_description(task_type, model_variant)
}
logger.debug(f"Model info set: {self.model_info}")
return True
else:
logger.error(f"Failed to get or validate model path for {task_type}/{model_variant}")
self.model_loaded = False
self.model_info = {}
return False
def get_model_info(self) -> Dict[str, Any]:
"""Get information about the loaded model."""
# Ensure info is consistent with loaded state
if self.model_loaded:
# Add the description to the model info
if self.task_type and self.model_variant:
description = get_model_description(self.task_type, self.model_variant)
self.model_info['description'] = description
return self.model_info
else:
# Return minimal info or empty dict if not loaded
return {
'task_type': self.task_type,
'model_variant': self.model_variant,
'model_path': None # Indicate path is not valid/loaded
}
# Removed original get_model_path, download_model etc. from global scope
# They are now integrated into the new get_model_path function.
@@ -0,0 +1 @@
"""Object Detection implementation using MediaPipe."""
@@ -0,0 +1,120 @@
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 typing import List
# Import new types
from ..types import ObjectDetectionResult, ObjectDetectionCategory, BoundingBox
from ..utils.timestamp_provider import TimestampProvider # Import TimestampProvider
class ObjectDetector:
"""Detects objects using MediaPipe ObjectDetector."""
def __init__(self, model_path: str):
"""Initialize the detector with the model path.
Args:
model_path: Path to the MediaPipe ObjectDetector .tflite 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 # 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(
base_options=base_options,
running_mode=mode_enum, # Use enum mode
score_threshold=score_threshold,
max_results=max_results
)
self._current_options = options
# Only create timestamp provider for VIDEO mode
self._timestamp_provider = TimestampProvider() if mode_enum == vision.RunningMode.VIDEO else None
return vision.ObjectDetector.create_from_options(options)
def detect(self, image: torch.Tensor, score_threshold: float = 0.5, 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
@@ -0,0 +1 @@
"""Pose Landmark Detection implementation using MediaPipe."""
@@ -0,0 +1,148 @@
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 typing import List, Tuple, Optional
# Import new types
from ..types import LandmarkPoint, PoseLandmarksResult
from ..utils.timestamp_provider import TimestampProvider # Import TimestampProvider
class PoseLandmarkDetector:
"""Detects pose landmarks in an image using MediaPipe PoseLandmarker."""
def __init__(self, model_path: str):
"""Initialize the detector with the model path.
Args:
model_path: Path to the MediaPipe PoseLandmarker .task 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 # 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(
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
)
self._current_options = options
# Only create timestamp provider for VIDEO mode
self._timestamp_provider = TimestampProvider() if mode_enum == vision.RunningMode.VIDEO else None
return vision.PoseLandmarker.create_from_options(options)
def detect(self, image: torch.Tensor, num_poses: int = 1, min_detection_confidence: float = 0.5,
min_presence_confidence: float = 0.5, min_tracking_confidence: float = 0.5,
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())
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
+150
View File
@@ -0,0 +1,150 @@
from typing import List, Optional, NamedTuple, Tuple
import numpy as np # Added numpy for matrix type
# Define a structure for a single landmark point
class LandmarkPoint(NamedTuple):
index: int
x: float
y: float
z: float
visibility: Optional[float] = None # Available in Pose landmarks
presence: Optional[float] = None # Available in Pose landmarks
# Define container structures for the output of each landmark task
# These will be returned by the respective detector nodes
class FaceLandmarksResult(NamedTuple):
landmarks: List[LandmarkPoint]
# Add other face-specific results if needed later (e.g., bounding box)
class HandLandmarksResult(NamedTuple):
landmarks: List[LandmarkPoint]
world_landmarks: Optional[List[LandmarkPoint]] # Keep world landmarks optional for now
handedness: Optional[str] # 'Left', 'Right'
class PoseLandmarksResult(NamedTuple):
landmarks: List[LandmarkPoint]
world_landmarks: Optional[List[LandmarkPoint]]
# Segmentation mask info could potentially go here if tightly coupled
# We might add more types later for detections, segmentations etc.
# For now, focus is on landmarks.
# Define the list types that will be used for batching
# These will be the types used in node RETURN_TYPES/INPUT_TYPES strings
FACE_LANDMARKS_LIST = List[List[FaceLandmarksResult]] # Batch[FacesPerImage[Result]]
HAND_LANDMARKS_LIST = List[List[HandLandmarksResult]] # Batch[HandsPerImage[Result]]
POSE_LANDMARKS_LIST = List[List[PoseLandmarksResult]] # Batch[PosesPerImage[Result]]
# --- Blendshape Result ---
class Blendshape(NamedTuple):
index: Optional[int]
score: float
category_name: str # e.g., mouthSmileLeft, eyeBlinkRight
display_name: Optional[str]
# --- Common Types ---
class BoundingBox(NamedTuple):
origin_x: int
origin_y: int
width: int
height: int
# --- Face Detection Result ---
class FaceKeypoint(NamedTuple):
label: Optional[str] # e.g., 'nose_tip', 'left_eye'
x: float
y: float
class FaceDetectionResult(NamedTuple):
bounding_box: BoundingBox
keypoints: Optional[List[FaceKeypoint]]
score: Optional[float] # Confidence score for the detection
# --- Object Detection Result ---
class ObjectDetectionCategory(NamedTuple):
index: Optional[int]
score: float
display_name: Optional[str]
category_name: Optional[str]
class ObjectDetectionResult(NamedTuple):
bounding_box: BoundingBox
categories: List[ObjectDetectionCategory]
# --- Batch List Types for Nodes ---
# Define the list types that will be used for batching
# These will be the types used in node RETURN_TYPES/INPUT_TYPES strings
FACE_LANDMARKS = List[List[FaceLandmarksResult]] # Batch[FacesPerImage[Result]]
HAND_LANDMARKS = List[List[HandLandmarksResult]] # Batch[HandsPerImage[Result]]
POSE_LANDMARKS = List[List[PoseLandmarksResult]] # Batch[PosesPerImage[Result]]
OBJECT_DETECTIONS = List[List[ObjectDetectionResult]] # Batch[DetectionsPerImage[Result]]
FACE_DETECTIONS = List[List[FaceDetectionResult]] # Batch[DetectionsPerImage[Result]]
# --- Gesture Recognition Result ---
class GestureCategory(NamedTuple):
index: Optional[int]
score: float
display_name: Optional[str]
category_name: Optional[str]
class GestureRecognitionResult(NamedTuple):
gestures: List[GestureCategory] # Top gesture(s) recognized for a hand
handedness: Optional[str] # Which hand this gesture belongs to
# Define Batch List Type
GESTURE_RECOGNITIONS = List[List[GestureRecognitionResult]] # Batch[Hands[Gestures]]
# --- Image Embedding Result ---
class EmbeddingEntry(NamedTuple):
index: int
score: float
display_name: Optional[str]
category_name: Optional[str]
class ImageEmbedderResult(NamedTuple):
float_embedding: Optional[List[float]]
quantized_embedding: Optional[bytes]
head_index: int
head_name: Optional[str]
# Define Batch List Type
IMAGE_EMBEDDINGS = List[List[ImageEmbedderResult]] # Batch[EmbeddingsPerHead[Result]]
# --- Interactive Segmentation Input/Result ---
class PointOfInterest(NamedTuple):
x: float
y: float
label: Optional[str] = None # Optional label for the point
# Result is typically a mask, potentially multiple if categories are involved
# Using existing MASK type might suffice, or define a specific one if needed.
# --- Holistic Landmark Result ---
class HolisticLandmarksResult(NamedTuple):
face_landmarks: Optional[List[LandmarkPoint]] # Only one face per person
pose_landmarks: Optional[List[LandmarkPoint]]
pose_world_landmarks: Optional[List[LandmarkPoint]]
left_hand_landmarks: Optional[List[LandmarkPoint]]
right_hand_landmarks: Optional[List[LandmarkPoint]]
# Pose segmentation mask could also be included here
# Define Batch List Type
HOLISTIC_LANDMARKS = List[HolisticLandmarksResult] # Batch[ResultPerPerson]
# =======================================
# --- Batch List Types for Node I/O ---
# =======================================
FACE_LANDMARKS = List[List[FaceLandmarksResult]] # Batch[FacesPerImage[Result]]
HAND_LANDMARKS = List[List[HandLandmarksResult]] # Batch[HandsPerImage[Result]]
POSE_LANDMARKS = List[List[PoseLandmarksResult]] # Batch[PosesPerImage[Result]]
OBJECT_DETECTIONS = List[List[ObjectDetectionResult]] # Batch[DetectionsPerImage[Result]]
FACE_DETECTIONS = List[List[FaceDetectionResult]] # Batch[DetectionsPerImage[Result]]
GESTURE_RECOGNITIONS = List[List[GestureRecognitionResult]] # Batch[Hands[Gestures]]
IMAGE_EMBEDDINGS = List[List[ImageEmbedderResult]] # Batch[EmbeddingsPerHead[Result]]
HOLISTIC_LANDMARKS = List[HolisticLandmarksResult] # Batch[ResultPerPerson]
# New Types - Defined AFTER Blendshape class
# Updated names to match FaceLandmarker output convention
BLENDSHAPES_LIST = List[List[List[Blendshape]]] # Batch[FacesPerImage[BlendshapeList[Blendshape]]]
TRANSFORM_MATRIX_LIST = List[List[np.ndarray]] # Batch[FacesPerImage[Matrix]]
+12
View File
@@ -0,0 +1,12 @@
"""Utility functions for MediaPipe Vision."""
from .download_utils import download_model
from .image_utils import convert_to_cv2, convert_to_tensor
from .delta_utils import calculate_euclidean_delta, scale_value
__all__ = [
"convert_to_cv2",
"convert_to_tensor",
"calculate_euclidean_delta",
"scale_value"
]
+11
View File
@@ -0,0 +1,11 @@
"""Utilities specific to ComfyUI integration."""
class AlwaysEqualProxy(str):
"""Helper class to allow connecting multiple specific input types to a single input node.
Borrowed from https://github.com/theUpsider/ComfyUI-Logic
"""
def __eq__(self, _):
return True
def __ne__(self, _):
return False
+35
View File
@@ -0,0 +1,35 @@
import math
from typing import Optional
from ..types import LandmarkPoint # Assuming LandmarkPoint is in src/types.py
def calculate_euclidean_delta(p1: Optional[LandmarkPoint], p2: Optional[LandmarkPoint]) -> Optional[float]:
"""Calculates the Euclidean distance between two LandmarkPoint objects."""
if p1 is None or p2 is None:
return None # Cannot calculate delta if one point is missing
delta_x = p1.x - p2.x
delta_y = p1.y - p2.y
delta_z = p1.z - p2.z
distance = math.sqrt(delta_x**2 + delta_y**2 + delta_z**2)
return distance
def scale_value(value: Optional[float],
delta_min: float, delta_max: float,
output_min: float, output_max: float,
clamp: bool = True) -> Optional[float]:
"""Linearly scales a value from one range to another."""
if value is None:
return None
if delta_max == delta_min: # Avoid division by zero
return output_min if value <= delta_min else output_max
# Scale
scaled = output_min + (value - delta_min) * (output_max - output_min) / (delta_max - delta_min)
# Clamp
if clamp:
scaled = max(output_min, min(output_max, scaled))
return scaled
@@ -0,0 +1,156 @@
"""Utilities for downloading MediaPipe models."""
import os
import requests
import logging
from tqdm import tqdm
from typing import Dict, Optional, Any, Tuple
from .model_utils import MEDIAPIPE_MODELS, get_model_path
# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
def download_model(task_type: str, model_variant: str) -> Optional[str]:
"""Download a MediaPipe model for a specific task and variant.
Args:
task_type: The task category (e.g., 'face_mesh', 'pose_detection')
model_variant: The specific model variant (e.g., 'standard', 'lite')
Returns:
Path to the downloaded model file, or None if download failed
"""
# Check if the task and variant exist
if task_type not in MEDIAPIPE_MODELS:
logger.error(f"Task '{task_type}' not found in available models.")
return None
if model_variant not in MEDIAPIPE_MODELS[task_type]:
logger.error(f"Model variant '{model_variant}' not found for task '{task_type}'.")
return None
# Get model details
model_info = MEDIAPIPE_MODELS[task_type][model_variant]
url = model_info["url"]
filename = model_info["filename"]
# Get the full path for the model
output_path = get_model_path(task_type, filename)
# If model already exists, return the path
if os.path.exists(output_path):
logger.info(f"Using existing model: {output_path}")
return output_path
# Download the model
logger.info(f"Downloading {task_type}/{model_variant} model from {url}...")
try:
# Ensure parent directory exists
os.makedirs(os.path.dirname(output_path), exist_ok=True)
# Stream the download with progress bar
response = requests.get(url, stream=True)
response.raise_for_status()
# Get content length for progress bar if available
total_size = int(response.headers.get('content-length', 0))
block_size = 8192 # 8 KB
with open(output_path, 'wb') as f:
with tqdm(total=total_size, unit='B', unit_scale=True, desc=f"{task_type}/{model_variant}") as pbar:
for chunk in response.iter_content(chunk_size=block_size):
if chunk:
f.write(chunk)
pbar.update(len(chunk))
logger.info(f"Successfully downloaded {task_type}/{model_variant} to {output_path}")
return output_path
except Exception as e:
logger.error(f"Error downloading {task_type}/{model_variant}: {str(e)}")
# Clean up partial download if it exists
if os.path.exists(output_path):
try:
os.remove(output_path)
except:
pass
return None
def get_model_options(task_type: Optional[str] = None) -> Dict[str, Any]:
"""Get available model options for a specific task or all tasks.
Args:
task_type: The task category to get options for, or None for all tasks
Returns:
Dictionary of available models and variants
"""
if task_type is None:
# Return all tasks and their variants
return {task: list(variants.keys()) for task, variants in MEDIAPIPE_MODELS.items()}
# Return variants for specific task
if task_type in MEDIAPIPE_MODELS:
return {task_type: list(MEDIAPIPE_MODELS[task_type].keys())}
# Task not found
return {}
def get_all_task_types() -> list:
"""Get a list of all available task types.
Returns:
List of task type strings
"""
return list(MEDIAPIPE_MODELS.keys())
def get_model_variants(task_type: str) -> list:
"""Get a list of all available model variants for a specific task.
Args:
task_type: The task category
Returns:
List of model variant strings, or empty list if task not found
"""
if task_type in MEDIAPIPE_MODELS:
return list(MEDIAPIPE_MODELS[task_type].keys())
return []
def get_model_info(task_type: str, model_variant: str) -> Tuple[Optional[str], Optional[Dict[str, Any]]]:
"""Get information about a specific model.
Args:
task_type: The task category
model_variant: The specific model variant
Returns:
Tuple containing:
- Path to the model file (downloaded if needed)
- Dictionary with model metadata (or None if task/variant not found)
"""
# Check if the task and variant exist
if task_type not in MEDIAPIPE_MODELS or model_variant not in MEDIAPIPE_MODELS[task_type]:
return None, None
# Get model details
model_info = MEDIAPIPE_MODELS[task_type][model_variant].copy()
# Download the model if needed and get the path
model_path = download_model(task_type, model_variant)
if not model_path:
return None, None
# Add the model path to the info
model_info["model_path"] = model_path
return model_path, model_info
+44
View File
@@ -0,0 +1,44 @@
import torch
import numpy as np
import cv2
def convert_to_cv2(tensor: torch.Tensor) -> list:
"""Converts a ComfyUI IMAGE tensor (BHWC, float32, 0-1) to a list of cv2 images (BGR, uint8)."""
if tensor.ndim != 4 or tensor.shape[3] != 3:
raise ValueError(f"Input tensor must be BHWC with 3 channels, got shape: {tensor.shape}")
# Convert entire batch to numpy, scale to 0-255, and convert to uint8
images_np = tensor.cpu().numpy()
images_np = (images_np * 255).clip(0, 255).astype(np.uint8)
# Convert RGB to BGR for cv2
cv2_images = [cv2.cvtColor(img, cv2.COLOR_RGB2BGR) for img in images_np]
return cv2_images
def convert_to_tensor(cv2_images) -> torch.Tensor:
"""Converts cv2 image(s) (BGR, uint8) to a ComfyUI IMAGE tensor (BHWC, float32, 0-1).
Args:
cv2_images: A single cv2 image or a list of cv2 images
"""
# Handle both single image and list of images
if not isinstance(cv2_images, list):
cv2_images = [cv2_images]
# Process each image
np_images = []
for cv2_image in cv2_images:
# Convert BGR to RGB
rgb_image = cv2.cvtColor(cv2_image, cv2.COLOR_BGR2RGB)
# Scale to 0-1
image_np = rgb_image.astype(np.float32) / 255.0
np_images.append(image_np)
# Stack all images into a batch
batch_np = np.stack(np_images, axis=0)
# Convert to tensor (already has batch dimension)
tensor = torch.from_numpy(batch_np)
return tensor
+26
View File
@@ -0,0 +1,26 @@
# --- Helper Function ---
import torch
import numpy as np
import cv2
from typing import List, Tuple
import logging
logger = logging.getLogger(__name__)
def create_mask_from_points(height: int, width: int, points: List[Tuple[int, int]], device='cpu') -> torch.Tensor:
"""Creates a binary mask tensor from a list of points using convex hull."""
mask = np.zeros((height, width), dtype=np.uint8)
if len(points) < 3:
logger.warning(f"Cannot create mask from {len(points)} points (need >= 3 for convex hull). Returning empty mask.")
return torch.from_numpy(mask).float().unsqueeze(0).to(device)
try:
np_points = np.array(points, dtype=np.int32)
hull = cv2.convexHull(np_points)
cv2.fillPoly(mask, [hull], 1)
except Exception as e:
logger.error(f"Error creating mask polygon: {e}")
return torch.from_numpy(mask).float().unsqueeze(0).to(device)
mask_tensor = torch.from_numpy(mask).float().unsqueeze(0).to(device)
return mask_tensor
+29
View File
@@ -0,0 +1,29 @@
# src/utils/misc_utils.py
import mediapipe as mp
import logging
from typing import List
logger = logging.getLogger(__name__)
def get_blendshape_categories() -> List[str]:
"""Returns the list of standard blendshape category names from MediaPipe."""
# Define blendshape categories directly
# This is based on the MediaPipe standard blendshapes
standard_blendshapes = [
"browDownLeft", "browDownRight", "browInnerUp", "browOuterUpLeft",
"browOuterUpRight", "cheekPuff", "cheekSquintLeft", "cheekSquintRight",
"eyeBlinkLeft", "eyeBlinkRight", "eyeLookDownLeft", "eyeLookDownRight",
"eyeLookInLeft", "eyeLookInRight", "eyeLookOutLeft", "eyeLookOutRight",
"eyeLookUpLeft", "eyeLookUpRight", "eyeSquintLeft", "eyeSquintRight",
"eyeWideLeft", "eyeWideRight", "jawForward", "jawLeft", "jawOpen",
"jawRight", "mouthClose", "mouthDimpleLeft", "mouthDimpleRight",
"mouthFrownLeft", "mouthFrownRight", "mouthFunnel", "mouthLeft",
"mouthLowerDownLeft", "mouthLowerDownRight", "mouthPressLeft",
"mouthPressRight", "mouthPucker", "mouthRight", "mouthRollLower",
"mouthRollUpper", "mouthShrugLower", "mouthShrugUpper", "mouthSmileLeft",
"mouthSmileRight", "mouthStretchLeft", "mouthStretchRight", "mouthUpperUpLeft",
"mouthUpperUpRight", "noseSneerLeft", "noseSneerRight"
]
logger.info("Using predefined list of standard blendshape categories.")
return standard_blendshapes
+148
View File
@@ -0,0 +1,148 @@
"""Utilities for managing MediaPipe Vision models."""
import os
import logging
from typing import Dict, Any, Optional
# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
# Import ComfyUI's folder_paths to use the models directory
try:
import folder_paths
HAS_COMFY_PATHS = True
except ImportError:
logger.warning("ComfyUI folder_paths module not available. Using local paths.")
HAS_COMFY_PATHS = False
def get_mediapipe_models_directory() -> str:
"""Get the directory where MediaPipe models are stored.
Uses ComfyUI's models directory structure when available.
Returns:
Path to the MediaPipe models directory
"""
if HAS_COMFY_PATHS:
base_path = folder_paths.models_dir
mediapipe_dir = os.path.join(base_path, "mediapipe", "vision")
else:
# Fallback to local path
base_path = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
mediapipe_dir = os.path.join(base_path, "models")
# Create the directory if it doesn't exist
os.makedirs(mediapipe_dir, exist_ok=True)
return mediapipe_dir
def get_model_path(model_type: str, model_name: str) -> str:
"""Get the path for a specific MediaPipe Vision model.
Args:
model_type: Type/category of model (e.g., 'face_mesh', 'pose')
model_name: Name of the specific model file
Returns:
Full path to the model file
"""
models_dir = get_mediapipe_models_directory()
model_type_dir = os.path.join(models_dir, model_type)
os.makedirs(model_type_dir, exist_ok=True)
return os.path.join(model_type_dir, model_name)
# Model information dictionary with URLs and files for each model type
# Using a nested dictionary structure to organize by task
MEDIAPIPE_MODELS = {
"face_detection": {
"short_range": {
"url": "https://storage.googleapis.com/mediapipe-models/face_detector/blaze_face_short_range/float16/1/blaze_face_short_range.task",
"filename": "blaze_face_short_range.task"
},
"full_range": {
"url": "https://storage.googleapis.com/mediapipe-models/face_detector/blaze_face/float16/1/blaze_face.task",
"filename": "blaze_face.task"
}
},
"face_mesh": {
"standard": {
"url": "https://storage.googleapis.com/mediapipe-models/face_landmarker/face_landmarker/float16/1/face_landmarker.task",
"filename": "face_landmarker.task",
"description": "Detects face landmarks and facial features (468 points)."
},
"attention": {
"url": "https://storage.googleapis.com/mediapipe-models/face_landmarker/face_landmarker/float16/2/face_landmarker.task",
"filename": "face_landmarker_v2.task",
"description": "Improved version with better accuracy for face landmarks."
}
},
"hand_tracking": {
"standard": {
"url": "https://storage.googleapis.com/mediapipe-models/hand_landmarker/hand_landmarker/float16/1/hand_landmarker.task",
"filename": "hand_landmarker.task",
"description": "Detects hand landmarks (21 points per hand)."
}
},
"pose_detection": {
"lite": {
"url": "https://storage.googleapis.com/mediapipe-models/pose_landmarker/pose_landmarker_lite/float16/1/pose_landmarker_lite.task",
"filename": "pose_landmarker_lite.task",
"description": "Lightweight pose detection model (33 points)."
},
"full": {
"url": "https://storage.googleapis.com/mediapipe-models/pose_landmarker/pose_landmarker_full/float16/1/pose_landmarker_full.task",
"filename": "pose_landmarker_full.task",
"description": "Full pose detection model with high accuracy."
},
"heavy": {
"url": "https://storage.googleapis.com/mediapipe-models/pose_landmarker/pose_landmarker_heavy/float16/1/pose_landmarker_heavy.task",
"filename": "pose_landmarker_heavy.task",
"description": "Most accurate pose detection model."
}
},
"image_segmentation": {
"selfie": {
"url": "https://storage.googleapis.com/mediapipe-models/image_segmenter/selfie_segmenter/float16/1/selfie_segmenter.task",
"filename": "selfie_segmenter.task",
"description": "Segments people from background."
},
"selfie_landscape": {
"url": "https://storage.googleapis.com/mediapipe-models/image_segmenter/selfie_segmenter_landscape/float16/1/selfie_segmenter_landscape.task",
"filename": "selfie_segmenter_landscape.task",
"description": "Optimized for landscape orientation."
},
"multiclass": {
"url": "https://storage.googleapis.com/mediapipe-models/image_segmenter/selfie_multiclass_256x256/float32/1/selfie_multiclass_256x256.task",
"filename": "selfie_multiclass_256x256.task",
"description": "Segments multiple body parts (face, hair, clothes, etc.)."
}
},
"object_detection": {
"efficientdet_lite0": {
"url": "https://storage.googleapis.com/mediapipe-models/object_detector/efficientdet_lite0/float16/1/efficientdet_lite0.task",
"filename": "efficientdet_lite0.task",
"description": "General object detection (80 classes)."
},
"efficientdet_lite2": {
"url": "https://storage.googleapis.com/mediapipe-models/object_detector/efficientdet_lite2/float16/1/efficientdet_lite2.task",
"filename": "efficientdet_lite2.task",
"description": "Higher accuracy object detection."
}
},
"image_classification": {
"efficientnet_lite0": {
"url": "https://storage.googleapis.com/mediapipe-models/image_classifier/efficientnet_lite0/float32/1/efficientnet_lite0.task",
"filename": "efficientnet_lite0.task",
"description": "General image classification (1000 classes)."
},
"efficientnet_lite4": {
"url": "https://storage.googleapis.com/mediapipe-models/image_classifier/efficientnet_lite4/float32/2/efficientnet_lite4.task",
"filename": "efficientnet_lite4.task",
"description": "Higher accuracy image classification."
}
}
}
@@ -0,0 +1,21 @@
import time
class TimestampProvider:
"""
Provides real-time, monotonic timestamps in milliseconds.
MediaPipe's RunningMode.VIDEO requires strictly increasing timestamp_ms values.
Using time.time() alone can result in duplicates during tight loops, which causes
MediaPipe to skip or freeze frames. This class ensures timestamps are always
increasing while staying aligned with real elapsed time.
"""
def __init__(self):
self._start_time = time.time()
self._last_timestamp = 0
def next(self) -> int:
now = int((time.time() - self._start_time) * 1000)
if now <= self._last_timestamp:
now = self._last_timestamp + 1
self._last_timestamp = now
return now