mirror of
https://github.com/luckyrobots/open_phantom.git
synced 2025-04-04 02:52:18 +00:00
524 lines
22 KiB
Python
524 lines
22 KiB
Python
import os
|
|
import cv2
|
|
import time
|
|
import torch
|
|
import numpy as np
|
|
import open3d as o3d
|
|
import mediapipe as mp
|
|
|
|
from tqdm import tqdm
|
|
from utils.visualizations import *
|
|
from phantom.robot_manager import RobotManager
|
|
# from diffusers import StableDiffusionInpaintPipeline
|
|
|
|
|
|
class HandProcessor:
|
|
def __init__(self) -> None:
|
|
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
|
|
self.mp_hands = mp.solutions.hands
|
|
self.mp_drawing = mp.solutions.drawing_utils
|
|
self.mp_drawing_styles = mp.solutions.drawing_styles
|
|
|
|
self.hands = self.mp_hands.Hands(
|
|
model_complexity=1,
|
|
max_num_hands=1,
|
|
static_image_mode=False,
|
|
)
|
|
|
|
# NOTE: Look into better depth estimation models
|
|
# Initialize MiDaS for depth estimation
|
|
print("Loading MiDaS model...")
|
|
self.midas = torch.hub.load("intel-isl/MiDaS", "DPT_Hybrid")
|
|
self.midas.to(self.device)
|
|
self.midas.eval()
|
|
|
|
# Load MiDaS transforms to resize and normalize input images
|
|
self.midas_transforms = torch.hub.load("intel-isl/MiDaS", "transforms")
|
|
self.transform = self.midas_transforms.dpt_transform
|
|
|
|
# Segment hand from image using MediaPipe Hands landmarks
|
|
def _create_mask(self, image: np.ndarray, landmarks: list, thickness: int=5, padding:int=10) -> np.ndarray:
|
|
height, width = image.shape[:2]
|
|
mask = np.zeros((height, width), dtype=np.uint8)
|
|
|
|
points = [(int(landmark.x * width), int(landmark.y * height)) for landmark in landmarks.landmark]
|
|
|
|
# Draw filled circles at each landmark
|
|
for point in points:
|
|
cv2.circle(mask, point, thickness, 255, -1)
|
|
|
|
# Connect all landmarks with thick lines
|
|
for connection in self.mp_hands.HAND_CONNECTIONS:
|
|
start_idx, end_idx = connection
|
|
cv2.line(mask, points[start_idx], points[end_idx], 255, thickness)
|
|
|
|
# Create palm by connecting base of fingers with wrist
|
|
palm_points = [points[0], points[1], points[5], points[9], points[13], points[17]]
|
|
cv2.fillPoly(mask, [np.array(palm_points)], 255)
|
|
|
|
# Create shape between fingers
|
|
finger_bases = [(1,5), (5,9), (9,13), (13,17)]
|
|
for base1, base2 in finger_bases:
|
|
triangle = np.array([points[0], points[base1], points[base2]])
|
|
cv2.fillPoly(mask, [triangle], 255)
|
|
|
|
# Dilate to smooth and expand slightly
|
|
kernel = np.ones((padding, padding), np.uint8)
|
|
mask = cv2.dilate(mask, kernel, iterations=1)
|
|
|
|
# Smooth the mask
|
|
mask = cv2.GaussianBlur(mask, (21, 21), 0)
|
|
_, mask = cv2.threshold(mask, 50, 255, cv2.THRESH_BINARY)
|
|
|
|
return mask
|
|
|
|
"""
|
|
Transform input image to match MiDaS model requirements
|
|
Estimate depth map using MiDaS model
|
|
"""
|
|
# TODO: Look into why screen goes black sometimes
|
|
def _estimate_depth(self, image: np.ndarray) -> tuple:
|
|
img = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
|
|
|
|
# Transform img for MiDaS model
|
|
input_batch = self.transform(img).to(self.device)
|
|
|
|
with torch.inference_mode():
|
|
prediction = self.midas(input_batch)
|
|
|
|
# Interpolate to original size
|
|
prediction = torch.nn.functional.interpolate(
|
|
prediction.unsqueeze(1),
|
|
size=image.shape[:2],
|
|
mode="bicubic",
|
|
align_corners=False,
|
|
).squeeze()
|
|
|
|
# Convert to numpy and normalize to 0-255 for visualization
|
|
depth_map = prediction.cpu().numpy()
|
|
depth_min = depth_map.min()
|
|
depth_max = depth_map.max()
|
|
depth_map_normalized = 255 * (depth_map - depth_min) / (depth_max - depth_min)
|
|
|
|
return depth_map, depth_map_normalized.astype(np.uint8)
|
|
|
|
# TODO: Look into better depth scaling
|
|
def _create_cloud(self, depth_map: np.ndarray, hand_mask: np.ndarray) -> np.ndarray:
|
|
focal_x, focal_y, center_x, center_y = self.camera_intrinsics
|
|
|
|
v_coords, u_coords = np.where(hand_mask > 0)
|
|
z_values = depth_map[v_coords, u_coords]
|
|
|
|
# Filter out zero depth values
|
|
valid_indices = z_values > 0
|
|
u_coords = u_coords[valid_indices]
|
|
v_coords = v_coords[valid_indices]
|
|
z_values = z_values[valid_indices]
|
|
|
|
# NOTE: Abritrary scaling factor for depth
|
|
z_metric = z_values * 0.5
|
|
|
|
# Back-project to 3D using camera intrinsics
|
|
x_values = (u_coords - center_x) * z_metric / focal_x
|
|
y_values = (v_coords - center_y) * z_metric / focal_y
|
|
|
|
points = np.column_stack((x_values, y_values, z_metric))
|
|
|
|
return points
|
|
|
|
# TODO: Look into better Z scaling
|
|
def _create_mesh(self, landmarks: list, image_dims: tuple) -> np.ndarray:
|
|
width, height = image_dims
|
|
|
|
vertices = []
|
|
for landmark in landmarks.landmark:
|
|
vertices.append([
|
|
landmark.x * width,
|
|
landmark.y * height,
|
|
landmark.z * width
|
|
])
|
|
|
|
# Define faces (triangles) connecting landmarks
|
|
faces = [
|
|
# Palm
|
|
[0, 1, 5], [0, 5, 9], [0, 9, 13], [0, 13, 17],
|
|
# Thumb
|
|
[1, 2, 3], [2, 3, 4],
|
|
# Index finger
|
|
[5, 6, 7], [6, 7, 8],
|
|
# Middle finger
|
|
[9, 10, 11], [10, 11, 12],
|
|
# Ring finger
|
|
[13, 14, 15], [14, 15, 16],
|
|
# Pinky
|
|
[17, 18, 19], [18, 19, 20]
|
|
]
|
|
|
|
dense_vertices = list(vertices)
|
|
|
|
# Add interpolated vertices along finger segments
|
|
connections = self.mp_hands.HAND_CONNECTIONS
|
|
for connection in connections:
|
|
start_idx, end_idx = connection
|
|
start_point = np.array(vertices[start_idx])
|
|
end_point = np.array(vertices[end_idx])
|
|
|
|
# Add 2 interpolated points between each connected pair
|
|
for t in [0.33, 0.66]:
|
|
interp_point = start_point * (1-t) + end_point * t
|
|
dense_vertices.append(interp_point.tolist())
|
|
|
|
return np.array(dense_vertices), np.array(faces)
|
|
|
|
def _icp_registration(self, cloud: np.ndarray, vertices: np.ndarray) -> np.ndarray:
|
|
# Normalize both point clouds to handle scale differences
|
|
cloud_centroid = np.mean(cloud, axis=0)
|
|
vertices_centroid = np.mean(vertices, axis=0)
|
|
|
|
cloud_centered = cloud - cloud_centroid
|
|
vertices_centered = vertices - vertices_centroid
|
|
|
|
cloud_scale = np.max(np.linalg.norm(cloud_centered, axis=1))
|
|
vertices_scale = np.max(np.linalg.norm(vertices_centered, axis=1))
|
|
|
|
cloud_normalized = cloud_centered / cloud_scale
|
|
vertices_normalized = vertices_centered / vertices_scale
|
|
|
|
# Create Open3D point clouds
|
|
source = o3d.geometry.PointCloud()
|
|
source.points = o3d.utility.Vector3dVector(vertices_normalized)
|
|
|
|
target = o3d.geometry.PointCloud()
|
|
target.points = o3d.utility.Vector3dVector(cloud_normalized)
|
|
|
|
# Optional: Downsample for better performance
|
|
source = source.voxel_down_sample(voxel_size=0.01)
|
|
target = target.voxel_down_sample(voxel_size=0.01)
|
|
|
|
# Initial identity transformation
|
|
trans_init = np.eye(4)
|
|
|
|
# Run ICP with better parameters
|
|
result = o3d.pipelines.registration.registration_icp(
|
|
source,
|
|
target,
|
|
max_correspondence_distance=0.2, # Increased for better matching
|
|
init=trans_init,
|
|
estimation_method=o3d.pipelines.registration.TransformationEstimationPointToPlane(),
|
|
criteria=o3d.pipelines.registration.ICPConvergenceCriteria(max_iteration=100)
|
|
)
|
|
|
|
# Create denormalization transformation
|
|
denorm_transform = np.eye(4)
|
|
denorm_transform[:3, :3] *= (cloud_scale / vertices_scale)
|
|
|
|
# Combine transformations
|
|
transform_to_origin = np.eye(4)
|
|
transform_to_origin[:3, 3] = -vertices_centroid
|
|
|
|
transform_from_origin = np.eye(4)
|
|
transform_from_origin[:3, 3] = cloud_centroid
|
|
|
|
final_transform = np.matmul(transform_from_origin,
|
|
np.matmul(denorm_transform,
|
|
np.matmul(result.transformation, transform_to_origin)))
|
|
|
|
return final_transform
|
|
|
|
def _refine_landmarks(self, landmarks: list, transform: int, image_dims: tuple):
|
|
width, height = image_dims
|
|
|
|
refined_landmarks = []
|
|
for landmark in landmarks.landmark:
|
|
point = np.array([
|
|
landmark.x * width,
|
|
landmark.y * height,
|
|
landmark.z * width, # TODO: Figure out better scaling factor
|
|
1.0
|
|
])
|
|
|
|
# Apply transformation to 3D point
|
|
transformed = np.dot(transform, point)
|
|
refined_landmarks.append(transformed[:3])
|
|
|
|
return refined_landmarks
|
|
|
|
def _apply_constraints(self, landmarks: list):
|
|
constrained = np.array(landmarks)
|
|
|
|
# Define finger joint indices
|
|
# MediaPipe hand model: Wrist is 0, thumb is 1-4, index is 5-8, etc.
|
|
thumb_indices = [1, 2, 3, 4]
|
|
index_indices = [5, 6, 7, 8]
|
|
|
|
# Constrain the last two joints of thumb and index finger
|
|
# as mentioned in the paper
|
|
for finger_indices in [thumb_indices, index_indices]:
|
|
# Get the last three joints (two segments)
|
|
if len(finger_indices) >= 3:
|
|
# Get joint positions
|
|
joint1 = constrained[finger_indices[-3]]
|
|
joint2 = constrained[finger_indices[-2]]
|
|
joint3 = constrained[finger_indices[-1]]
|
|
|
|
# Direction of the first segment
|
|
dir1 = joint2 - joint1
|
|
dir1 = dir1 / np.linalg.norm(dir1)
|
|
|
|
# Instead of full ball joint, constrain the last joint's direction
|
|
# to follow similar direction as the previous segment
|
|
ideal_dir = dir1.copy()
|
|
actual_dir = joint3 - joint2
|
|
actual_length = np.linalg.norm(actual_dir)
|
|
|
|
# Force the direction to be within a cone of the previous segment
|
|
# (limiting to single degree of freedom approximately)
|
|
corrected_dir = ideal_dir * actual_length
|
|
|
|
# Apply the correction
|
|
constrained[finger_indices[-1]] = joint2 + corrected_dir
|
|
|
|
return constrained
|
|
|
|
def _get_robot_params(self, refined_landmarks) -> tuple:
|
|
# Extract key landmarks
|
|
wrist = refined_landmarks[0] # Wrist landmark
|
|
thumb_tip = refined_landmarks[4] # Thumb tip
|
|
index_tip = refined_landmarks[8] # Index finger tip
|
|
|
|
# Calculate end effector position (midpoint between thumb and index tips)
|
|
position = (thumb_tip + index_tip) / 2
|
|
|
|
# Calculate vectors for orientation
|
|
v1 = thumb_tip - wrist # Vector from wrist to thumb tip
|
|
v2 = index_tip - wrist # Vector from wrist to index tip
|
|
|
|
# Calculate normal to hand plane
|
|
normal = np.cross(v1, v2)
|
|
if np.linalg.norm(normal) > 0:
|
|
normal = normal / np.linalg.norm(normal)
|
|
else:
|
|
# Default if vectors are collinear
|
|
normal = np.array([0, 0, 1])
|
|
|
|
# Calculate main direction along thumb
|
|
direction = thumb_tip - wrist
|
|
if np.linalg.norm(direction) > 0:
|
|
direction = direction / np.linalg.norm(direction)
|
|
else:
|
|
# Default if thumb is at wrist (unlikely)
|
|
direction = np.array([1, 0, 0])
|
|
|
|
# Calculate gripper width
|
|
gripper_width = np.linalg.norm(thumb_tip - index_tip)
|
|
|
|
return position, direction, normal, gripper_width
|
|
|
|
def record_video(self) -> str:
|
|
output_dir = os.path.join(os.getcwd(), 'recordings')
|
|
os.makedirs(output_dir, exist_ok=True)
|
|
|
|
timestamp = time.strftime("%Y%m%d-%H%M%S")
|
|
video_filename = os.path.join(output_dir, f"recorded_video_{timestamp}.mp4")
|
|
|
|
cap = cv2.VideoCapture(0)
|
|
|
|
assert cap.isOpened(), "Failed to open camera."
|
|
|
|
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
|
|
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
|
|
fps = int(cap.get(cv2.CAP_PROP_FPS))
|
|
|
|
# Update camera intrinsics based on video dimensions
|
|
self.camera_intrinsics = (width * 0.8, height * 0.8, width / 2, height / 2)
|
|
|
|
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
|
|
out = None
|
|
|
|
recording = False
|
|
did_record = False
|
|
|
|
print("Camera is active. Press 'r' to start/stop recording, 'q' to quit.")
|
|
|
|
while cap.isOpened():
|
|
success, frame = cap.read()
|
|
assert success, "Failed to read from camera."
|
|
|
|
# Mirror the image for more intuitive viewing
|
|
frame = cv2.flip(frame, 1)
|
|
|
|
# Create a separate display frame for showing the recording indicator
|
|
display_frame = frame.copy()
|
|
|
|
# Display recording indicator ONLY on the display frame
|
|
if recording:
|
|
cv2.circle(display_frame, (30, 30), 15, (0, 0, 255), -1)
|
|
cv2.putText(display_frame, "RECORDING", (50, 40), cv2.FONT_HERSHEY_SIMPLEX,
|
|
1, (0, 0, 255), 2, cv2.LINE_AA)
|
|
|
|
# Show the display frame (with indicator if recording)
|
|
cv2.imshow('Video Recording', display_frame)
|
|
|
|
# Write the original frame (without indicator) to video file if recording
|
|
if recording and out is not None:
|
|
out.write(frame)
|
|
|
|
key = cv2.waitKey(1) & 0xFF
|
|
|
|
# NOTE: Haven't tested what happens if user records multiple videos in one session (probably overwrites?)
|
|
if key == ord('q'):
|
|
break
|
|
elif key == ord('r'):
|
|
recording = not recording
|
|
if recording:
|
|
print(f"Started recording to {video_filename}")
|
|
out = cv2.VideoWriter(video_filename, fourcc, fps, (width, height))
|
|
else:
|
|
if out is not None:
|
|
out.release()
|
|
print(f"Stopped recording. Video saved to {video_filename}")
|
|
did_record = True
|
|
|
|
if out is not None:
|
|
out.release()
|
|
cap.release()
|
|
cv2.destroyAllWindows()
|
|
|
|
return video_filename if did_record else None
|
|
|
|
def process_video(self, video_path: str) -> tuple:
|
|
assert video_path, "Video path is required."
|
|
|
|
cap = cv2.VideoCapture(video_path)
|
|
|
|
assert cap.isOpened(), f"Failed to open video file {video_path}"
|
|
|
|
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
|
|
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
|
|
fps = int(cap.get(cv2.CAP_PROP_FPS))
|
|
total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
|
|
|
|
# Update camera intrinsics based on video dimensions
|
|
self.camera_intrinsics = (width * 0.8, height * 0.8, width / 2, height / 2)
|
|
|
|
base_path = os.path.splitext(video_path)[0]
|
|
segmented_output_path = f"{base_path}_masked.mp4"
|
|
depth_output_path = f"{base_path}_depth.mp4"
|
|
mesh_output_path = f"{base_path}_mesh.mp4"
|
|
registration_output_path = f"{base_path}_registration.mp4"
|
|
constraints_output_path = f"{base_path}_constraints.mp4"
|
|
robot_output_path = f"{base_path}_robot.mp4"
|
|
|
|
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
|
|
segmented_out = cv2.VideoWriter(segmented_output_path, fourcc, fps, (width, height))
|
|
depth_out = cv2.VideoWriter(depth_output_path, fourcc, fps, (width, height))
|
|
mesh_out = cv2.VideoWriter(mesh_output_path, fourcc, fps, (width, height))
|
|
reg_out = cv2.VideoWriter(registration_output_path, fourcc, fps, (640, 480)) # Fixed size
|
|
constraints_out = cv2.VideoWriter(constraints_output_path, fourcc, fps, (width, height))
|
|
robot_out = cv2.VideoWriter(robot_output_path, fourcc, fps, (width, height))
|
|
|
|
print(f"Processing video with {total_frames} frames...")
|
|
for _ in tqdm(range(total_frames)):
|
|
success, frame = cap.read()
|
|
assert success, "Failed to read frame from video."
|
|
|
|
# Convert image to RGB for MediaPipe
|
|
image_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
|
|
|
|
# Process with MediaPipe Hands
|
|
results = self.hands.process(image_rgb)
|
|
|
|
# Initialize output frames
|
|
segmented_frame = frame.copy()
|
|
depth_frame = np.zeros((height, width, 3), dtype=np.uint8)
|
|
mesh_frame = frame.copy()
|
|
reg_frame = np.ones((480, 640, 3), dtype=np.uint8) * 255 # Fixed size, white background
|
|
constraints_frame = frame.copy()
|
|
robot_frame = frame.copy()
|
|
|
|
# Process if hand is detected
|
|
if results.multi_hand_landmarks:
|
|
for hand_landmarks in results.multi_hand_landmarks:
|
|
# Segment hand
|
|
hand_mask = self._create_mask(frame, hand_landmarks)
|
|
mask_overlay = frame.copy()
|
|
mask_overlay[hand_mask > 0] = [0, 0, 255] # Red color for mask
|
|
segmented_frame = cv2.addWeighted(frame, 0.7, mask_overlay, 0.3, 0)
|
|
|
|
# Depth estimation
|
|
depth_map, depth_vis = self._estimate_depth(frame)
|
|
depth_colored = cv2.applyColorMap(depth_vis, cv2.COLORMAP_JET)
|
|
depth_frame = depth_colored.copy()
|
|
|
|
# Create hand mesh
|
|
hand_vertices, hand_faces = self._create_mesh(hand_landmarks, (width, height))
|
|
mesh_frame = visualize_mesh(frame, hand_vertices, hand_faces)
|
|
|
|
# Create point cloud from depth
|
|
cloud = self._create_cloud(depth_map, hand_mask)
|
|
|
|
# Perform ICP registration
|
|
icp_transform = self._icp_registration(cloud, hand_vertices)
|
|
reg_frame = visualize_registration(cloud, hand_vertices, icp_transform)
|
|
|
|
# Refine landmarks using ICP transformation
|
|
refined_landmarks = self._refine_landmarks(hand_landmarks, icp_transform, (width, height))
|
|
|
|
# Store pre-constraint landmarks for visualization
|
|
original_refined = refined_landmarks.copy()
|
|
|
|
# Apply anatomical constraints
|
|
constrained_landmarks = self._apply_constraints(refined_landmarks)
|
|
constraints_frame = visualize_constraints(frame, original_refined, constrained_landmarks, self.camera_intrinsics)
|
|
|
|
# # Calculate robot parameters
|
|
# position, direction, normal, gripper_width = self._get_robot_params(constrained_landmarks)
|
|
# robot_frame = visualize_robot_params(frame, position, direction, normal, gripper_width)
|
|
|
|
segmented_out.write(segmented_frame)
|
|
depth_out.write(depth_frame)
|
|
mesh_out.write(mesh_frame)
|
|
reg_out.write(reg_frame)
|
|
constraints_out.write(constraints_frame)
|
|
robot_out.write(robot_frame)
|
|
|
|
display_scale = 0.5
|
|
display_size = (int(width * display_scale), int(height * display_scale))
|
|
reg_display_size = (int(640 * display_scale), int(480 * display_scale))
|
|
|
|
cv2.imshow('Segmented', cv2.resize(segmented_frame, display_size))
|
|
cv2.imshow('Depth', cv2.resize(depth_frame, display_size))
|
|
cv2.imshow('Mesh', cv2.resize(mesh_frame, display_size))
|
|
cv2.imshow('Registration', cv2.resize(reg_frame, reg_display_size))
|
|
cv2.imshow('Constraints', cv2.resize(constraints_frame, display_size))
|
|
cv2.imshow('Robot Parameters', cv2.resize(robot_frame, display_size))
|
|
|
|
if cv2.waitKey(1) & 0xFF == 27: # ESC to exit
|
|
break
|
|
|
|
cap.release()
|
|
segmented_out.release()
|
|
depth_out.release()
|
|
mesh_out.release()
|
|
reg_out.release()
|
|
constraints_out.release()
|
|
robot_out.release()
|
|
cv2.destroyAllWindows()
|
|
|
|
print(f"Processing complete. Results saved to:")
|
|
print(f"- Hand mask: {segmented_output_path}")
|
|
print(f"- Depth visualization: {depth_output_path}")
|
|
print(f"- Mesh visualization: {mesh_output_path}")
|
|
print(f"- Registration visualization: {registration_output_path}")
|
|
print(f"- Constraints visualization: {constraints_output_path}")
|
|
print(f"- Robot parameters: {robot_output_path}")
|
|
|
|
return {
|
|
"segmented": segmented_output_path,
|
|
"depth": depth_output_path,
|
|
"mesh": mesh_output_path,
|
|
"registration": registration_output_path,
|
|
"constraints": constraints_output_path,
|
|
"robot": robot_output_path
|
|
} |