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
@@ -27,6 +27,7 @@ ENV/
|
||||
# IDE files
|
||||
.idea/
|
||||
.vscode/
|
||||
.cursor/
|
||||
*.swp
|
||||
*.swo
|
||||
|
||||
|
||||
+24
-3
@@ -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 = " 🕒🅡🅣🅝"
|
||||
|
||||
+30
-1
@@ -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"""
|
||||
|
||||
@@ -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):
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
@@ -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)
|
||||
Reference in New Issue
Block a user