diff --git a/node_wrappers/mediapipe_vision/face/landmark/blendshape.py b/node_wrappers/mediapipe_vision/face/landmark/blendshape.py index 7ea855d..61a79b1 100644 --- a/node_wrappers/mediapipe_vision/face/landmark/blendshape.py +++ b/node_wrappers/mediapipe_vision/face/landmark/blendshape.py @@ -6,7 +6,7 @@ from .....src.mediapipe_vision.types import BLENDSHAPES_LIST # Updated type alia 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 ...location_landmark.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__) diff --git a/node_wrappers/mediapipe_vision/face/landmark/head_pose.py b/node_wrappers/mediapipe_vision/face/landmark/head_pose.py index 8eff28d..4a0b891 100644 --- a/node_wrappers/mediapipe_vision/face/landmark/head_pose.py +++ b/node_wrappers/mediapipe_vision/face/landmark/head_pose.py @@ -6,7 +6,7 @@ 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 +from ...location_landmark.universal_landmark_nodes import FLOAT_EQUALITY_TOLERANCE # For trigger node comparisons logger = logging.getLogger(__name__) diff --git a/node_wrappers/mediapipe_vision/face/landmark/position.py b/node_wrappers/mediapipe_vision/face/landmark/position.py index 449e50c..5f60f34 100644 --- a/node_wrappers/mediapipe_vision/face/landmark/position.py +++ b/node_wrappers/mediapipe_vision/face/landmark/position.py @@ -1,5 +1,5 @@ import logging -from ...common.universal_landmark_nodes import LandmarkPositionBaseNode +from ...location_landmark.universal_landmark_nodes import LandmarkPositionBaseNode # Import the landmark definitions from .....src.mediapipe_vision.landmark_definitions import FACE_LANDMARK_TOOLTIP @@ -22,6 +22,7 @@ class FaceLandmarkPositionNode(LandmarkPositionBaseNode): "required": { "landmarks": (cls.LANDMARKS_TYPE, {"forceInput": True}), # Use the class variable "landmark_index": ("INT", {"default": 0, "min": 0, "tooltip": FACE_LANDMARK_TOOLTIP}), + "result_index": ("INT", {"default": 0, "min": 0, "tooltip": "Index of the face detection to use (0=first detected face)"}), } } diff --git a/node_wrappers/mediapipe_vision/hand/landmark/position.py b/node_wrappers/mediapipe_vision/hand/landmark/position.py index af924ab..0ce6805 100644 --- a/node_wrappers/mediapipe_vision/hand/landmark/position.py +++ b/node_wrappers/mediapipe_vision/hand/landmark/position.py @@ -3,7 +3,7 @@ import logging from typing import List, Tuple, Any, Optional # Import the base class -from ...common.universal_landmark_nodes import LandmarkPositionBaseNode +from ...location_landmark.universal_landmark_nodes import LandmarkPositionBaseNode # Import the specific landmark type for clarity and potential future use from .....src.mediapipe_vision.types import HAND_LANDMARKS # Import the landmark definitions @@ -27,6 +27,7 @@ class HandLandmarkPositionNode(LandmarkPositionBaseNode): "required": { "landmarks": (cls.LANDMARKS_TYPE, {"forceInput": True}), "landmark_index": ("INT", {"default": 8, "min": 0, "max": 20, "tooltip": HAND_LANDMARK_TOOLTIP}), + "result_index": ("INT", {"default": 0, "min": 0, "tooltip": "Index of the hand detection to use (0=first detected hand)"}), "use_world_coordinates": ("BOOLEAN", {"default": False, "tooltip": "Use world coordinates if available"}), } } diff --git a/node_wrappers/mediapipe_vision/location_landmark/drawing_nodes.py b/node_wrappers/mediapipe_vision/location_landmark/drawing_nodes.py new file mode 100644 index 0000000..5c411d7 --- /dev/null +++ b/node_wrappers/mediapipe_vision/location_landmark/drawing_nodes.py @@ -0,0 +1,219 @@ +import torch +import logging +from typing import Union, List + +from ....src.mediapipe_vision.location_landmark.drawing_engine import DrawingEngine + +logger = logging.getLogger(__name__) + +class RTDrawPointsNode: + """ + High-performance point drawing node optimized for real-time applications. + Draws points (circles) on images with minimal overhead. + """ + CATEGORY = "Realtime Nodes/MediaPipe Vision/Common/Drawing" + FUNCTION = "draw_points" + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image_out",) + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "x": ("FLOAT", {"default": 0.5, "forceInput": True, "tooltip": "X coordinate(s) (0.0-1.0 normalized by default)"}), + "y": ("FLOAT", {"default": 0.5, "forceInput": True, "tooltip": "Y coordinate(s) (0.0-1.0 normalized by default)"}), + "radius": ("INT", {"default": 5, "min": 1, "max": 100, "tooltip": "Radius of points in pixels"}), + "color_hex": ("STRING", {"default": "#FF0000", "tooltip": "Color in #RRGGBB format"}), + }, + "optional": { + "is_normalized": ("BOOLEAN", {"default": True, "tooltip": "Whether coordinates are normalized (0-1) or pixel values"}), + "batch_mapping": (["broadcast", "one-to-one", "all-on-first"], {"default": "broadcast", "tooltip": "How to map coords to batch images"}) + } + } + + def draw_points(self, image: torch.Tensor, x: Union[float, List[float]], y: Union[float, List[float]], + radius: int, color_hex: str, is_normalized: bool = True, batch_mapping: str = "broadcast"): + """Draw points efficiently using the DrawingEngine.""" + drawing_engine = DrawingEngine() + return (drawing_engine.draw_points( + image=image, + points_x=x, + points_y=y, + radius=radius, + color=color_hex, + is_normalized=is_normalized, + batch_mapping=batch_mapping + ),) + +class RTDrawLinesNode: + """ + High-performance line drawing node optimized for real-time applications. + Draws lines on images with minimal overhead. + """ + CATEGORY = "Realtime Nodes/MediaPipe Vision/Common/Drawing" + FUNCTION = "draw_lines" + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image_out",) + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "x1": ("FLOAT", {"default": 0.25, "forceInput": True, "tooltip": "Start X coordinate(s) (0.0-1.0 normalized by default)"}), + "y1": ("FLOAT", {"default": 0.25, "forceInput": True, "tooltip": "Start Y coordinate(s) (0.0-1.0 normalized by default)"}), + "x2": ("FLOAT", {"default": 0.75, "forceInput": True, "tooltip": "End X coordinate(s) (0.0-1.0 normalized by default)"}), + "y2": ("FLOAT", {"default": 0.75, "forceInput": True, "tooltip": "End Y coordinate(s) (0.0-1.0 normalized by default)"}), + "thickness": ("INT", {"default": 2, "min": 1, "max": 20, "tooltip": "Line thickness in pixels"}), + "color_hex": ("STRING", {"default": "#00FF00", "tooltip": "Color in #RRGGBB format"}), + }, + "optional": { + "is_normalized": ("BOOLEAN", {"default": True, "tooltip": "Whether coordinates are normalized (0-1) or pixel values"}), + "batch_mapping": (["broadcast", "one-to-one", "all-on-first"], {"default": "broadcast"}), + "draw_endpoints": ("BOOLEAN", {"default": False, "tooltip": "Draw circles at line endpoints"}), + "point_radius": ("INT", {"default": 3, "min": 1, "max": 20, "tooltip": "Radius of endpoint circles in pixels"}), + "draw_label": ("BOOLEAN", {"default": False, "tooltip": "Draw text label on line"}), + "label_text": ("STRING", {"default": "Line", "tooltip": "Text to display as label"}), + "label_position": (["Midpoint", "Start", "End"], {"default": "Midpoint"}), + "font_scale": ("FLOAT", {"default": 0.5, "min": 0.1, "max": 2.0, "step": 0.1}) + } + } + + def draw_lines(self, image: torch.Tensor, + x1: Union[float, List[float]], y1: Union[float, List[float]], + x2: Union[float, List[float]], y2: Union[float, List[float]], + thickness: int, color_hex: str, + is_normalized: bool = True, batch_mapping: str = "broadcast", + draw_endpoints: bool = False, point_radius: int = 3, + draw_label: bool = False, label_text: str = "Line", + label_position: str = "Midpoint", font_scale: float = 0.5): + """Draw lines efficiently using the DrawingEngine.""" + drawing_engine = DrawingEngine() + return (drawing_engine.draw_lines( + image=image, + x1=x1, y1=y1, x2=x2, y2=y2, + thickness=thickness, + color=color_hex, + is_normalized=is_normalized, + batch_mapping=batch_mapping, + draw_endpoints=draw_endpoints, + point_radius=point_radius, + draw_label=draw_label, + label_text=label_text, + label_position=label_position, + font_scale=font_scale + ),) + +class RTDrawPolygonNode: + """ + High-performance polygon drawing node optimized for real-time applications. + Draws polygons on images with minimal overhead. + """ + CATEGORY = "Realtime Nodes/MediaPipe Vision/Common/Drawing" + FUNCTION = "draw_polygon" + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image_out",) + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "points_x": ("FLOAT", {"default": [0.25, 0.75, 0.5], "forceInput": True, "tooltip": "List of X coordinates for polygon vertices"}), + "points_y": ("FLOAT", {"default": [0.25, 0.25, 0.75], "forceInput": True, "tooltip": "List of Y coordinates for polygon vertices"}), + "color_hex": ("STRING", {"default": "#0000FF", "tooltip": "Color in #RRGGBB format"}), + "thickness": ("INT", {"default": 2, "min": 1, "max": 20, "tooltip": "Line thickness for polygon outline"}), + }, + "optional": { + "fill": ("BOOLEAN", {"default": False, "tooltip": "Whether to fill the polygon"}), + "is_normalized": ("BOOLEAN", {"default": True, "tooltip": "Whether coordinates are normalized (0-1) or pixel values"}), + "batch_mapping": (["broadcast", "one-to-one", "all-on-first"], {"default": "broadcast"}), + "draw_vertices": ("BOOLEAN", {"default": False, "tooltip": "Draw points at polygon vertices"}), + "vertex_radius": ("INT", {"default": 3, "min": 1, "max": 20, "tooltip": "Radius of vertex points"}), + "draw_label": ("BOOLEAN", {"default": False, "tooltip": "Draw text label on polygon"}), + "label_text": ("STRING", {"default": "Polygon", "tooltip": "Text to display as label"}), + "font_scale": ("FLOAT", {"default": 0.5, "min": 0.1, "max": 2.0, "step": 0.1}) + } + } + + def draw_polygon(self, image: torch.Tensor, + points_x: List[float], points_y: List[float], + color_hex: str, thickness: int, + fill: bool = False, is_normalized: bool = True, + batch_mapping: str = "broadcast", draw_vertices: bool = False, + vertex_radius: int = 3, draw_label: bool = False, + label_text: str = "Polygon", font_scale: float = 0.5): + """Draw polygon efficiently using the DrawingEngine.""" + drawing_engine = DrawingEngine() + return (drawing_engine.draw_polygon( + image=image, + points_x=points_x, + points_y=points_y, + color=color_hex, + thickness=thickness, + fill=fill, + is_normalized=is_normalized, + batch_mapping=batch_mapping, + draw_vertices=draw_vertices, + vertex_radius=vertex_radius, + draw_label=draw_label, + label_text=label_text, + label_position="Center", # Only center supported for polygons currently + font_scale=font_scale + ),) + +class RTCoordinateConverterNode: + """ + Fast coordinate conversion between normalized (0-1) and pixel space. + Optimized for real-time applications with minimal overhead. + """ + CATEGORY = "Realtime Nodes/MediaPipe Vision/Common/Coordinates" + FUNCTION = "convert_coordinates" + RETURN_TYPES = ("FLOAT", "FLOAT") + RETURN_NAMES = ("x_out", "y_out") + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "x": ("FLOAT", {"default": 0.0, "forceInput": True, "tooltip": "X coordinate(s) to convert"}), + "y": ("FLOAT", {"default": 0.0, "forceInput": True, "tooltip": "Y coordinate(s) to convert"}), + "image_for_dimensions": ("IMAGE", {"tooltip": "Reference image for dimensions"}), + "mode": (["Pixel to Normalized", "Normalized to Pixel"],), + } + } + + def convert_coordinates(self, x: Union[float, List[float]], y: Union[float, List[float]], + image_for_dimensions: torch.Tensor, mode: str): + """Convert coordinates efficiently using DrawingEngine methods.""" + if image_for_dimensions.dim() != 4: + logger.error("Input image must be BHWC format.") + return ([0.0] * len(x), [0.0] * len(y)) if isinstance(x, list) else (0.0, 0.0) + + _, height, width, _ = image_for_dimensions.shape + drawing_engine = DrawingEngine() + + if mode == "Pixel to Normalized": + x_out = drawing_engine.normalize_coords(x, width) + y_out = drawing_engine.normalize_coords(y, height) + else: # "Normalized to Pixel" + x_out = drawing_engine.denormalize_coords(x, width) + y_out = drawing_engine.denormalize_coords(y, height) + + return (x_out, y_out) + +# Register nodes +NODE_CLASS_MAPPINGS = { + "RTDrawPoints": RTDrawPointsNode, + "RTDrawLines": RTDrawLinesNode, + "RTDrawPolygon": RTDrawPolygonNode, + "RTCoordinateConverter": RTCoordinateConverterNode, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "RTDrawPoints": "Draw Points (Fast)", + "RTDrawLines": "Draw Lines (Fast)", + "RTDrawPolygon": "Draw Polygon (Fast)", + "RTCoordinateConverter": "Coordinate Converter (Fast)", +} \ No newline at end of file diff --git a/node_wrappers/mediapipe_vision/common/universal_landmark_nodes.py b/node_wrappers/mediapipe_vision/location_landmark/universal_landmark_nodes.py similarity index 89% rename from node_wrappers/mediapipe_vision/common/universal_landmark_nodes.py rename to node_wrappers/mediapipe_vision/location_landmark/universal_landmark_nodes.py index 2ceb8dc..b284201 100644 --- a/node_wrappers/mediapipe_vision/common/universal_landmark_nodes.py +++ b/node_wrappers/mediapipe_vision/location_landmark/universal_landmark_nodes.py @@ -7,6 +7,8 @@ import cv2 import collections import typing # Add typing for Dict hint +#TODO: consider moving base nodes to src + # Assuming types are correctly imported relative to this file's location from ....src.mediapipe_vision.types import ( FACE_LANDMARKS, HAND_LANDMARKS, POSE_LANDMARKS, @@ -28,7 +30,7 @@ class LandmarkPositionBaseNode: 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. + Processes only the detected object at the specified result_index per batch item. RETURN_TYPES declare the semantic type (FLOAT/BOOLEAN); the actual return value may be a list. """ FUNCTION = "extract_positions" @@ -42,7 +44,7 @@ class LandmarkPositionBaseNode: # 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]]: + def extract_positions(self, landmarks: list, landmark_index: int, result_index: int = 0, 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.""" @@ -69,38 +71,42 @@ class LandmarkPositionBaseNode: # 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 + # Get the result at the specified index if it exists + if result_index < len(item_list): + result_item = item_list[result_index] - 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.") + 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 + # 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 + # --- 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: + logger.debug(f"{self.__class__.__name__}: Result index {result_index} exceeds available detections for batch item {batch_idx} (has {len(item_list)})") # else: item_list was not a list or was empty, keep defaults/invalid # --- Extract Values or Use Defaults, Set Validity --- @@ -123,6 +129,10 @@ class LandmarkPositionBaseNode: pres_coords.append(current_pres) is_valid_list.append(is_valid) + # Check if each list has only one element, if so return single values instead of lists + if len(x_coords) == 1: + return (x_coords[0], y_coords[0], z_coords[0], vis_coords[0], pres_coords[0], is_valid_list[0]) + return (x_coords, y_coords, z_coords, vis_coords, pres_coords, is_valid_list) diff --git a/node_wrappers/mediapipe_vision/pose/position.py b/node_wrappers/mediapipe_vision/pose/position.py index 0d39751..5b4788d 100644 --- a/node_wrappers/mediapipe_vision/pose/position.py +++ b/node_wrappers/mediapipe_vision/pose/position.py @@ -2,7 +2,7 @@ import logging from typing import List, Tuple, Any, Optional # Import the base class -from ..common.universal_landmark_nodes import LandmarkPositionBaseNode +from ..location_landmark.universal_landmark_nodes import LandmarkPositionBaseNode # Import the specific landmark type for clarity and potential future use from ....src.mediapipe_vision.types import POSE_LANDMARKS # Import the landmark definitions @@ -26,6 +26,7 @@ class PoseLandmarkPositionNode(LandmarkPositionBaseNode): "required": { "landmarks": (cls.LANDMARKS_TYPE, {"forceInput": True}), "landmark_index": ("INT", {"default": 0, "min": 0, "max": 32, "tooltip": POSE_LANDMARK_TOOLTIP}), + "result_index": ("INT", {"default": 0, "min": 0, "tooltip": "Index of the pose detection to use (0=first detected pose)"}), "use_world_coordinates": ("BOOLEAN", {"default": False, "tooltip": "Use world coordinates if available"}), } } diff --git a/node_wrappers/utils/misc_nodes.py b/node_wrappers/utils/misc_nodes.py index 3779896..7ce8959 100644 --- a/node_wrappers/utils/misc_nodes.py +++ b/node_wrappers/utils/misc_nodes.py @@ -1011,3 +1011,234 @@ class FeedbackEffect: return result +class ToString: + """A node that converts any input to a string.""" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "any": (AlwaysEqualProxy("*"),{}), # Accept any input type + "format_string": ("STRING", { + "default": "{}", + "multiline": False, + "tooltip": "Optional format string (e.g., '{:.2f}' for 2 decimal place float)" + }) + }, + "optional": { + "prepend": ("STRING", { + "default": "", + "multiline": False, + "tooltip": "Text to add before the converted string" + }), + "append": ("STRING", { + "default": "", + "multiline": False, + "tooltip": "Text to add after the converted string" + }) + } + } + + RETURN_TYPES = ("STRING",) + FUNCTION = "to_string" + CATEGORY = "utils" + + def to_string(self, any, format_string="{}", prepend="", append=""): + try: + # Try to use the format string if provided + result = format_string.format(any) + except (ValueError, TypeError): + # Fall back to simple string conversion if format fails + result = str(any) + + # Add prepend and append text + result = prepend + result + append + + return (result,) + +class RoundNode: + """A node that rounds a float to a specified number of decimal places.""" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "value": ("FLOAT", { + "default": 0.0, + "tooltip": "Float value to round" + }), + "decimal_places": ("INT", { + "default": 2, + "min": 0, + "max": 10, + "step": 1, + "tooltip": "Number of decimal places to round to" + }) + } + } + + RETURN_TYPES = ("FLOAT",) + RETURN_NAMES = ("rounded_value",) + FUNCTION = "round_value" + CATEGORY = "utils" + + def round_value(self, value, decimal_places): + if isinstance(value, (list, tuple, np.ndarray, torch.Tensor)): + rounded = [round(v, decimal_places) for v in value] + else: + rounded = round(value, decimal_places) + return (rounded,) + +class RenormalizeFloat: + """Renormalizes a float value from one range to another""" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "value": ("FLOAT", { + "default": 0.5, + "forceInput": True, + "tooltip": "Float value to renormalize" + }), + "input_min": ("FLOAT", { + "default": 0.0, + "step": 0.01, + "tooltip": "Minimum value of the input range" + }), + "input_max": ("FLOAT", { + "default": 1.0, + "step": 0.01, + "tooltip": "Maximum value of the input range" + }), + "output_min": ("FLOAT", { + "default": 0.0, + "step": 0.01, + "tooltip": "Minimum value of the output range" + }), + "output_max": ("FLOAT", { + "default": 1.0, + "step": 0.01, + "tooltip": "Maximum value of the output range" + }), + "clamp": ("BOOLEAN", { + "default": True, + "tooltip": "Clamp output to the output range" + }) + } + } + + RETURN_TYPES = ("FLOAT",) + RETURN_NAMES = ("renormalized",) + FUNCTION = "renormalize" + CATEGORY = "utils" + + def renormalize(self, value, input_min, input_max, output_min, output_max, clamp): + """Rescales a value from an input range to an output range""" + # Check if value is a list, tuple, numpy array, or torch tensor + if isinstance(value, (list, tuple, np.ndarray, torch.Tensor)): + # Process element-wise + result = [] + for v in value: + result.append(self._renormalize_single(v, input_min, input_max, output_min, output_max, clamp)) + return (result,) + else: + # Process single value + result = self._renormalize_single(value, input_min, input_max, output_min, output_max, clamp) + return (result,) + + def _renormalize_single(self, value, input_min, input_max, output_min, output_max, clamp): + """Helper method to renormalize a single value""" + # Handle cases where input_min equals input_max + if abs(input_max - input_min) < 1e-6: + return output_min + + # Calculate normalized position in input range + normalized = (value - input_min) / (input_max - input_min) + + # Map to output range + result = output_min + normalized * (output_max - output_min) + + # Apply clamping if requested + if clamp: + if output_min <= output_max: + result = max(output_min, min(output_max, result)) + else: + result = max(output_max, min(output_min, result)) + + return result + +class RenormalizeInt: + """Renormalizes an integer value from one range to another""" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "value": ("INT", { + "default": 50, + "forceInput": True, + "tooltip": "Integer value to renormalize" + }), + "input_min": ("INT", { + "default": 0, + "tooltip": "Minimum value of the input range" + }), + "input_max": ("INT", { + "default": 100, + "tooltip": "Maximum value of the input range" + }), + "output_min": ("INT", { + "default": 0, + "tooltip": "Minimum value of the output range" + }), + "output_max": ("INT", { + "default": 100, + "tooltip": "Maximum value of the output range" + }), + "clamp": ("BOOLEAN", { + "default": True, + "tooltip": "Clamp output to the output range" + }) + } + } + + RETURN_TYPES = ("INT",) + RETURN_NAMES = ("renormalized",) + FUNCTION = "renormalize" + CATEGORY = "utils" + + def renormalize(self, value, input_min, input_max, output_min, output_max, clamp): + """Rescales an integer value from an input range to an output range""" + # Check if value is a list, tuple, numpy array, or torch tensor + if isinstance(value, (list, tuple, np.ndarray, torch.Tensor)): + # Process element-wise + result = [] + for v in value: + result.append(self._renormalize_single(v, input_min, input_max, output_min, output_max, clamp)) + return (result,) + else: + # Process single value + result = self._renormalize_single(value, input_min, input_max, output_min, output_max, clamp) + return (result,) + + def _renormalize_single(self, value, input_min, input_max, output_min, output_max, clamp): + """Helper method to renormalize a single integer value""" + # Handle cases where input_min equals input_max + if input_max == input_min: + return output_min + + # Calculate normalized position in input range + normalized = (value - input_min) / (input_max - input_min) + + # Map to output range and round to integer + result = round(output_min + normalized * (output_max - output_min)) + + # Apply clamping if requested + if clamp: + if output_min <= output_max: + result = max(output_min, min(output_max, result)) + else: + result = max(output_max, min(output_min, result)) + + return result diff --git a/src/mediapipe_vision/location_landmark/drawing_engine.py b/src/mediapipe_vision/location_landmark/drawing_engine.py new file mode 100644 index 0000000..a3bd59a --- /dev/null +++ b/src/mediapipe_vision/location_landmark/drawing_engine.py @@ -0,0 +1,425 @@ +import torch +import numpy as np +import cv2 +from typing import Union, List, Tuple, Dict, Optional, Any +import logging + +logger = logging.getLogger(__name__) + +class DrawingEngine: + """High-performance drawing engine for real-time applications. + Optimized for minimal tensor-numpy conversions and efficient batch operations. + """ + + @staticmethod + def hex_to_bgr(hex_color: str) -> Tuple[int, int, int]: + """Converts hex color to BGR tuple with validation.""" + hex_color = hex_color.lstrip('#') + if len(hex_color) != 6: + return (0, 0, 0) # Default black + try: + rgb = tuple(int(hex_color[i:i+2], 16) for i in (0, 2, 4)) + return (rgb[2], rgb[1], rgb[0]) # RGB to BGR + except ValueError: + return (0, 0, 0) + + @staticmethod + def normalize_coords(coords: Union[float, List[float]], dim_size: int) -> Union[float, List[float]]: + """Normalize pixel coordinates to 0.0-1.0 range.""" + if isinstance(coords, list): + return [float(c) / dim_size for c in coords] + return float(coords) / dim_size + + @staticmethod + def denormalize_coords(coords: Union[float, List[float]], dim_size: int) -> Union[float, List[float]]: + """Convert normalized coordinates to pixel space.""" + if isinstance(coords, list): + return [float(max(0, min(dim_size - 1, c * dim_size))) for c in coords] + return float(max(0, min(dim_size - 1, coords * dim_size))) + + @staticmethod + def prep_batch_image(image: torch.Tensor, batch_idx: int) -> Tuple[np.ndarray, torch.device, torch.dtype]: + """Fast preparation of single image from batch for drawing - minimal conversion.""" + if image.dim() != 4: + raise ValueError("Input image must be BHWC format") + + device, dtype = image.device, image.dtype + # Get single image, convert to float32 for processing + img_tensor = image[batch_idx].to(torch.float32) + # Single efficient conversion to numpy + img_np = (img_tensor.cpu().numpy() * 255).astype(np.uint8) + # Ensure contiguous for CV2 operations + return np.ascontiguousarray(img_np), device, dtype + + @staticmethod + def tensor_from_numpy(img_np: np.ndarray, device: torch.device, dtype: torch.dtype) -> torch.Tensor: + """Fast conversion from numpy back to tensor with proper device/dtype.""" + return torch.from_numpy(img_np.astype(np.float32) / 255.0).to(device=device, dtype=dtype) + + def draw_points(self, + image: torch.Tensor, + points_x: Union[float, List[float]], + points_y: Union[float, List[float]], + radius: int = 5, + color: str = "#FF0000", + is_normalized: bool = True, + batch_mapping: str = "broadcast") -> torch.Tensor: + """ + Draw points on images efficiently. + + Args: + image: BHWC tensor + points_x/y: Coordinate values (normalized 0-1 or pixel) + radius: Point radius in pixels + color: Hex color + is_normalized: If True, coords are 0-1 normalized + batch_mapping: How to map points to batch ("broadcast", "one-to-one", "all-on-first") + + Returns: + Modified image tensor + """ + batch_size, height, width, _ = image.shape + output_batch = image.clone() # Avoid modifying input + bgr_color = self.hex_to_bgr(color) + + # Standardize input to lists for unified processing + x_list = points_x if isinstance(points_x, list) else [points_x] + y_list = points_y if isinstance(points_y, list) else [points_y] + + # Determine batch mapping strategy + points_len = len(x_list) + if batch_mapping == "one-to-one" and points_len != batch_size and points_len != 1: + logger.warning(f"One-to-one mapping requested but point list length ({points_len}) != batch size ({batch_size})") + batch_mapping = "all-on-first" + + # Process each image in batch + for b in range(batch_size): + # Skip processing based on mapping strategy + if batch_mapping == "all-on-first" and b > 0: + continue + + # Get points for this batch item based on mapping strategy + if batch_mapping == "one-to-one" and points_len == batch_size: + # One point per image + curr_x = [x_list[b]] + curr_y = [y_list[b]] + elif batch_mapping == "broadcast" or (batch_mapping == "one-to-one" and points_len == 1): + # Same point(s) on all images + curr_x = x_list + curr_y = y_list + else: # "all-on-first" + # All points on first image only + curr_x = x_list + curr_y = y_list + + # Skip if no points to draw + if not curr_x: + continue + + # Single efficient conversion to numpy + img_np, device, dtype = self.prep_batch_image(output_batch, b) + + # Draw all points for this image + for xi, yi in zip(curr_x, curr_y): + if not isinstance(xi, (float, int)) or not isinstance(yi, (float, int)): + continue + + # Convert to pixel coordinates if needed + if is_normalized: + px = int(max(0, min(width - 1, xi * width))) + py = int(max(0, min(height - 1, yi * height))) + else: + px = int(max(0, min(width - 1, xi))) + py = int(max(0, min(height - 1, yi))) + + # Draw circle + cv2.circle(img_np, (px, py), radius, bgr_color, -1) + + # Convert back to tensor and update batch + output_batch[b] = self.tensor_from_numpy(img_np, device, dtype) + + return output_batch + + def draw_lines(self, + image: torch.Tensor, + x1: Union[float, List[float]], + y1: Union[float, List[float]], + x2: Union[float, List[float]], + y2: Union[float, List[float]], + thickness: int = 2, + color: str = "#00FF00", + is_normalized: bool = True, + batch_mapping: str = "broadcast", + draw_endpoints: bool = False, + point_radius: int = 3, + draw_label: bool = False, + label_text: str = "", + label_position: str = "Midpoint", + font_scale: float = 0.5) -> torch.Tensor: + """ + Draw lines on images efficiently. + + Args: + image: BHWC tensor + x1/y1/x2/y2: Line endpoint coordinates + thickness: Line thickness in pixels + color: Hex color + is_normalized: If True, coords are 0-1 normalized + batch_mapping: How to map lines to batch + draw_endpoints: Whether to draw circles at endpoints + Other params: Styling options + + Returns: + Modified image tensor + """ + batch_size, height, width, _ = image.shape + output_batch = image.clone() + bgr_color = self.hex_to_bgr(color) + + # Standardize input to lists + x1_list = x1 if isinstance(x1, list) else [x1] + y1_list = y1 if isinstance(y1, list) else [y1] + x2_list = x2 if isinstance(x2, list) else [x2] + y2_list = y2 if isinstance(y2, list) else [y2] + + # Validate input lengths + line_len = len(x1_list) + if not all(len(coord) == line_len for coord in [y1_list, x2_list, y2_list]): + raise ValueError("All coordinate lists must have the same length") + + # Determine batch mapping + if batch_mapping == "one-to-one" and line_len != batch_size and line_len != 1: + logger.warning(f"One-to-one mapping requested but line list length ({line_len}) != batch size ({batch_size})") + batch_mapping = "all-on-first" + + # Process each image in batch + for b in range(batch_size): + # Skip processing based on mapping strategy + if batch_mapping == "all-on-first" and b > 0: + continue + + # Get lines for this batch item + if batch_mapping == "one-to-one" and line_len == batch_size: + indices = [b] + elif batch_mapping == "broadcast" or (batch_mapping == "one-to-one" and line_len == 1): + indices = range(line_len) + else: # "all-on-first" + indices = range(line_len) + + # Skip if no lines to draw + if not indices: + continue + + # Single efficient conversion to numpy + img_np, device, dtype = self.prep_batch_image(output_batch, b) + + # Draw all lines for this image + for i in indices: + _x1, _y1 = x1_list[i], y1_list[i] + _x2, _y2 = x2_list[i], y2_list[i] + + if not all(isinstance(c, (float, int)) for c in [_x1, _y1, _x2, _y2]): + continue + + # Convert to pixel coordinates if needed + if is_normalized: + px1 = int(max(0, min(width - 1, _x1 * width))) + py1 = int(max(0, min(height - 1, _y1 * height))) + px2 = int(max(0, min(width - 1, _x2 * width))) + py2 = int(max(0, min(height - 1, _y2 * height))) + else: + px1 = int(max(0, min(width - 1, _x1))) + py1 = int(max(0, min(height - 1, _y1))) + px2 = int(max(0, min(width - 1, _x2))) + py2 = int(max(0, min(height - 1, _y2))) + + # Draw line + cv2.line(img_np, (px1, py1), (px2, py2), bgr_color, thickness) + + # Optional features + if draw_endpoints: + cv2.circle(img_np, (px1, py1), point_radius, bgr_color, -1) + cv2.circle(img_np, (px2, py2), point_radius, bgr_color, -1) + + if draw_label and label_text: + # Calculate label position + if label_position == "Midpoint": + label_x = int((px1 + px2) / 2) + label_y = int((py1 + py2) / 2) + elif label_position == "Start": + label_x, label_y = px1, py1 + else: # End or fallback + label_x, label_y = px2, py2 + + # Add offset and draw text + label_y += 10 + font = cv2.FONT_HERSHEY_SIMPLEX + cv2.putText(img_np, label_text, (label_x, label_y), + font, font_scale, bgr_color, 1, cv2.LINE_AA) + + # Convert back to tensor and update batch + output_batch[b] = self.tensor_from_numpy(img_np, device, dtype) + + return output_batch + + def draw_polygon(self, + image: torch.Tensor, + points_x: List[float], + points_y: List[float], + color: str = "#0000FF", + thickness: int = 2, + fill: bool = False, + is_normalized: bool = True, + batch_mapping: str = "broadcast", + draw_vertices: bool = False, + vertex_radius: int = 3, + draw_label: bool = False, + label_text: str = "", + label_position: str = "Center", + font_scale: float = 0.5) -> torch.Tensor: + """ + Draw polygons on images efficiently. + + Args: + image: BHWC tensor + points_x/y: Lists of polygon vertex coordinates + color: Hex color + thickness: Line thickness for polygon outline + fill: Whether to fill the polygon + is_normalized: If True, coords are 0-1 normalized + batch_mapping: How to map polygons to batch + draw_vertices: Whether to draw points at polygon vertices + vertex_radius: Radius of vertex points + draw_label: Whether to draw a text label + label_text: Text to display + label_position: Position of the label ("Center" only for now) + font_scale: Font scale + + Returns: + Modified image tensor + """ + batch_size, height, width, _ = image.shape + output_batch = image.clone() + bgr_color = self.hex_to_bgr(color) + + # Validate input + if not isinstance(points_x, list) or not isinstance(points_y, list): + logger.error("Polygon requires lists of points") + return output_batch + + if len(points_x) != len(points_y): + logger.error("points_x and points_y must have the same length") + return output_batch + + if len(points_x) < 3: + logger.error("Polygon requires at least 3 points") + return output_batch + + # Check if we have a list of polygon point sets or just a single polygon + is_multi_polygon = False + if isinstance(points_x[0], list): + is_multi_polygon = True + if not all(isinstance(p, list) for p in points_x + points_y): + logger.error("If providing multiple polygons, all entries in points_x/y must be lists") + return output_batch + if len(points_x) != len(points_y): + logger.error("For multiple polygons, points_x and points_y lists must have the same length") + return output_batch + + poly_count = len(points_x) + # Validate each polygon has matching x,y counts + for i in range(poly_count): + if len(points_x[i]) != len(points_y[i]): + logger.error(f"Polygon {i}: points_x and points_y must have the same length") + return output_batch + if len(points_x[i]) < 3: + logger.error(f"Polygon {i} requires at least 3 points") + return output_batch + else: + # Single polygon case - wrap in outer list for consistent processing + points_x = [points_x] + points_y = [points_y] + poly_count = 1 + + # Determine batch mapping + if batch_mapping == "one-to-one" and poly_count != batch_size and poly_count != 1: + logger.warning(f"One-to-one mapping requested but polygon count ({poly_count}) != batch size ({batch_size})") + batch_mapping = "all-on-first" + + # Process each image in batch + for b in range(batch_size): + # Skip processing based on mapping strategy + if batch_mapping == "all-on-first" and b > 0: + continue + + # Get polygons for this batch item + if batch_mapping == "one-to-one" and poly_count == batch_size: + poly_indices = [b] + elif batch_mapping == "broadcast" or (batch_mapping == "one-to-one" and poly_count == 1): + poly_indices = range(poly_count) + else: # "all-on-first" + poly_indices = range(poly_count) + + # Skip if no polygons to draw + if not poly_indices: + continue + + # Single efficient conversion to numpy + img_np, device, dtype = self.prep_batch_image(output_batch, b) + + # Draw all polygons for this image + for i in poly_indices: + px = points_x[i] + py = points_y[i] + + # Convert to pixel coordinates if needed + if is_normalized: + # Vectorized conversion + px_pixels = [int(max(0, min(width - 1, x * width))) for x in px] + py_pixels = [int(max(0, min(height - 1, y * height))) for y in py] + else: + # Ensure within bounds + px_pixels = [int(max(0, min(width - 1, x))) for x in px] + py_pixels = [int(max(0, min(height - 1, y))) for y in py] + + # Combine x,y into points format needed by OpenCV + points = np.array(list(zip(px_pixels, py_pixels)), np.int32) + points = points.reshape((-1, 1, 2)) + + # Draw filled polygon + if fill: + cv2.fillPoly(img_np, [points], bgr_color) + + # Draw polygon outline + cv2.polylines(img_np, [points], True, bgr_color, thickness) + + # Draw vertices if requested + if draw_vertices: + for vx, vy in zip(px_pixels, py_pixels): + cv2.circle(img_np, (vx, vy), vertex_radius, bgr_color, -1) + + # Draw label if requested + if draw_label and label_text: + # Calculate centroid for label position + centroid_x = sum(px_pixels) // len(px_pixels) + centroid_y = sum(py_pixels) // len(py_pixels) + + # Add offset and draw text with outline for better visibility + font = cv2.FONT_HERSHEY_SIMPLEX + text_size = cv2.getTextSize(label_text, font, font_scale, 1)[0] + + # Center text + text_x = max(0, min(width - text_size[0], centroid_x - text_size[0] // 2)) + text_y = centroid_y + + # Draw text outline/shadow for better visibility + cv2.putText(img_np, label_text, (text_x, text_y), + font, font_scale, (0, 0, 0), 2, cv2.LINE_AA) # Thicker black outline + # Draw main text + cv2.putText(img_np, label_text, (text_x, text_y), + font, font_scale, bgr_color, 1, cv2.LINE_AA) + + # Convert back to tensor and update batch + output_batch[b] = self.tensor_from_numpy(img_np, device, dtype) + + return output_batch \ No newline at end of file