RyanOnTheInside
9a9754de3c
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.
212 lines
8.2 KiB
Python
212 lines
8.2 KiB
Python
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) |