coordinate to src
This commit is contained in:
RyanOnTheInside
@@ -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",
|
||||
}
|
||||
@@ -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']
|
||||
@@ -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
|
||||
@@ -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]
|
||||
Reference in New Issue
Block a user