hand tracking, state management
This commit includes updates to global state management, and a state reset node. It includes new tracking nodes powered by media pipe. Includes an interactive mask node.
This commit is contained in:
RyanOnTheInside
@@ -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),)
|
||||
@@ -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"],)
|
||||
|
||||
@@ -271,3 +271,5 @@ class LazyCondition(ControlNodeBase):
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user