diff --git a/.gitignore b/.gitignore index e1251c7..fbdf4f3 100644 --- a/.gitignore +++ b/.gitignore @@ -27,6 +27,7 @@ ENV/ # IDE files .idea/ .vscode/ +.cursor/ *.swp *.swo diff --git a/__init__.py b/__init__.py index d7c660e..4629c9c 100644 --- a/__init__.py +++ b/__init__.py @@ -1,11 +1,23 @@ from .controls.value_controls import FloatControl, IntControl, StringControl from .controls.sequence_controls import FloatSequence, IntSequence, StringSequence from .controls.utility_controls import FPSMonitor, SimilarityFilter, LazyCondition +from .controls.state_management_controls import StateResetNode, StateTestNode from .controls.motion_controls import MotionController, ROINode, IntegerMotionController -from .misc_nodes import DTypeConverter, FastWebcamCapture, YOLOSimilarityCompare, TextRenderer, QuickShapeMask, MultilineText, LoadImageFromPath_ +from .misc_nodes import( + DTypeConverter, + FastWebcamCapture, + YOLOSimilarityCompare, + TextRenderer, + QuickShapeMask, + MultilineText, + LoadImageFromPath_ +) +from .media_pipe_nodes import HandTrackingNode, HandMaskNode +from .controls.mask_controls import RepulsiveMaskNode, ResizeMaskNode import re + NODE_CLASS_MAPPINGS = { "FloatControl": FloatControl, "IntControl": IntControl, @@ -21,15 +33,24 @@ NODE_CLASS_MAPPINGS = { "YOLOSimilarityCompare": YOLOSimilarityCompare, "TextRenderer": TextRenderer, "ROINode": ROINode, - #"IntervalControl": IntervalCo ntrol, - #"DeltaControl": DeltaControl, "QuickShapeMask": QuickShapeMask, "DTypeConverter": DTypeConverter, "FastWebcamCapture": FastWebcamCapture, "MultilineText": MultilineText, "LoadImageFromPath_": LoadImageFromPath_, + "HandTrackingNode": HandTrackingNode, + "HandMaskNode": HandMaskNode, + #"RepulsiveMaskNode": RepulsiveMaskNode, + "ResizeMaskNode": ResizeMaskNode, + "StateResetNode": StateResetNode, + "StateTestNode": StateTestNode, } + + + + + NODE_DISPLAY_NAME_MAPPINGS = {} suffix = " 🕒🅡🅣🅝" diff --git a/base/control_base.py b/base/control_base.py index f14fada..9af10e3 100644 --- a/base/control_base.py +++ b/base/control_base.py @@ -16,13 +16,35 @@ class Pattern(ABC): class StateManager: """Handles persistent state between executions""" def __init__(self): - self._states = {} + self._states = {} def get_state(self, node_id: str, default=None): return self._states.get(node_id, default) def set_state(self, node_id: str, state): self._states[node_id] = state + + def clear_state(self, node_id: str) -> bool: + """Clear state for a specific node. Returns True if state was found and cleared.""" + if node_id in self._states: + del self._states[node_id] + return True + return False + + def clear_all_states(self): + """Clear all states""" + self._states.clear() + + def get_all_node_ids(self): + """Get list of all node IDs with state""" + return list(self._states.keys()) + + def cleanup_orphaned_states(self, active_node_ids: set[str]): + """Remove states for nodes that no longer exist""" + orphaned = [nid for nid in self._states.keys() if nid not in active_node_ids] + for nid in orphaned: + del self._states[nid] + return len(orphaned) class ControlNodeBase(ABC): """Base class for all control nodes""" @@ -57,6 +79,13 @@ class ControlNodeBase(ABC): def set_state(self, state): self.state_manager.set_state(self.node_id, state) + def cleanup(self): + """Clean up node state when node is destroyed""" + self.state_manager.clear_state(self.node_id) + + def __del__(self): + self.cleanup() + @abstractmethod def update(self, *args, **kwargs): """Main update method to be implemented by control nodes""" diff --git a/controls/mask_controls.py b/controls/mask_controls.py new file mode 100644 index 0000000..18e420a --- /dev/null +++ b/controls/mask_controls.py @@ -0,0 +1,244 @@ +import torch +import numpy as np +import cv2 +from ..base.control_base import ControlNodeBase + +class MaskControlMixin: + """Mixin providing common mask functionality""" + def create_circle_mask(self, height, width, center_y, center_x, size): + """Create a circular mask with anti-aliased edges""" + y, x = np.ogrid[:height, :width] + radius = size * min(height, width) / 2 + dist = np.sqrt((x - center_x)**2 + (y - center_y)**2) + mask = np.clip(radius + 1 - dist, 0, 1).astype(np.float32) + return mask + + def get_initial_state(self, x_pos, y_pos, size, min_size=None, max_size=None): + """Get initial state dictionary""" + state = { + "x": x_pos, + "y": y_pos, + "size": size, + "initialized": True + } + if min_size is not None: + state["min_size"] = min_size + if max_size is not None: + state["max_size"] = max_size + return state + +class RepulsiveMaskNode(ControlNodeBase, MaskControlMixin): + """Node that maintains a mask that repulses or attracts to input masks""" + @classmethod + def INPUT_TYPES(cls): + inputs = super().INPUT_TYPES() + inputs["required"].update({ + "input_mask": ("MASK",), + "x_pos": ("FLOAT", { + "default": 0.5, + "min": 0.0, + "max": 1.0, + "step": 0.01, + "tooltip": "Initial X position (0-1)" + }), + "y_pos": ("FLOAT", { + "default": 0.5, + "min": 0.0, + "max": 1.0, + "step": 0.01, + "tooltip": "Initial Y position (0-1)" + }), + "size": ("FLOAT", { + "default": 0.2, + "min": 0.01, + "max": 1.0, + "step": 0.01, + "tooltip": "Size of the mask relative to image size" + }), + "interaction_type": (["repulse", "attract"], { + "default": "repulse", + "tooltip": "Whether to move away from or toward input masks" + }), + "interaction_strength": ("FLOAT", { + "default": 0.1, + "min": 0.0, + "max": 1.0, + "step": 0.01, + "tooltip": "Strength of the interaction effect" + }) + }) + return inputs + + RETURN_TYPES = ("MASK",) + FUNCTION = "update" + CATEGORY = "real-time/control/mask" + + def interact(self, mask, input_mask_np, state, strength, height, width, is_attract): + if np.max(input_mask_np) < 0.1: + return state + + y_indices, x_indices = np.nonzero(input_mask_np) + if len(y_indices) == 0: + return state + + weights = input_mask_np[y_indices, x_indices] + com_y = np.average(y_indices, weights=weights) + com_x = np.average(x_indices, weights=weights) + + center_y = state["y"] * height + center_x = state["x"] * width + + dir_y = center_y - com_y + dir_x = center_x - com_x + + dist = np.sqrt(dir_x**2 + dir_y**2) + if dist < 1: + return state + + scale = min(1.0, 20.0/dist) + dir_y = (dir_y / dist) * scale + dir_x = (dir_x / dist) * scale + + # Flip direction if attracting + if is_attract: + dir_y *= -1 + dir_x *= -1 + + new_y = state["y"] + (dir_y * strength) + new_x = state["x"] + (dir_x * strength) + + state["y"] = np.clip(new_y, 0.0, 1.0) + state["x"] = np.clip(new_x, 0.0, 1.0) + + return state + + def update(self, input_mask, x_pos, y_pos, size, interaction_type, interaction_strength, always_execute=True): + was_reset = len(self.state_manager._states) == 0 + + if was_reset: + state = self.get_initial_state(x_pos, y_pos, size) + else: + state = self.get_state(self.get_initial_state(x_pos, y_pos, size)) + + batch_size, height, width = input_mask.shape + masks = np.zeros((batch_size, height, width), dtype=np.float32) + + for b in range(batch_size): + center_y, center_x = int(state["y"] * height), int(state["x"] * width) + mask = self.create_circle_mask(height, width, center_y, center_x, state["size"]) + + state = self.interact(mask, input_mask[b].cpu().numpy(), state, + interaction_strength, height, width, + interaction_type == "attract") + + center_y, center_x = int(state["y"] * height), int(state["x"] * width) + masks[b] = self.create_circle_mask(height, width, center_y, center_x, state["size"]) + + self.set_state(state) + return (torch.from_numpy(masks),) + +class ResizeMaskNode(ControlNodeBase, MaskControlMixin): + """Node that maintains a mask that resizes based on proximity to input masks""" + @classmethod + def INPUT_TYPES(cls): + inputs = super().INPUT_TYPES() + inputs["required"].update({ + "input_mask": ("MASK",), + "x_pos": ("FLOAT", { + "default": 0.5, + "min": 0.0, + "max": 1.0, + "step": 0.01, + "tooltip": "Initial X position (0-1)" + }), + "y_pos": ("FLOAT", { + "default": 0.5, + "min": 0.0, + "max": 1.0, + "step": 0.01, + "tooltip": "Initial Y position (0-1)" + }), + "size": ("FLOAT", { + "default": 0.2, + "min": 0.01, + "max": 1.0, + "step": 0.01, + "tooltip": "Size of the mask relative to image size" + }), + "min_size": ("FLOAT", { + "default": 0.1, + "min": 0.01, + "max": 1.0, + "step": 0.01, + "tooltip": "Minimum size the mask can shrink to" + }), + "max_size": ("FLOAT", { + "default": 0.4, + "min": 0.01, + "max": 1.0, + "step": 0.01, + "tooltip": "Maximum size the mask can grow to" + }), + "resize_strength": ("FLOAT", { + "default": 0.1, + "min": 0.0, + "max": 1.0, + "step": 0.01, + "tooltip": "Strength of the resize effect" + }) + }) + return inputs + + RETURN_TYPES = ("MASK",) + FUNCTION = "update" + CATEGORY = "real-time/control/mask" + + def resize(self, mask, input_mask_np, state, strength, height, width): + if np.max(input_mask_np) < 0.1: + # No mask detected - use max size + state["size"] = state["max_size"] + return state + + # Get distance from center to nearest mask pixel + center_y = int(state["y"] * height) + center_x = int(state["x"] * width) + y_indices, x_indices = np.nonzero(input_mask_np > 0.1) + + if len(y_indices) == 0: + return state + + distances = np.sqrt((y_indices - center_y)**2 + (x_indices - center_x)**2) + min_dist = np.min(distances) + + # Map distance directly to size + # When dist=0, size=min_size + # When dist=100, size=max_size + dist_range = 100.0 # pixels + size_range = state["max_size"] - state["min_size"] + size = state["min_size"] + (min_dist / dist_range) * size_range + state["size"] = np.clip(size, state["min_size"], state["max_size"]) + + return state + + def update(self, input_mask, x_pos, y_pos, size, min_size, max_size, resize_strength, always_execute=True): + was_reset = len(self.state_manager._states) == 0 + + if was_reset: + state = self.get_initial_state(x_pos, y_pos, size, min_size, max_size) + else: + state = self.get_state(self.get_initial_state(x_pos, y_pos, size, min_size, max_size)) + + batch_size, height, width = input_mask.shape + masks = np.zeros((batch_size, height, width), dtype=np.float32) + + for b in range(batch_size): + center_y, center_x = int(state["y"] * height), int(state["x"] * width) + mask = self.create_circle_mask(height, width, center_y, center_x, state["size"]) + + state = self.resize(mask, input_mask[b].cpu().numpy(), state, resize_strength, height, width) + + center_y, center_x = int(state["y"] * height), int(state["x"] * width) + masks[b] = self.create_circle_mask(height, width, center_y, center_x, state["size"]) + + self.set_state(state) + return (torch.from_numpy(masks),) \ No newline at end of file diff --git a/controls/state_management_controls.py b/controls/state_management_controls.py new file mode 100644 index 0000000..c15542b --- /dev/null +++ b/controls/state_management_controls.py @@ -0,0 +1,70 @@ +from ..base.control_base import ControlNodeBase + +class StateResetNode(ControlNodeBase): + """Node that resets all control node states when triggered""" + + OUTPUT_NODE = True + + @classmethod + def INPUT_TYPES(cls): + inputs = super().INPUT_TYPES() + inputs["required"].update({ + "trigger": ("BOOLEAN", { + "default": False, + "tooltip": "Set to True to reset all states" + }) + }) + return inputs + + RETURN_TYPES = ("BOOLEAN",) + FUNCTION = "update" + CATEGORY = "real-time/control/utility" + + def update(self, trigger, always_execute=True): + print(f"\n=== StateResetNode UPDATE - node_id: {self.node_id} ===") + print(f"States before potential reset: {self.state_manager._states}") + if trigger: + print("RESETTING ALL STATES") + self.state_manager.clear_all_states() + print(f"States after reset: {self.state_manager._states}") + return (True,) + return (False,) + +class StateTestNode(ControlNodeBase): + """Simple node that maintains a counter to test state management""" + + @classmethod + def INPUT_TYPES(cls): + inputs = super().INPUT_TYPES() + inputs["required"].update({ + "increment": ("INT", { + "default": 1, + "min": 1, + "max": 100, + "step": 1, + "tooltip": "Amount to increment counter by" + }) + }) + return inputs + + RETURN_TYPES = ("INT",) + FUNCTION = "update" + CATEGORY = "real-time/control/utility" + + def update(self, increment, always_execute=True): + print(f"\n=== StateTestNode UPDATE - node_id: {self.node_id} ===") + print(f"All states before get: {self.state_manager._states}") + + state = self.get_state({ + "counter": 0 + }) + print(f"Retrieved state: {state}") + + state["counter"] += increment + print(f"Updated state before save: {state}") + + self.set_state(state) + print(f"All states after save: {self.state_manager._states}") + + return (state["counter"],) + diff --git a/controls/utility_controls.py b/controls/utility_controls.py index 9b3f382..3ee9592 100644 --- a/controls/utility_controls.py +++ b/controls/utility_controls.py @@ -271,3 +271,5 @@ class LazyCondition(ControlNodeBase): + + diff --git a/media_pipe_nodes.py b/media_pipe_nodes.py new file mode 100644 index 0000000..b13bcee --- /dev/null +++ b/media_pipe_nodes.py @@ -0,0 +1,212 @@ +import torch +import numpy as np +import mediapipe as mp +import cv2 + +class HandTrackingNode: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "detection_confidence": ("FLOAT", { + "default": 0.5, + "min": 0.1, + "max": 1.0, + "step": 0.05 + }), + "tracking_confidence": ("FLOAT", { + "default": 0.5, + "min": 0.1, + "max": 1.0, + "step": 0.05 + }), + }, + "optional": { + "max_num_hands": ("INT", { + "default": 2, + "min": 1, + "max": 4, + "step": 1 + }), + "static_image_mode": ("BOOLEAN", { + "default": False + }), + "draw_debug": ("BOOLEAN", { + "default": False + }) + } + } + + RETURN_TYPES = ("IMAGE", "HAND_DATA") + RETURN_NAMES = ("debug_image", "hand_data") + FUNCTION = "track_hands" + CATEGORY = "image/processing" + + def __init__(self): + self.last_detection_conf = 0.5 + self.last_tracking_conf = 0.5 + self.last_max_hands = 2 + self.last_static_mode = True + + self.mp_hands = mp.solutions.hands + self.hands = self.mp_hands.Hands( + static_image_mode=self.last_static_mode, + max_num_hands=self.last_max_hands, + min_detection_confidence=self.last_detection_conf, + min_tracking_confidence=self.last_tracking_conf + ) + + def track_hands(self, image, detection_confidence, tracking_confidence, max_num_hands=2, static_image_mode=True, draw_debug=True): + # Check if we need to reinitialize MediaPipe hands + if (detection_confidence != self.last_detection_conf or + tracking_confidence != self.last_tracking_conf or + max_num_hands != self.last_max_hands or + static_image_mode != self.last_static_mode): + + self.hands = self.mp_hands.Hands( + static_image_mode=static_image_mode, + max_num_hands=max_num_hands, + min_detection_confidence=detection_confidence, + min_tracking_confidence=tracking_confidence + ) + self.last_detection_conf = detection_confidence + self.last_tracking_conf = tracking_confidence + self.last_max_hands = max_num_hands + self.last_static_mode = static_image_mode + + # Handle input + image_np = image.cpu().numpy() + image_np = (image_np * 255).astype(np.uint8) + batch_size = image_np.shape[0] + + # Process each image in batch + processed_images = [] + hand_data_batch = [] + + for i in range(batch_size): + single_image = image_np[i] + image_rgb = cv2.cvtColor(single_image, cv2.COLOR_BGR2RGB) + results = self.hands.process(image_rgb) + + hand_data = { + "left_hand": None, + "right_hand": None, + "num_hands": 0, + "hands_present": False + } + + if results.multi_hand_landmarks: + hand_data["num_hands"] = len(results.multi_hand_landmarks) + hand_data["hands_present"] = True + + # Create single debug image for all hands + debug_image = single_image.copy() if draw_debug else single_image + + # Process all detected hands + for hand_idx, (hand_landmarks, handedness) in enumerate(zip(results.multi_hand_landmarks, results.multi_handedness)): + hand_type = handedness.classification[0].label.lower() + landmarks = np.array([[lm.x, lm.y, lm.z] for lm in hand_landmarks.landmark]) + + # Store hand data based on type + if hand_type == "left" and hand_data["left_hand"] is None: + hand_data["left_hand"] = landmarks + elif hand_type == "right" and hand_data["right_hand"] is None: + hand_data["right_hand"] = landmarks + + # Draw landmarks for this hand on the shared debug image + if draw_debug: + mp.solutions.drawing_utils.draw_landmarks( + debug_image, + hand_landmarks, + self.mp_hands.HAND_CONNECTIONS + ) + + processed_images.append(debug_image) + else: + processed_images.append(single_image) + + hand_data_batch.append(hand_data) + + # Convert back to tensors + processed_batch = np.stack(processed_images, axis=0) + output_tensor = torch.from_numpy(processed_batch).float() / 255.0 + + return (output_tensor, hand_data_batch) + + def __del__(self): + if hasattr(self, 'hands'): + self.hands.close() + +class HandMaskNode: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "hand_data": ("HAND_DATA",), + "mask_type": (["full_hand", "palm_only", "fingers_only", "thumb", "index", "middle", "ring", "pinky"],), + } + } + + RETURN_TYPES = ("IMAGE", "MASK") + FUNCTION = "create_mask" + CATEGORY = "image/processing" + + def create_mask(self, image, hand_data, mask_type): + batch_size, height, width, channels = image.shape + masks = np.zeros((batch_size, height, width), dtype=np.float32) + + # Finger indices in MediaPipe hand landmarks + finger_indices = { + "thumb": [1,2,3,4], + "index": [5,6,7,8], + "middle": [9,10,11,12], + "ring": [13,14,15,16], + "pinky": [17,18,19,20] + } + palm_indices = [0,1,5,9,13,17] # Wrist and base of each finger + + for i in range(batch_size): + frame_hand_data = hand_data[i] + if not frame_hand_data["hands_present"]: + continue + + frame_mask = np.zeros((height, width), dtype=np.float32) + + for hand_key in ["right_hand", "left_hand"]: + hand = frame_hand_data[hand_key] + if hand is None: + continue + + points = hand[:, :2] # Take only x,y + points = (points * np.array([width, height])).astype(np.int32) + + if mask_type == "full_hand": + hull = cv2.convexHull(points) + cv2.fillConvexPoly(frame_mask, hull, 1.0) + + elif mask_type == "palm_only": + palm_points = points[palm_indices] + hull = cv2.convexHull(palm_points) + cv2.fillConvexPoly(frame_mask, hull, 1.0) + + elif mask_type == "fingers_only": + for finger in finger_indices.values(): + finger_points = points[finger] + cv2.fillConvexPoly(frame_mask, cv2.convexHull(finger_points), 1.0) + + elif mask_type in finger_indices: + finger_points = points[finger_indices[mask_type]] + cv2.fillConvexPoly(frame_mask, cv2.convexHull(finger_points), 1.0) + + masks[i] = frame_mask + + # # Debug output + # if frame_hand_data["right_hand"] is not None: + # print(f"Right hand points range: x({np.min(points[:, 0]):.1f}-{np.max(points[:, 0]):.1f}), y({np.min(points[:, 1]):.1f}-{np.max(points[:, 1]):.1f})") + # if frame_hand_data["left_hand"] is not None: + # print(f"Left hand points range: x({np.min(points[:, 0]):.1f}-{np.max(points[:, 0]):.1f}), y({np.min(points[:, 1]):.1f}-{np.max(points[:, 1]):.1f})") + + mask_tensor = torch.from_numpy(masks) + return (image, mask_tensor) \ No newline at end of file