diff --git a/node_wrappers/mediapipe_vision/location_landmark/universal_landmark_nodes.py b/node_wrappers/mediapipe_vision/location_landmark/universal_landmark_nodes.py deleted file mode 100644 index b5b1e50..0000000 --- a/node_wrappers/mediapipe_vision/location_landmark/universal_landmark_nodes.py +++ /dev/null @@ -1,607 +0,0 @@ -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 - -#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, - LandmarkPoint, FaceLandmarksResult, HandLandmarksResult, PoseLandmarksResult -) - -# Import from the new consolidated utilities -from ....src.utils.math import scale_value -from ....src.utils.general 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 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" - # 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, 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.""" - - 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: - # 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.") - - # 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: - 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 --- - 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) - - # 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) - - -# --- 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", -} \ No newline at end of file diff --git a/src/coordinates/__init__.py b/src/coordinates/__init__.py new file mode 100644 index 0000000..2503eb2 --- /dev/null +++ b/src/coordinates/__init__.py @@ -0,0 +1,12 @@ +""" +Coordinates module for ComfyUI RealTimeNodes. + +This package provides unified coordinate system handling for ComfyUI RealTimeNodes, +allowing seamless conversion between different coordinate spaces (normalized, pixel) +with support for single values, batches, and point collections. +""" + +from .coordinate_system import CoordinateSystem, Point, PointList +from .drawing_engine import DrawingEngine + +__all__ = ['CoordinateSystem', 'Point', 'PointList', 'DrawingEngine'] \ No newline at end of file diff --git a/src/coordinates/coordinate_delta.py b/src/coordinates/coordinate_delta.py new file mode 100644 index 0000000..25cfb39 --- /dev/null +++ b/src/coordinates/coordinate_delta.py @@ -0,0 +1,61 @@ +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 + +#TODO: consider moving base nodes to src + +# Import from the new consolidated utilities +from ...src.utils.math import scale_value + +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 + + +# --- Base Position Delta Logic (Handling Multiple Histories) --- +class BaseCoordinateDelta: + """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 diff --git a/src/coordinates/coordinate_system.py b/src/coordinates/coordinate_system.py new file mode 100644 index 0000000..a00b240 --- /dev/null +++ b/src/coordinates/coordinate_system.py @@ -0,0 +1,422 @@ +""" +Unified coordinate system for ComfyUI RealTimeNodes. + +This module provides a centralized system for handling coordinates across different +spaces (normalized, pixel) with support for both single values and batches. +""" + +import torch +import numpy as np +from typing import Union, List, Tuple, Optional, Dict, Any +import logging + +logger = logging.getLogger(__name__) + +class CoordinateSystem: + """Unified coordinate system handling for coordinate space conversions.""" + + # Constants for coordinate spaces + NORMALIZED = "normalized" # 0.0-1.0 range + PIXEL = "pixel" # Pixel coordinates + + @staticmethod + def normalize(coords: Union[float, List[float], torch.Tensor], + dimensions: Union[int, Tuple[int, ...]], + input_space: str = PIXEL) -> Union[float, List[float], torch.Tensor]: + """ + Normalize coordinates to 0.0-1.0 range. + + Args: + coords: Coordinates to normalize + dimensions: Dimensions to normalize by (width, height) + input_space: Current coordinate space + + Returns: + Normalized coordinates (0.0-1.0) + """ + if input_space == CoordinateSystem.NORMALIZED: + return coords + + # Handle different dimension specifications + if isinstance(dimensions, tuple): + dim_size = dimensions[0] if len(dimensions) == 1 else dimensions + else: + dim_size = dimensions + + # Handle different input types + if isinstance(coords, list): + if isinstance(dim_size, tuple): + # Multiple dimensions + if len(coords) % len(dim_size) != 0: + logger.warning(f"Coordinates length {len(coords)} not divisible by dimensions length {len(dim_size)}") + result = [] + for i, c in enumerate(coords): + dim_idx = i % len(dim_size) + result.append(float(c) / dim_size[dim_idx]) + return result + else: + # Single dimension + return [float(c) / dim_size for c in coords] + elif isinstance(coords, torch.Tensor): + if isinstance(dim_size, tuple): + # Tensor with shape handling + if coords.shape[-1] != len(dim_size): + raise ValueError(f"Tensor last dimension {coords.shape[-1]} doesn't match dimensions {len(dim_size)}") + return coords.float() / torch.tensor(dim_size, device=coords.device) + else: + return coords.float() / dim_size + else: + return float(coords) / dim_size + + @staticmethod + def denormalize(coords: Union[float, List[float], torch.Tensor], + dimensions: Union[int, Tuple[int, ...]], + output_space: str = PIXEL) -> Union[float, List[float], torch.Tensor]: + """ + Convert normalized coordinates to target space. + + Args: + coords: Normalized coordinates (0.0-1.0) + dimensions: Target dimensions + output_space: Target coordinate space + + Returns: + Denormalized coordinates in target space + """ + if output_space == CoordinateSystem.NORMALIZED: + return coords + + # Handle different dimension specifications + if isinstance(dimensions, tuple): + dim_size = dimensions[0] if len(dimensions) == 1 else dimensions + else: + dim_size = dimensions + + # Handle different input types + if isinstance(coords, list): + if isinstance(dim_size, tuple): + # Multiple dimensions + if len(coords) % len(dim_size) != 0: + logger.warning(f"Coordinates length {len(coords)} not divisible by dimensions length {len(dim_size)}") + result = [] + for i, c in enumerate(coords): + dim_idx = i % len(dim_size) + result.append(float(max(0, min(dim_size[dim_idx] - 1, c * dim_size[dim_idx])))) + return result + else: + # Single dimension + return [float(max(0, min(dim_size - 1, c * dim_size))) for c in coords] + elif isinstance(coords, torch.Tensor): + if isinstance(dim_size, tuple): + # Tensor with shape handling + if coords.shape[-1] != len(dim_size): + raise ValueError(f"Tensor last dimension {coords.shape[-1]} doesn't match dimensions {len(dim_size)}") + dim_tensor = torch.tensor(dim_size, device=coords.device) + return torch.clamp(coords * dim_tensor, 0, dim_tensor - 1) + else: + return torch.clamp(coords * dim_size, 0, dim_size - 1) + else: + return float(max(0, min(dim_size - 1, coords * dim_size))) + + @staticmethod + def convert(coords: Union[float, List[float], torch.Tensor], + dimensions: Union[int, Tuple[int, ...]], + from_space: str, + to_space: str) -> Union[float, List[float], torch.Tensor]: + """ + Convert between coordinate spaces. + + Args: + coords: Coordinates to convert + dimensions: Target dimensions + from_space: Source coordinate space + to_space: Target coordinate space + + Returns: + Converted coordinates + """ + if from_space == to_space: + return coords + + # First normalize if needed + if from_space != CoordinateSystem.NORMALIZED: + normalized = CoordinateSystem.normalize(coords, dimensions, from_space) + else: + normalized = coords + + # Then denormalize to target space if needed + if to_space != CoordinateSystem.NORMALIZED: + return CoordinateSystem.denormalize(normalized, dimensions, to_space) + else: + return normalized + + @staticmethod + def get_dimensions_from_tensor(tensor: torch.Tensor) -> Tuple[int, int]: + """ + Extract dimensions from tensor in BHWC format. + + Args: + tensor: Input tensor in BHWC format + + Returns: + Tuple of (width, height) + """ + if tensor.dim() != 4: + raise ValueError(f"Expected 4D tensor (BHWC), got shape {tensor.shape}") + + _, height, width, _ = tensor.shape + return (width, height) + +class Point: + """ + Represents a 2D or 3D point with coordinate space awareness. + """ + + def __init__(self, + x: float, + y: float, + z: Optional[float] = None, + space: str = CoordinateSystem.NORMALIZED): + """ + Initialize a point with coordinates and space information. + + Args: + x: X coordinate + y: Y coordinate + z: Z coordinate (optional) + space: Coordinate space + """ + self.x = x + self.y = y + self.z = z + self.space = space + self.is_3d = z is not None + + def to_normalized(self, dimensions: Tuple[int, ...]) -> 'Point': + """ + Convert point to normalized space. + + Args: + dimensions: Dimensions for conversion + + Returns: + New point in normalized space + """ + if self.space == CoordinateSystem.NORMALIZED: + return self + + norm_x = CoordinateSystem.normalize(self.x, dimensions[0], self.space) + norm_y = CoordinateSystem.normalize(self.y, dimensions[1], self.space) + norm_z = None + + if self.is_3d and len(dimensions) > 2: + norm_z = CoordinateSystem.normalize(self.z, dimensions[2], self.space) + + return Point(norm_x, norm_y, norm_z, CoordinateSystem.NORMALIZED) + + def to_pixel(self, dimensions: Tuple[int, ...]) -> 'Point': + """ + Convert point to pixel space. + + Args: + dimensions: Dimensions for conversion + + Returns: + New point in pixel space + """ + if self.space == CoordinateSystem.PIXEL: + return self + + # First normalize if not already + if self.space != CoordinateSystem.NORMALIZED: + normalized = self.to_normalized(dimensions) + else: + normalized = self + + pixel_x = CoordinateSystem.denormalize(normalized.x, dimensions[0], CoordinateSystem.PIXEL) + pixel_y = CoordinateSystem.denormalize(normalized.y, dimensions[1], CoordinateSystem.PIXEL) + pixel_z = None + + if self.is_3d and len(dimensions) > 2: + pixel_z = CoordinateSystem.denormalize(normalized.z, dimensions[2], CoordinateSystem.PIXEL) + + return Point(pixel_x, pixel_y, pixel_z, CoordinateSystem.PIXEL) + + def to_space(self, dimensions: Tuple[int, ...], target_space: str) -> 'Point': + """ + Convert point to target space. + + Args: + dimensions: Dimensions for conversion + target_space: Target coordinate space + + Returns: + New point in target space + """ + if self.space == target_space: + return self + + if target_space == CoordinateSystem.NORMALIZED: + return self.to_normalized(dimensions) + elif target_space == CoordinateSystem.PIXEL: + return self.to_pixel(dimensions) + else: + raise ValueError(f"Unsupported target space: {target_space}") + + def to_tuple(self, include_z: bool = False) -> Tuple[float, ...]: + """ + Convert point to tuple representation. + + Args: + include_z: Whether to include Z coordinate if available + + Returns: + Tuple of coordinates + """ + if include_z and self.is_3d: + return (self.x, self.y, self.z) + else: + return (self.x, self.y) + + def to_list(self, include_z: bool = False) -> List[float]: + """ + Convert point to list representation. + + Args: + include_z: Whether to include Z coordinate if available + + Returns: + List of coordinates + """ + return list(self.to_tuple(include_z)) + + def __repr__(self) -> str: + """String representation of the point.""" + if self.is_3d: + return f"Point({self.x}, {self.y}, {self.z}, space='{self.space}')" + else: + return f"Point({self.x}, {self.y}, space='{self.space}')" + +class PointList: + """ + Represents a list of 2D or 3D points with coordinate space awareness. + Useful for handling batches of points or landmark lists. + """ + + def __init__(self, + points: List[Point] = None, + space: str = CoordinateSystem.NORMALIZED): + """ + Initialize a point list. + + Args: + points: List of Point objects + space: Coordinate space for all points + """ + self.points = points or [] + self.space = space + + @classmethod + def from_coordinates(cls, + x_coords: List[float], + y_coords: List[float], + z_coords: Optional[List[float]] = None, + space: str = CoordinateSystem.NORMALIZED) -> 'PointList': + """ + Create a PointList from separate coordinate lists. + + Args: + x_coords: List of X coordinates + y_coords: List of Y coordinates + z_coords: List of Z coordinates (optional) + space: Coordinate space + + Returns: + New PointList object + """ + if len(x_coords) != len(y_coords): + raise ValueError(f"Coordinate lists must have same length: x={len(x_coords)}, y={len(y_coords)}") + + has_z = z_coords is not None + if has_z and len(z_coords) != len(x_coords): + raise ValueError(f"Z coordinates must match length of X/Y: x={len(x_coords)}, z={len(z_coords)}") + + points = [] + for i in range(len(x_coords)): + z = z_coords[i] if has_z else None + points.append(Point(x_coords[i], y_coords[i], z, space)) + + return cls(points, space) + + def to_normalized(self, dimensions: Tuple[int, ...]) -> 'PointList': + """ + Convert all points to normalized space. + + Args: + dimensions: Dimensions for conversion + + Returns: + New PointList in normalized space + """ + if self.space == CoordinateSystem.NORMALIZED: + return self + + normalized_points = [point.to_normalized(dimensions) for point in self.points] + return PointList(normalized_points, CoordinateSystem.NORMALIZED) + + def to_pixel(self, dimensions: Tuple[int, ...]) -> 'PointList': + """ + Convert all points to pixel space. + + Args: + dimensions: Dimensions for conversion + + Returns: + New PointList in pixel space + """ + if self.space == CoordinateSystem.PIXEL: + return self + + pixel_points = [point.to_pixel(dimensions) for point in self.points] + return PointList(pixel_points, CoordinateSystem.PIXEL) + + def to_space(self, dimensions: Tuple[int, ...], target_space: str) -> 'PointList': + """ + Convert all points to target space. + + Args: + dimensions: Dimensions for conversion + target_space: Target coordinate space + + Returns: + New PointList in target space + """ + if self.space == target_space: + return self + + if target_space == CoordinateSystem.NORMALIZED: + return self.to_normalized(dimensions) + elif target_space == CoordinateSystem.PIXEL: + return self.to_pixel(dimensions) + else: + raise ValueError(f"Unsupported target space: {target_space}") + + def get_x_list(self) -> List[float]: + """Get list of X coordinates.""" + return [point.x for point in self.points] + + def get_y_list(self) -> List[float]: + """Get list of Y coordinates.""" + return [point.y for point in self.points] + + def get_z_list(self) -> List[Optional[float]]: + """Get list of Z coordinates (may contain None values).""" + return [point.z for point in self.points] + + def __len__(self) -> int: + """Get number of points in the list.""" + return len(self.points) + + def __getitem__(self, idx: int) -> Point: + """Access point by index.""" + return self.points[idx] \ No newline at end of file