455 lines
18 KiB
Python
455 lines
18 KiB
Python
import logging
|
|
import math
|
|
from typing import List, Tuple, Union
|
|
|
|
import cv2
|
|
import numpy as np
|
|
import torch
|
|
|
|
from ...src.coordinates.coordinate_delta import BaseCoordinateDelta, MAX_POSITION_HISTORY, FLOAT_EQUALITY_TOLERANCE
|
|
from ...src.coordinates import coordinate_utils
|
|
|
|
# TODO: consider moving base nodes to src
|
|
# Import from the new consolidated utilities
|
|
from ...src.utils.math import scale_value
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
# --- Position Delta Control Nodes (Int/Float) (Handling Lists) ---
|
|
class CoordinateDeltaControlNode(BaseCoordinateDelta):
|
|
"""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)
|
|
|
|
# Check required arguments
|
|
required_args = {
|
|
"x": x_in, "y": y_in, "z": z_in, "delta_min": delta_min,
|
|
"delta_max": delta_max, "output_min": output_min, "output_max": output_max, "clamp": clamp
|
|
}
|
|
if None in required_args.values():
|
|
return (fallback_value,)
|
|
|
|
# Validate input types
|
|
is_list = coordinate_utils.validate_coordinate_input_types(x_in, y_in, z_in, self.__class__.__name__)
|
|
|
|
if is_list:
|
|
current_batch_size = len(x_in)
|
|
# Clean up histories
|
|
coordinate_utils.cleanup_coordinate_histories(self._histories, current_batch_size, self.__class__.__name__)
|
|
|
|
# Handle empty lists
|
|
if current_batch_size == 0:
|
|
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]):
|
|
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:
|
|
# Clean up single input histories
|
|
coordinate_utils.cleanup_single_input_histories(self._histories, self.__class__.__name__)
|
|
|
|
# Validate single coordinate types
|
|
if not all(isinstance(c, (float, int)) for c in [x_in, y_in, z_in]):
|
|
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 CoordinateDeltaIntControlNode(CoordinateDeltaControlNode):
|
|
RETURN_TYPES = ("INT",) # Semantic type is INT
|
|
RETURN_NAMES = ("value",)
|
|
IS_INT_OUTPUT = True
|
|
DEFAULT_OUTPUT_MIN = 0
|
|
DEFAULT_OUTPUT_MAX = 100
|
|
|
|
|
|
class CoordinateDeltaFloatControlNode(CoordinateDeltaControlNode):
|
|
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 CoordinateDeltaTriggerNode(BaseCoordinateDelta):
|
|
"""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
|
|
|
|
# Check required arguments
|
|
required_args = {
|
|
"x": x_in, "y": y_in, "z": z_in, "threshold": threshold, "condition": condition
|
|
}
|
|
if None in required_args.values():
|
|
return (fallback_value,)
|
|
|
|
# Validate input types
|
|
is_list = coordinate_utils.validate_coordinate_input_types(x_in, y_in, z_in, self.__class__.__name__)
|
|
|
|
if is_list:
|
|
current_batch_size = len(x_in)
|
|
# Clean up histories
|
|
coordinate_utils.cleanup_coordinate_histories(self._histories, current_batch_size, self.__class__.__name__)
|
|
|
|
# Handle empty lists
|
|
if current_batch_size == 0:
|
|
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]):
|
|
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:
|
|
# Clean up single input histories
|
|
coordinate_utils.cleanup_single_input_histories(self._histories, self.__class__.__name__)
|
|
|
|
# Validate single coordinate types
|
|
if not all(isinstance(c, (float, int)) for c in [x_in, y_in, z_in]):
|
|
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 CoordinateProximityNode:
|
|
"""
|
|
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
|
|
# Validate inputs
|
|
inputs = [x1, y1, z1, x2, y2, z2]
|
|
is_empty, list_len = coordinate_utils.validate_proximity_inputs(inputs, self.__class__.__name__)
|
|
fallback_value = 0.0
|
|
|
|
if is_empty:
|
|
return ([],) # Return empty list for FLOAT type
|
|
|
|
is_list_input = isinstance(x1, list)
|
|
|
|
if is_list_input:
|
|
# Process List Input
|
|
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
|
|
|
|
|
|
class MaskFromCoordinate:
|
|
"""
|
|
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]:
|
|
# Validate coordinate inputs
|
|
is_empty, num_coords = coordinate_utils.validate_mask_coordinate_inputs(x, y, self.__class__.__name__)
|
|
|
|
if is_empty:
|
|
# 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),)
|
|
|
|
is_x_list = isinstance(x, list)
|
|
if is_x_list:
|
|
# Input is List[float]
|
|
x_list = x
|
|
y_list = y
|
|
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)):
|
|
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,)
|