1681 lines
74 KiB
Python
1681 lines
74 KiB
Python
"""
|
||
Improved HaMeR pipeline with 3-pass architecture using YOLO hand detector:
|
||
Pass 1: Extract all raw bboxes (YOLO → Hand bboxes directly)
|
||
Pass 2: Clean bbox sequences globally (detect anomalies, smooth, interpolate)
|
||
Pass 3: Run HaMeR on cleaned bboxes
|
||
|
||
This uses YOLO hand detector (like WiLoR) instead of Detectron2 + ViTPose.
|
||
"""
|
||
|
||
from pathlib import Path
|
||
import torch
|
||
import argparse
|
||
import os
|
||
import sys
|
||
import cv2
|
||
import numpy as np
|
||
import pickle
|
||
from typing import Dict, List, Tuple, Optional
|
||
from tqdm import tqdm
|
||
import json
|
||
import time
|
||
import imageio
|
||
from hamer.configs import CACHE_DIR_HAMER
|
||
from hamer.models import HAMER, download_models, load_hamer
|
||
from hamer.utils import recursive_to
|
||
from hamer.utils.geometry import aa_to_rotmat, perspective_projection
|
||
from hamer.datasets.vitdet_dataset import ViTDetDataset, DEFAULT_MEAN, DEFAULT_STD
|
||
from hamer.utils.renderer import Renderer, cam_crop_to_full
|
||
from ultralytics import YOLO
|
||
import hamer
|
||
|
||
LIGHT_BLUE = (0.65098039, 0.74117647, 0.85882353)
|
||
|
||
openpose_indices = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20]
|
||
gt_indices = openpose_indices
|
||
|
||
|
||
def compute_bbox_from_keypoints(keypoints, img_shape, scale_factor=1.0, max_size_ratio=0.4):
|
||
"""Compute tight bbox from keypoints. Scaling is applied later when creating HaMeR input patches."""
|
||
valid = keypoints[:, 2] > 0.5
|
||
if valid.sum() < 10:
|
||
return None
|
||
|
||
valid_kp = keypoints[valid, :2]
|
||
x_min, y_min = valid_kp.min(axis=0)
|
||
x_max, y_max = valid_kp.max(axis=0)
|
||
|
||
width = x_max - x_min
|
||
height = y_max - y_min
|
||
cx = (x_min + x_max) / 2
|
||
cy = (y_min + y_max) / 2
|
||
|
||
width *= scale_factor
|
||
height *= scale_factor
|
||
|
||
img_h, img_w = img_shape[:2]
|
||
max_width = img_w * max_size_ratio
|
||
max_height = img_h * max_size_ratio
|
||
|
||
width = min(width, max_width)
|
||
height = min(height, max_height)
|
||
size = max(width, height)
|
||
|
||
x_min = cx - size / 2
|
||
y_min = cy - size / 2
|
||
x_max = cx + size / 2
|
||
y_max = cy + size / 2
|
||
|
||
return np.array([x_min, y_min, x_max, y_max])
|
||
|
||
|
||
def compute_iou(bbox1, bbox2):
|
||
"""Compute IoU between two bboxes."""
|
||
x1_min, y1_min, x1_max, y1_max = bbox1
|
||
x2_min, y2_min, x2_max, y2_max = bbox2
|
||
|
||
inter_x_min = max(x1_min, x2_min)
|
||
inter_y_min = max(y1_min, y2_min)
|
||
inter_x_max = min(x1_max, x2_max)
|
||
inter_y_max = min(y1_max, y2_max)
|
||
|
||
inter_area = max(0, inter_x_max - inter_x_min) * max(0, inter_y_max - inter_y_min)
|
||
bbox1_area = (x1_max - x1_min) * (y1_max - y1_min)
|
||
bbox2_area = (x2_max - x2_min) * (y2_max - y2_min)
|
||
union_area = bbox1_area + bbox2_area - inter_area
|
||
|
||
return inter_area / union_area if union_area > 0 else 0
|
||
|
||
|
||
def compute_containment_ratio(bbox1, bbox2):
|
||
"""
|
||
Compute how much bbox1 is contained within bbox2.
|
||
Returns the ratio of bbox1's area that overlaps with bbox2.
|
||
"""
|
||
x1_min, y1_min, x1_max, y1_max = bbox1
|
||
x2_min, y2_min, x2_max, y2_max = bbox2
|
||
|
||
inter_x_min = max(x1_min, x2_min)
|
||
inter_y_min = max(y1_min, y2_min)
|
||
inter_x_max = min(x1_max, x2_max)
|
||
inter_y_max = min(y1_max, y2_max)
|
||
|
||
inter_area = max(0, inter_x_max - inter_x_min) * max(0, inter_y_max - inter_y_min)
|
||
bbox1_area = (x1_max - x1_min) * (y1_max - y1_min)
|
||
|
||
return inter_area / bbox1_area if bbox1_area > 0 else 0
|
||
|
||
|
||
def create_video_from_images(image_folder, output_video_path, fps=30):
|
||
"""Create MP4 video from images in a folder using imageio (same as src_cam video)."""
|
||
image_folder = Path(image_folder)
|
||
if not image_folder.exists():
|
||
print(f"Warning: {image_folder} does not exist, skipping video creation")
|
||
return
|
||
|
||
# Get all jpg images sorted by filename
|
||
image_files = sorted(list(image_folder.glob('*.jpg')))
|
||
if len(image_files) == 0:
|
||
print(f"Warning: No images found in {image_folder}, skipping video creation")
|
||
return
|
||
|
||
print(f"Creating video: {output_video_path}")
|
||
|
||
# Load all images
|
||
frames = []
|
||
for img_file in tqdm(image_files, desc=f" Loading frames"):
|
||
img = imageio.imread(str(img_file))
|
||
if img is not None:
|
||
frames.append(img)
|
||
|
||
if len(frames) == 0:
|
||
print(f"Warning: Could not load any images from {image_folder}")
|
||
return
|
||
|
||
# Write video using imageio (same method as src_cam video generation)
|
||
imageio.mimwrite(str(output_video_path), frames, fps=fps)
|
||
print(f" Video saved: {output_video_path} ({len(frames)} frames @ {fps} fps)")
|
||
|
||
|
||
# ============================================================================
|
||
# PASS 1: Extract all raw bboxes using YOLO hand detector
|
||
# ============================================================================
|
||
|
||
def extract_raw_bboxes(img_paths, detector, vis_dir=None):
|
||
"""
|
||
Pass 1: Extract raw hand bboxes using YOLO hand detector (like WiLoR).
|
||
YOLO directly detects hand bboxes with left/right classification.
|
||
No ViTPose needed - much simpler and more reliable!
|
||
|
||
Args:
|
||
img_paths: List of image paths
|
||
detector: YOLO hand detector model
|
||
vis_dir: Optional directory for visualizations
|
||
|
||
Returns:
|
||
raw_data: List of dicts for each frame with bbox data (no keypoints from YOLO)
|
||
"""
|
||
print("\n" + "="*80)
|
||
print("PASS 1: Extracting raw hand bboxes (YOLO hand detector)")
|
||
print("="*80)
|
||
|
||
if vis_dir is not None:
|
||
pass1_dir = os.path.join(vis_dir, 'pass1_raw_bboxes')
|
||
os.makedirs(pass1_dir, exist_ok=True)
|
||
|
||
raw_data = []
|
||
|
||
for frame_idx, img_path in enumerate(tqdm(sorted(img_paths), desc="Pass 1: Extracting hand bboxes")):
|
||
img_path = str(img_path)
|
||
img_cv2 = cv2.imread(img_path)
|
||
|
||
frame_data = {
|
||
'frame_idx': frame_idx,
|
||
'img_path': img_path,
|
||
'left_bbox': None,
|
||
'right_bbox': None,
|
||
'left_keypoints': None, # Will be None for YOLO version
|
||
'right_keypoints': None, # Will be None for YOLO version
|
||
'left_conf': 0.0,
|
||
'right_conf': 0.0
|
||
}
|
||
|
||
# YOLO direct hand detection (following WiLoR demo.py)
|
||
# Class 0 = left hand, Class 1 = right hand
|
||
detections = detector(img_cv2, conf=0.5, verbose=False)[0]
|
||
|
||
# Track best detection for each hand (in case of duplicates)
|
||
left_detections = []
|
||
right_detections = []
|
||
|
||
# Helper function to expand bbox by a given scale (e.g., 1.2x) and make it SQUARE
|
||
def enlarge_bbox(bbox, scale=1.2, img_shape=None):
|
||
x1, y1, x2, y2 = bbox
|
||
w = x2 - x1
|
||
h = y2 - y1
|
||
cx = x1 + w / 2.0
|
||
cy = y1 + h / 2.0
|
||
|
||
# Make it square by taking the max dimension
|
||
size = max(w, h) * scale
|
||
|
||
new_x1 = cx - size / 2.0
|
||
new_y1 = cy - size / 2.0
|
||
new_x2 = cx + size / 2.0
|
||
new_y2 = cy + size / 2.0
|
||
|
||
# Optionally clip to image boundaries
|
||
if img_shape is not None:
|
||
H, W = img_shape[:2]
|
||
new_x1 = max(0, min(W-1, new_x1))
|
||
new_y1 = max(0, min(H-1, new_y1))
|
||
new_x2 = max(0, min(W-1, new_x2))
|
||
new_y2 = max(0, min(H-1, new_y2))
|
||
return np.array([new_x1, new_y1, new_x2, new_y2], dtype=bbox.dtype)
|
||
|
||
img_shape = img_cv2.shape
|
||
for det in detections:
|
||
bbox = det.boxes.xyxy[0].cpu().numpy() # x1, y1, x2, y2
|
||
conf = float(det.boxes.conf[0])
|
||
cls = int(det.boxes.cls[0])
|
||
|
||
# Enlarge bbox by 1.2x, clamp to image
|
||
bbox = enlarge_bbox(bbox, scale=1.2, img_shape=img_shape)
|
||
|
||
if cls == 0: # left hand
|
||
left_detections.append({'bbox': bbox, 'conf': conf})
|
||
elif cls == 1: # right hand
|
||
right_detections.append({'bbox': bbox, 'conf': conf})
|
||
|
||
# Keep highest confidence detection for each hand
|
||
if len(left_detections) > 0:
|
||
best_left = max(left_detections, key=lambda x: x['conf'])
|
||
frame_data['left_bbox'] = best_left['bbox']
|
||
frame_data['left_conf'] = best_left['conf']
|
||
|
||
if len(right_detections) > 0:
|
||
best_right = max(right_detections, key=lambda x: x['conf'])
|
||
frame_data['right_bbox'] = best_right['bbox']
|
||
frame_data['right_conf'] = best_right['conf']
|
||
|
||
# Visualize raw bboxes
|
||
if vis_dir is not None:
|
||
img_vis = img_cv2.copy()
|
||
if frame_data['left_bbox'] is not None:
|
||
x1, y1, x2, y2 = frame_data['left_bbox'].astype(int)
|
||
cv2.rectangle(img_vis, (x1, y1), (x2, y2), (255, 0, 0), 2) # Blue for left
|
||
cv2.putText(img_vis, f"L (conf={frame_data['left_conf']:.2f})",
|
||
(x1, y1-10), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 0, 0), 2)
|
||
if frame_data['right_bbox'] is not None:
|
||
x1, y1, x2, y2 = frame_data['right_bbox'].astype(int)
|
||
cv2.rectangle(img_vis, (x1, y1), (x2, y2), (0, 255, 0), 2) # Green for right
|
||
cv2.putText(img_vis, f"R (conf={frame_data['right_conf']:.2f})",
|
||
(x1, y1-10), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)
|
||
|
||
img_fn = os.path.splitext(os.path.basename(img_path))[0]
|
||
cv2.imwrite(os.path.join(pass1_dir, f'{img_fn}_raw.jpg'), img_vis)
|
||
|
||
raw_data.append(frame_data)
|
||
|
||
# Print statistics with missing frame details
|
||
left_count = sum(1 for f in raw_data if f['left_bbox'] is not None)
|
||
right_count = sum(1 for f in raw_data if f['right_bbox'] is not None)
|
||
missing_left = [i for i, f in enumerate(raw_data) if f['left_bbox'] is None]
|
||
missing_right = [i for i, f in enumerate(raw_data) if f['right_bbox'] is None]
|
||
|
||
print(f"\nExtracted bboxes for {len(raw_data)} frames:")
|
||
print(f" Left hand detected in {left_count} frames ({100*left_count/len(raw_data):.1f}%)")
|
||
print(f" Right hand detected in {right_count} frames ({100*right_count/len(raw_data):.1f}%)")
|
||
|
||
if len(missing_left) > 0:
|
||
print(f" Missing left hand in {len(missing_left)} frames")
|
||
if len(missing_left) <= 20:
|
||
print(f" Frame indices: {missing_left}")
|
||
else:
|
||
print(f" First 10: {missing_left[:10]}")
|
||
print(f" Last 10: {missing_left[-10:]}")
|
||
|
||
if len(missing_right) > 0:
|
||
print(f" Missing right hand in {len(missing_right)} frames")
|
||
if len(missing_right) <= 20:
|
||
print(f" Frame indices: {missing_right}")
|
||
else:
|
||
print(f" First 10: {missing_right[:10]}")
|
||
print(f" Last 10: {missing_right[-10:]}")
|
||
|
||
# Create video from Pass 1 visualizations
|
||
if vis_dir is not None:
|
||
pass1_dir = os.path.join(vis_dir, 'pass1_raw_bboxes')
|
||
video_path = os.path.join(vis_dir, 'pass1_raw_bboxes.mp4')
|
||
create_video_from_images(pass1_dir, video_path, fps=30)
|
||
|
||
return raw_data
|
||
|
||
|
||
# ============================================================================
|
||
# PASS 2: Clean bbox sequences globally
|
||
# ============================================================================
|
||
|
||
def detect_overlapping_bboxes(raw_data, iou_threshold=0.7, containment_threshold=0.7):
|
||
"""
|
||
Detect and remove overlapping bboxes (hallucinations) - YOLO version without keypoints.
|
||
Uses simplified criteria:
|
||
1. High bbox IoU (> 0.7) - very overlapping bboxes
|
||
2. High containment (> 0.7) - one bbox is largely contained in another
|
||
|
||
Since YOLO doesn't provide keypoints, we only use bbox-based overlap detection.
|
||
This is actually more reliable since YOLO is specifically trained for hand detection.
|
||
|
||
Also tracks which hand was removed due to overlap and which hand "caused" the removal,
|
||
so we can potentially restore the removed hand if the winner is later invalidated.
|
||
"""
|
||
print("\n" + "-"*80)
|
||
print("Step 1: Detecting overlapping bboxes (hallucinations)")
|
||
print("-"*80)
|
||
|
||
overlap_count = 0
|
||
|
||
for frame_data in raw_data:
|
||
# Initialize removal tracking flags
|
||
frame_data['left_removed_due_to_overlap_with'] = None
|
||
frame_data['right_removed_due_to_overlap_with'] = None
|
||
if frame_data['left_bbox'] is not None and frame_data['right_bbox'] is not None:
|
||
# Criterion 1: Bbox IoU
|
||
iou = compute_iou(frame_data['left_bbox'], frame_data['right_bbox'])
|
||
|
||
# Criterion 2: Containment ratio (either bbox contains the other)
|
||
left_in_right = compute_containment_ratio(frame_data['left_bbox'], frame_data['right_bbox'])
|
||
right_in_left = compute_containment_ratio(frame_data['right_bbox'], frame_data['left_bbox'])
|
||
max_containment = max(left_in_right, right_in_left)
|
||
|
||
# Determine if this is an overlap/hallucination
|
||
is_overlap = False
|
||
reason = []
|
||
|
||
if iou > iou_threshold:
|
||
is_overlap = True
|
||
reason.append(f"IoU={iou:.3f}")
|
||
|
||
if max_containment > containment_threshold:
|
||
is_overlap = True
|
||
which_contained = "L_in_R" if left_in_right > right_in_left else "R_in_L"
|
||
reason.append(f"containment={max_containment:.2f} ({which_contained})")
|
||
|
||
if is_overlap:
|
||
overlap_count += 1
|
||
print(f" Frame {frame_data['frame_idx']:04d}: {', '.join(reason)} - OVERLAP DETECTED")
|
||
|
||
# Keep hand with higher confidence, track which hand was removed and why
|
||
if frame_data['left_conf'] > frame_data['right_conf']:
|
||
print(f" → Keeping LEFT (conf={frame_data['left_conf']:.3f}), removing RIGHT")
|
||
frame_data['right_bbox'] = None
|
||
frame_data['right_keypoints'] = None
|
||
frame_data['right_conf'] = 0.0
|
||
frame_data['right_removed_due_to_overlap_with'] = 'left'
|
||
else:
|
||
print(f" → Keeping RIGHT (conf={frame_data['right_conf']:.3f}), removing LEFT")
|
||
frame_data['left_bbox'] = None
|
||
frame_data['left_keypoints'] = None
|
||
frame_data['left_conf'] = 0.0
|
||
frame_data['left_removed_due_to_overlap_with'] = 'right'
|
||
|
||
print(f"Removed overlapping bboxes in {overlap_count} frames")
|
||
return raw_data
|
||
|
||
|
||
def fix_handedness_swaps_by_trajectory(raw_data, position_threshold=200):
|
||
"""
|
||
Detect handedness swaps by analyzing trajectory jumps.
|
||
If one hand suddenly jumps to where the other hand was, it's likely mislabeled.
|
||
|
||
Example: Green curve (right) drops to blue curve's position while blue disappears
|
||
→ The "right" detection is actually the left hand with wrong label
|
||
"""
|
||
print("\n" + "-"*80)
|
||
print(f"Step 1.4a: Detecting handedness swaps via trajectory analysis")
|
||
print("-"*80)
|
||
|
||
swap_count = 0
|
||
n = len(raw_data)
|
||
|
||
for i in range(1, n):
|
||
left_bbox = raw_data[i]['left_bbox']
|
||
right_bbox = raw_data[i]['right_bbox']
|
||
prev_left = raw_data[i-1]['left_bbox']
|
||
prev_right = raw_data[i-1]['right_bbox']
|
||
|
||
# Case 1: Right hand jumps to where left hand was (left disappears)
|
||
if right_bbox is not None and left_bbox is None and prev_left is not None and prev_right is not None:
|
||
# Calculate positions
|
||
prev_left_cx = (prev_left[0] + prev_left[2]) / 2
|
||
prev_right_cx = (prev_right[0] + prev_right[2]) / 2
|
||
curr_right_cx = (right_bbox[0] + right_bbox[2]) / 2
|
||
|
||
# Distance from current "right" to previous left/right
|
||
dist_to_prev_left = abs(curr_right_cx - prev_left_cx)
|
||
dist_to_prev_right = abs(curr_right_cx - prev_right_cx)
|
||
|
||
# If current "right" is much closer to previous left than previous right, it's swapped
|
||
if dist_to_prev_left < dist_to_prev_right and dist_to_prev_right > position_threshold:
|
||
print(f" Frame {i:04d}: RIGHT jumped to LEFT position (dist to prev_left={dist_to_prev_left:.0f}, dist to prev_right={dist_to_prev_right:.0f}) - Swapping RIGHT → LEFT")
|
||
raw_data[i]['left_bbox'] = right_bbox
|
||
raw_data[i]['left_keypoints'] = raw_data[i]['right_keypoints']
|
||
raw_data[i]['left_conf'] = raw_data[i]['right_conf']
|
||
raw_data[i]['right_bbox'] = None
|
||
raw_data[i]['right_keypoints'] = None
|
||
raw_data[i]['right_conf'] = 0.0
|
||
swap_count += 1
|
||
|
||
# Case 2: Left hand jumps to where right hand was (right disappears)
|
||
elif left_bbox is not None and right_bbox is None and prev_right is not None and prev_left is not None:
|
||
# Calculate positions
|
||
prev_left_cx = (prev_left[0] + prev_left[2]) / 2
|
||
prev_right_cx = (prev_right[0] + prev_right[2]) / 2
|
||
curr_left_cx = (left_bbox[0] + left_bbox[2]) / 2
|
||
|
||
# Distance from current "left" to previous left/right
|
||
dist_to_prev_left = abs(curr_left_cx - prev_left_cx)
|
||
dist_to_prev_right = abs(curr_left_cx - prev_right_cx)
|
||
|
||
# If current "left" is much closer to previous right than previous left, it's swapped
|
||
if dist_to_prev_right < dist_to_prev_left and dist_to_prev_left > position_threshold:
|
||
print(f" Frame {i:04d}: LEFT jumped to RIGHT position (dist to prev_right={dist_to_prev_right:.0f}, dist to prev_left={dist_to_prev_left:.0f}) - Swapping LEFT → RIGHT")
|
||
raw_data[i]['right_bbox'] = left_bbox
|
||
raw_data[i]['right_keypoints'] = raw_data[i]['left_keypoints']
|
||
raw_data[i]['right_conf'] = raw_data[i]['left_conf']
|
||
raw_data[i]['left_bbox'] = None
|
||
raw_data[i]['left_keypoints'] = None
|
||
raw_data[i]['left_conf'] = 0.0
|
||
swap_count += 1
|
||
|
||
if swap_count > 0:
|
||
print(f" Fixed {swap_count} trajectory-based handedness swaps")
|
||
else:
|
||
print(f" No trajectory-based swaps detected")
|
||
|
||
return raw_data
|
||
|
||
|
||
def fix_handedness_swaps(raw_data, context_window=10):
|
||
"""
|
||
Detect and fix handedness swaps where YOLO misclassifies hand labels.
|
||
For example, detecting left hand as right hand for a few frames.
|
||
|
||
Strategy: If a hand appears/disappears while the other hand simultaneously
|
||
disappears/appears in a short sequence, it's likely a handedness swap.
|
||
"""
|
||
print("\n" + "-"*80)
|
||
print(f"Step 1.4: Detecting handedness swaps (context window={context_window})")
|
||
print("-"*80)
|
||
|
||
swap_count = 0
|
||
n = len(raw_data)
|
||
|
||
for i in range(n):
|
||
left_bbox = raw_data[i]['left_bbox']
|
||
right_bbox = raw_data[i]['right_bbox']
|
||
|
||
# Case 1: Only right hand detected, but neighbors mostly have only left hand
|
||
if right_bbox is not None and left_bbox is None:
|
||
# Count neighbors with left/right
|
||
left_count = 0
|
||
right_count = 0
|
||
|
||
for j in range(max(0, i - context_window), min(n, i + context_window + 1)):
|
||
if j == i:
|
||
continue
|
||
if raw_data[j]['left_bbox'] is not None and raw_data[j]['right_bbox'] is None:
|
||
left_count += 1
|
||
elif raw_data[j]['right_bbox'] is not None and raw_data[j]['left_bbox'] is None:
|
||
right_count += 1
|
||
|
||
# If neighbors are mostly left-only, this right is probably mislabeled left
|
||
if left_count > 0 and left_count > right_count * 2:
|
||
print(f" Frame {i:04d}: Swapping RIGHT → LEFT (neighbors: {left_count} left-only, {right_count} right-only)")
|
||
raw_data[i]['left_bbox'] = right_bbox
|
||
raw_data[i]['left_keypoints'] = raw_data[i]['right_keypoints']
|
||
raw_data[i]['left_conf'] = raw_data[i]['right_conf']
|
||
raw_data[i]['right_bbox'] = None
|
||
raw_data[i]['right_keypoints'] = None
|
||
raw_data[i]['right_conf'] = 0.0
|
||
swap_count += 1
|
||
|
||
# Case 2: Only left hand detected, but neighbors mostly have only right hand
|
||
elif left_bbox is not None and right_bbox is None:
|
||
# Count neighbors with left/right
|
||
left_count = 0
|
||
right_count = 0
|
||
|
||
for j in range(max(0, i - context_window), min(n, i + context_window + 1)):
|
||
if j == i:
|
||
continue
|
||
if raw_data[j]['left_bbox'] is not None and raw_data[j]['right_bbox'] is None:
|
||
left_count += 1
|
||
elif raw_data[j]['right_bbox'] is not None and raw_data[j]['left_bbox'] is None:
|
||
right_count += 1
|
||
|
||
# If neighbors are mostly right-only, this left is probably mislabeled right
|
||
if right_count > 0 and right_count > left_count * 2:
|
||
print(f" Frame {i:04d}: Swapping LEFT → RIGHT (neighbors: {left_count} left-only, {right_count} right-only)")
|
||
raw_data[i]['right_bbox'] = left_bbox
|
||
raw_data[i]['right_keypoints'] = raw_data[i]['left_keypoints']
|
||
raw_data[i]['right_conf'] = raw_data[i]['left_conf']
|
||
raw_data[i]['left_bbox'] = None
|
||
raw_data[i]['left_keypoints'] = None
|
||
raw_data[i]['left_conf'] = 0.0
|
||
swap_count += 1
|
||
|
||
if swap_count > 0:
|
||
print(f" Fixed {swap_count} handedness swaps")
|
||
else:
|
||
print(f" No handedness swaps detected")
|
||
|
||
return raw_data
|
||
|
||
|
||
def fix_handedness_swaps_frame_to_frame(raw_data, iou_threshold=0.7, max_gap=10):
|
||
"""
|
||
Detect frame-to-frame handedness swaps: same bbox position but different label.
|
||
|
||
Example problem:
|
||
- Frame N: Right hand at position X (green box)
|
||
- Frame N+1: Left hand at position X (blue box) - SAME BBOX, WRONG LABEL!
|
||
|
||
This catches YOLO randomly mislabeling the same physical hand.
|
||
|
||
Strategy:
|
||
If a hand suddenly appears while the other hand recently disappeared (within max_gap frames),
|
||
AND the new hand's bbox has high IoU with the disappeared hand's LAST VALID bbox,
|
||
→ It's the same hand with wrong label! Swap it back!
|
||
|
||
Args:
|
||
max_gap: Maximum frames to look back for last valid other hand (default 10)
|
||
"""
|
||
print("\n" + "-"*80)
|
||
print(f"Step 1.4: Detecting frame-to-frame handedness swaps (IoU threshold={iou_threshold}, max_gap={max_gap})")
|
||
print("-"*80)
|
||
|
||
swap_count = 0
|
||
|
||
for i in range(1, len(raw_data)):
|
||
curr = raw_data[i]
|
||
|
||
for hand_name in ['left', 'right']:
|
||
bbox_key = f'{hand_name}_bbox'
|
||
keyp_key = f'{hand_name}_keypoints'
|
||
conf_key = f'{hand_name}_conf'
|
||
other_hand = 'right' if hand_name == 'left' else 'left'
|
||
other_bbox_key = f'{other_hand}_bbox'
|
||
other_keyp_key = f'{other_hand}_keypoints'
|
||
other_conf_key = f'{other_hand}_conf'
|
||
|
||
curr_bbox = curr[bbox_key]
|
||
curr_other_bbox = curr[other_bbox_key]
|
||
|
||
# Case: Current hand exists and other hand is missing
|
||
if curr_bbox is not None and curr_other_bbox is None:
|
||
# Look back to find LAST VALID bbox of OTHER hand (within max_gap frames)
|
||
last_valid_other_bbox = None
|
||
last_valid_other_idx = None
|
||
|
||
for j in range(i - 1, max(0, i - max_gap - 1), -1):
|
||
if raw_data[j][other_bbox_key] is not None:
|
||
last_valid_other_bbox = raw_data[j][other_bbox_key]
|
||
last_valid_other_idx = j
|
||
break
|
||
|
||
# If found a recent OTHER hand detection, check IoU
|
||
print(f" Frame {curr['frame_idx']:04d}: {last_valid_other_bbox is not None}")
|
||
if last_valid_other_bbox is not None:
|
||
gap = i - last_valid_other_idx
|
||
iou = compute_iou(curr_bbox, last_valid_other_bbox)
|
||
print(f" Frame {curr['frame_idx']:04d}: IoU = {iou:.3f}")
|
||
if iou > iou_threshold:
|
||
print(f" Frame {curr['frame_idx']:04d}: {hand_name} at same position as {other_hand} "
|
||
f"(last seen at frame {raw_data[last_valid_other_idx]['frame_idx']}, gap={gap}, IoU={iou:.3f}) "
|
||
f"- SWAPPING {hand_name} → {other_hand}")
|
||
|
||
# Swap: this is actually the OTHER hand
|
||
curr[other_bbox_key] = curr_bbox.copy()
|
||
curr[other_keyp_key] = curr[keyp_key].copy() if curr[keyp_key] is not None else None
|
||
curr[other_conf_key] = curr[conf_key]
|
||
|
||
# Remove from current hand
|
||
curr[bbox_key] = None
|
||
curr[keyp_key] = None
|
||
curr[conf_key] = 0.0
|
||
|
||
swap_count += 1
|
||
|
||
if swap_count > 0:
|
||
print(f" Fixed {swap_count} frame-to-frame handedness swaps")
|
||
else:
|
||
print(f" No frame-to-frame swaps detected")
|
||
|
||
return raw_data
|
||
|
||
|
||
def fix_handedness_inconsistencies(raw_data, original_data, context_window=5):
|
||
"""
|
||
Detect handedness swaps using SPATIAL-TEMPORAL consistency.
|
||
|
||
Example problem: [L, L, L, _, _, R, R, R, L, L]
|
||
- The 3 Rs are spatially in the LEFT hand's trajectory but mislabeled as RIGHT
|
||
|
||
Strategy:
|
||
1. For each detection, compute IoU with recent detections of SAME hand
|
||
2. Also compute IoU with recent detections of OTHER hand
|
||
3. If IoU with OTHER hand's trajectory is much higher → handedness is swapped!
|
||
"""
|
||
print("\n" + "-"*80)
|
||
print(f"Step 1.5: Fixing handedness inconsistencies via spatial-temporal consistency (window={context_window})")
|
||
print("-"*80)
|
||
|
||
swap_count = 0
|
||
|
||
for i, frame_data in enumerate(raw_data):
|
||
for hand_name in ['left', 'right']:
|
||
bbox_key = f'{hand_name}_bbox'
|
||
keyp_key = f'{hand_name}_keypoints'
|
||
conf_key = f'{hand_name}_conf'
|
||
other_hand = 'right' if hand_name == 'left' else 'left'
|
||
other_bbox_key = f'{other_hand}_bbox'
|
||
other_keyp_key = f'{other_hand}_keypoints'
|
||
other_conf_key = f'{other_hand}_conf'
|
||
|
||
current_bbox = frame_data[bbox_key]
|
||
if current_bbox is None:
|
||
continue
|
||
|
||
# Collect recent bboxes for SAME hand and OTHER hand
|
||
same_hand_bboxes = []
|
||
other_hand_bboxes = []
|
||
|
||
# Look at context window (both past and future)
|
||
for j in range(max(0, i - context_window), min(len(raw_data), i + context_window + 1)):
|
||
if j == i:
|
||
continue # Skip current frame
|
||
|
||
# Collect SAME hand bboxes
|
||
if raw_data[j][bbox_key] is not None:
|
||
same_hand_bboxes.append(raw_data[j][bbox_key])
|
||
|
||
# Collect OTHER hand bboxes
|
||
if raw_data[j][other_bbox_key] is not None:
|
||
other_hand_bboxes.append(raw_data[j][other_bbox_key])
|
||
|
||
# Need enough context to make a decision
|
||
if len(same_hand_bboxes) < 2 and len(other_hand_bboxes) < 2:
|
||
continue # Not enough spatial context
|
||
|
||
# Compute average IoU with SAME hand trajectory
|
||
same_hand_ious = []
|
||
for bbox in same_hand_bboxes:
|
||
iou = compute_iou(current_bbox, bbox)
|
||
same_hand_ious.append(iou)
|
||
avg_iou_same = np.mean(same_hand_ious) if len(same_hand_ious) > 0 else 0.0
|
||
|
||
# Compute average IoU with OTHER hand trajectory
|
||
other_hand_ious = []
|
||
for bbox in other_hand_bboxes:
|
||
iou = compute_iou(current_bbox, bbox)
|
||
other_hand_ious.append(iou)
|
||
avg_iou_other = np.mean(other_hand_ious) if len(other_hand_ious) > 0 else 0.0
|
||
|
||
# DETECTION LOGIC: If current bbox fits OTHER hand's trajectory much better, it's swapped!
|
||
# Criteria:
|
||
# 1. OTHER hand trajectory has high overlap (>0.5) with current bbox
|
||
# 2. OTHER hand trajectory overlap is significantly higher than SAME hand (1.5x)
|
||
# 3. We have enough OTHER hand detections nearby (at least 3)
|
||
if (len(other_hand_bboxes) >= 3 and
|
||
avg_iou_other > 0.5 and
|
||
avg_iou_other > avg_iou_same * 1.5):
|
||
|
||
print(f" Frame {frame_data['frame_idx']:04d} ({hand_name}): "
|
||
f"Bbox fits {other_hand} trajectory better! "
|
||
f"(IoU with {hand_name}={avg_iou_same:.3f}, IoU with {other_hand}={avg_iou_other:.3f}, "
|
||
f"neighbors: {len(same_hand_bboxes)} {hand_name}, {len(other_hand_bboxes)} {other_hand}) "
|
||
f"- SWAPPING {hand_name} → {other_hand}")
|
||
|
||
# Swap: move current bbox to OTHER hand
|
||
frame_data[other_bbox_key] = current_bbox.copy()
|
||
frame_data[other_keyp_key] = frame_data[keyp_key].copy() if frame_data[keyp_key] is not None else None
|
||
frame_data[other_conf_key] = frame_data[conf_key]
|
||
|
||
# Remove from current hand
|
||
frame_data[bbox_key] = None
|
||
frame_data[keyp_key] = None
|
||
frame_data[conf_key] = 0.0
|
||
|
||
swap_count += 1
|
||
|
||
if swap_count > 0:
|
||
print(f" Fixed {swap_count} handedness swaps via spatial-temporal consistency")
|
||
else:
|
||
print(f" No spatial-temporal inconsistencies found")
|
||
|
||
return raw_data
|
||
|
||
|
||
def detect_bbox_jumps(raw_data, base_iou_threshold=0.5):
|
||
"""
|
||
Detect and fix bbox jumps with ADAPTIVE threshold based on detection gaps.
|
||
|
||
Adaptive threshold formula: adjusted_threshold = max(0.0, base_threshold - gap_frames * 0.15)
|
||
- Gap 0 (consecutive): threshold = 0.5 (strict, catches hallucinations)
|
||
- Gap 1: threshold = 0.35 (slightly lenient)
|
||
- Gap 2: threshold = 0.2 (more lenient)
|
||
- Gap 3+: threshold = 0.05/0.0 (very lenient, accepts genuine movement)
|
||
|
||
This allows genuine hand movement after detection gaps while still protecting
|
||
against spurious jumps in consecutive frames.
|
||
"""
|
||
print("\n" + "-"*80)
|
||
print(f"Step 2: Detecting bbox jumps with adaptive threshold (base={base_iou_threshold})")
|
||
print("-"*80)
|
||
|
||
for hand_name in ['left', 'right']:
|
||
bbox_key = f'{hand_name}_bbox'
|
||
hand_id = 0 if hand_name == 'left' else 1
|
||
|
||
last_valid_bbox = None
|
||
last_valid_idx = None
|
||
frames_since_last_detection = 0
|
||
jump_count = 0
|
||
|
||
for i, frame_data in enumerate(raw_data):
|
||
current_bbox = frame_data[bbox_key]
|
||
|
||
if current_bbox is not None:
|
||
if last_valid_bbox is not None:
|
||
# Calculate IoU with last valid bbox
|
||
iou = compute_iou(current_bbox, last_valid_bbox)
|
||
|
||
# Adaptive threshold based on gap
|
||
gap_frames = frames_since_last_detection
|
||
adjusted_threshold = max(0.0, base_iou_threshold - (gap_frames * 0.01))
|
||
|
||
if iou < adjusted_threshold:
|
||
# Jump detected!
|
||
jump_count += 1
|
||
print(f" Frame {frame_data['frame_idx']:04d} ({hand_name}): "
|
||
f"IoU={iou:.3f} < threshold={adjusted_threshold:.3f} (gap={gap_frames}) - JUMP DETECTED")
|
||
print(f" → Current bbox: [{current_bbox[0]:.1f}, {current_bbox[1]:.1f}, {current_bbox[2]:.1f}, {current_bbox[3]:.1f}]")
|
||
print(f" → Last valid bbox (frame {last_valid_idx}): [{last_valid_bbox[0]:.1f}, {last_valid_bbox[1]:.1f}, {last_valid_bbox[2]:.1f}, {last_valid_bbox[3]:.1f}]")
|
||
print(f" → Replacing with last valid bbox from frame {last_valid_idx}")
|
||
frame_data[bbox_key] = last_valid_bbox.copy()
|
||
else:
|
||
# Valid detection
|
||
if gap_frames > 0:
|
||
print(f" Frame {frame_data['frame_idx']:04d} ({hand_name}): "
|
||
f"IoU={iou:.3f} > threshold={adjusted_threshold:.3f} (gap={gap_frames}) - ACCEPTING movement")
|
||
last_valid_bbox = current_bbox.copy()
|
||
last_valid_idx = i
|
||
frames_since_last_detection = 0
|
||
else:
|
||
# First valid bbox
|
||
last_valid_bbox = current_bbox.copy()
|
||
last_valid_idx = i
|
||
frames_since_last_detection = 0
|
||
else:
|
||
# Hand is missing
|
||
frames_since_last_detection += 1
|
||
|
||
if jump_count > 0:
|
||
print(f" {hand_name.capitalize()}: Fixed {jump_count} jumps")
|
||
|
||
return raw_data
|
||
|
||
|
||
def interpolate_missing_bboxes(raw_data, max_gap=5):
|
||
"""Interpolate missing bboxes for short gaps and create dummy keypoints."""
|
||
print("\n" + "-"*80)
|
||
print(f"Step 3: Interpolating missing bboxes (max gap={max_gap} frames)")
|
||
print("-"*80)
|
||
|
||
for hand_name in ['left', 'right']:
|
||
bbox_key = f'{hand_name}_bbox'
|
||
keyp_key = f'{hand_name}_keypoints'
|
||
conf_key = f'{hand_name}_conf'
|
||
|
||
# Find valid bbox indices
|
||
valid_indices = [i for i, f in enumerate(raw_data) if f[bbox_key] is not None]
|
||
|
||
if len(valid_indices) < 2:
|
||
continue
|
||
|
||
interpolated_count = 0
|
||
|
||
for i in range(len(valid_indices) - 1):
|
||
start_idx = valid_indices[i]
|
||
end_idx = valid_indices[i + 1]
|
||
gap = end_idx - start_idx - 1
|
||
|
||
if 0 < gap <= max_gap:
|
||
start_bbox = raw_data[start_idx][bbox_key]
|
||
end_bbox = raw_data[end_idx][bbox_key]
|
||
|
||
for j in range(1, gap + 1):
|
||
alpha = j / (gap + 1)
|
||
interp_bbox = (1 - alpha) * start_bbox + alpha * end_bbox
|
||
raw_data[start_idx + j][bbox_key] = interp_bbox
|
||
|
||
# Create dummy keypoints for interpolated frames
|
||
dummy_keypoints = np.zeros((21, 3))
|
||
dummy_keypoints[:, 2] = 0.5
|
||
raw_data[start_idx + j][keyp_key] = dummy_keypoints
|
||
raw_data[start_idx + j][conf_key] = 0.5
|
||
|
||
interpolated_count += 1
|
||
|
||
print(f" {hand_name.capitalize()}: Interpolated frames {start_idx+1}-{end_idx-1} (gap={gap})")
|
||
|
||
if interpolated_count > 0:
|
||
print(f" {hand_name.capitalize()}: Interpolated {interpolated_count} frames total")
|
||
|
||
return raw_data
|
||
|
||
|
||
def smooth_bbox_sequence(raw_data, window_size=5):
|
||
"""Apply temporal smoothing using moving average."""
|
||
print("\n" + "-"*80)
|
||
print(f"Step 4: Smoothing bbox sequences (window={window_size})")
|
||
print("-"*80)
|
||
|
||
for hand_name in ['left', 'right']:
|
||
bbox_key = f'{hand_name}_bbox'
|
||
|
||
bboxes = [f[bbox_key] for f in raw_data]
|
||
smoothed_bboxes = []
|
||
|
||
for i in range(len(bboxes)):
|
||
if bboxes[i] is None:
|
||
smoothed_bboxes.append(None)
|
||
continue
|
||
|
||
window_start = max(0, i - window_size // 2)
|
||
window_end = min(len(bboxes), i + window_size // 2 + 1)
|
||
window_bboxes = [bboxes[j] for j in range(window_start, window_end) if bboxes[j] is not None]
|
||
|
||
if len(window_bboxes) > 0:
|
||
smoothed_bbox = np.mean(window_bboxes, axis=0)
|
||
smoothed_bboxes.append(smoothed_bbox)
|
||
else:
|
||
smoothed_bboxes.append(bboxes[i])
|
||
|
||
for i, frame_data in enumerate(raw_data):
|
||
if smoothed_bboxes[i] is not None:
|
||
frame_data[bbox_key] = smoothed_bboxes[i]
|
||
|
||
print(" Smoothing complete")
|
||
return raw_data
|
||
|
||
|
||
def check_if_bboxes_overlap(bbox1, bbox2, keypoints1=None, keypoints2=None,
|
||
iou_threshold=0.7, containment_threshold=0.7,
|
||
wrist_dist_ratio=0.15, keypoint_sim_ratio=0.10):
|
||
"""
|
||
Check if two bboxes overlap using the same comprehensive criteria as detect_overlapping_bboxes.
|
||
Returns True if the bboxes are considered overlapping (hallucination).
|
||
"""
|
||
# Compute average bbox size for normalization
|
||
size1 = max(bbox1[2] - bbox1[0], bbox1[3] - bbox1[1])
|
||
size2 = max(bbox2[2] - bbox2[0], bbox2[3] - bbox2[1])
|
||
avg_bbox_size = (size1 + size2) / 2
|
||
|
||
# Criterion 1: Bbox IoU
|
||
iou = compute_iou(bbox1, bbox2)
|
||
if iou > iou_threshold:
|
||
return True
|
||
|
||
# Criterion 2: Containment ratio
|
||
containment_1_in_2 = compute_containment_ratio(bbox1, bbox2)
|
||
containment_2_in_1 = compute_containment_ratio(bbox2, bbox1)
|
||
max_containment = max(containment_1_in_2, containment_2_in_1)
|
||
if max_containment > containment_threshold:
|
||
return True
|
||
|
||
# Criterion 3 & 4: Keypoint-based checks (if keypoints available)
|
||
if keypoints1 is not None and keypoints2 is not None:
|
||
# Wrist distance
|
||
wrist1 = keypoints1[0, :2]
|
||
wrist2 = keypoints2[0, :2]
|
||
wrist_dist = np.linalg.norm(wrist1 - wrist2)
|
||
wrist_dist_normalized = wrist_dist / avg_bbox_size
|
||
if wrist_dist_normalized < wrist_dist_ratio:
|
||
return True
|
||
|
||
# Average keypoint distance
|
||
keypoint_dists = []
|
||
for i in range(min(len(keypoints1), len(keypoints2))):
|
||
kp1 = keypoints1[i, :2]
|
||
kp2 = keypoints2[i, :2]
|
||
dist = np.linalg.norm(kp1 - kp2)
|
||
keypoint_dists.append(dist)
|
||
avg_keypoint_dist = np.mean(keypoint_dists)
|
||
avg_keypoint_dist_normalized = avg_keypoint_dist / avg_bbox_size
|
||
if avg_keypoint_dist_normalized < keypoint_sim_ratio:
|
||
return True
|
||
|
||
return False
|
||
|
||
|
||
def plot_bbox_trajectories(raw_data, vis_dir, filename='bbox_trajectories.png'):
|
||
"""
|
||
Plot left and right hand bbox center trajectories over time.
|
||
This helps visualize handedness consistency and detect label swaps.
|
||
"""
|
||
import matplotlib
|
||
matplotlib.use('Agg') # Use non-interactive backend
|
||
import matplotlib.pyplot as plt
|
||
|
||
frames = []
|
||
left_centers_x = []
|
||
left_centers_y = []
|
||
right_centers_x = []
|
||
right_centers_y = []
|
||
|
||
for i, frame_data in enumerate(raw_data):
|
||
frames.append(i)
|
||
|
||
# Left hand
|
||
if frame_data['left_bbox'] is not None:
|
||
left_cx = (frame_data['left_bbox'][0] + frame_data['left_bbox'][2]) / 2
|
||
left_cy = (frame_data['left_bbox'][1] + frame_data['left_bbox'][3]) / 2
|
||
left_centers_x.append(left_cx)
|
||
left_centers_y.append(left_cy)
|
||
else:
|
||
left_centers_x.append(None)
|
||
left_centers_y.append(None)
|
||
|
||
# Right hand
|
||
if frame_data['right_bbox'] is not None:
|
||
right_cx = (frame_data['right_bbox'][0] + frame_data['right_bbox'][2]) / 2
|
||
right_cy = (frame_data['right_bbox'][1] + frame_data['right_bbox'][3]) / 2
|
||
right_centers_x.append(right_cx)
|
||
right_centers_y.append(right_cy)
|
||
else:
|
||
right_centers_x.append(None)
|
||
right_centers_y.append(None)
|
||
|
||
# Create figure with 2 subplots
|
||
fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(20, 10))
|
||
|
||
# Plot X coordinates (horizontal position)
|
||
ax1.plot(frames, left_centers_x, 'b-', label='Left hand', linewidth=1.5, alpha=0.8)
|
||
ax1.plot(frames, right_centers_x, 'g-', label='Right hand', linewidth=1.5, alpha=0.8)
|
||
ax1.set_xlabel('Frame', fontsize=12)
|
||
ax1.set_ylabel('X coordinate (pixels)', fontsize=12)
|
||
ax1.set_title('Hand Bbox Center X-Position Over Time\n(Left hand should be consistently < Right hand)', fontsize=14, fontweight='bold')
|
||
ax1.legend(fontsize=11)
|
||
ax1.grid(True, alpha=0.3)
|
||
|
||
# Highlight regions where left > right (potential swaps)
|
||
swap_count = 0
|
||
for i in range(len(frames)):
|
||
if left_centers_x[i] is not None and right_centers_x[i] is not None:
|
||
if left_centers_x[i] > right_centers_x[i]:
|
||
ax1.axvspan(i-0.5, i+0.5, alpha=0.3, color='red')
|
||
swap_count += 1
|
||
|
||
# Plot Y coordinates (vertical position)
|
||
ax2.plot(frames, left_centers_y, 'b-', label='Left hand', linewidth=1.5, alpha=0.8)
|
||
ax2.plot(frames, right_centers_y, 'g-', label='Right hand', linewidth=1.5, alpha=0.8)
|
||
ax2.set_xlabel('Frame', fontsize=12)
|
||
ax2.set_ylabel('Y coordinate (pixels)', fontsize=12)
|
||
ax2.set_title('Hand Bbox Center Y-Position Over Time', fontsize=14, fontweight='bold')
|
||
ax2.legend(fontsize=11)
|
||
ax2.grid(True, alpha=0.3)
|
||
ax2.invert_yaxis() # Invert Y axis since image coordinates go top-to-bottom
|
||
|
||
plt.tight_layout()
|
||
save_path = os.path.join(vis_dir, filename)
|
||
plt.savefig(save_path, dpi=150, bbox_inches='tight')
|
||
plt.close()
|
||
print(f"\n{'='*80}")
|
||
print(f"Bbox trajectory plot saved to: {save_path}")
|
||
print(f" Red regions indicate frames where left_x > right_x (potential label swaps)")
|
||
print(f" Total frames with left_x > right_x: {swap_count}")
|
||
print(f"{'='*80}\n")
|
||
|
||
|
||
def clean_bbox_sequences(raw_data, vis_dir=None):
|
||
"""Pass 2: Clean bbox sequences using global temporal information."""
|
||
print("\n" + "="*80)
|
||
print("PASS 2: Cleaning bbox sequences")
|
||
print("="*80)
|
||
|
||
if vis_dir is not None:
|
||
pass2_dir = os.path.join(vis_dir, 'pass2_cleaned_bboxes')
|
||
os.makedirs(pass2_dir, exist_ok=True)
|
||
|
||
# Store original bboxes and keypoints for comparison and restoration
|
||
original_data = [{
|
||
'left_bbox': f['left_bbox'].copy() if f['left_bbox'] is not None else None,
|
||
'right_bbox': f['right_bbox'].copy() if f['right_bbox'] is not None else None,
|
||
'left_keypoints': f['left_keypoints'].copy() if f['left_keypoints'] is not None else None,
|
||
'right_keypoints': f['right_keypoints'].copy() if f['right_keypoints'] is not None else None,
|
||
'left_conf': f['left_conf'] if f['left_conf'] is not None else None,
|
||
'right_conf': f['right_conf'] if f['right_conf'] is not None else None,
|
||
} for f in raw_data]
|
||
|
||
# Plot trajectories BEFORE cleaning
|
||
if vis_dir is not None:
|
||
plot_bbox_trajectories(raw_data, vis_dir, filename='bbox_trajectories_before_cleaning.png')
|
||
|
||
# Remove oversized bboxes (> 50% of image area)
|
||
print("\n" + "-"*80)
|
||
print("Step 1.5: Removing oversized bboxes (> 50% of image area)")
|
||
print("-"*80)
|
||
MAX_BBOX_AREA_RATIO = 0.5
|
||
removed_oversized = 0
|
||
|
||
# Load first image to get dimensions (all images should have same size)
|
||
first_img = cv2.imread(raw_data[0]['img_path'])
|
||
img_area = first_img.shape[0] * first_img.shape[1]
|
||
|
||
for frame_data in raw_data:
|
||
for hand_name in ['left', 'right']:
|
||
bbox_key = f'{hand_name}_bbox'
|
||
keyp_key = f'{hand_name}_keypoints'
|
||
conf_key = f'{hand_name}_conf'
|
||
|
||
bbox = frame_data[bbox_key]
|
||
if bbox is not None:
|
||
bbox_width = bbox[2] - bbox[0]
|
||
bbox_height = bbox[3] - bbox[1]
|
||
bbox_area = bbox_width * bbox_height
|
||
area_ratio = bbox_area / img_area
|
||
|
||
if area_ratio > MAX_BBOX_AREA_RATIO:
|
||
print(f" Frame {frame_data['frame_idx']:04d} ({hand_name}): "
|
||
f"Bbox too large ({area_ratio:.1%} of image) - REMOVING")
|
||
frame_data[bbox_key] = None
|
||
frame_data[keyp_key] = None
|
||
frame_data[conf_key] = 0.0
|
||
removed_oversized += 1
|
||
|
||
if removed_oversized > 0:
|
||
print(f" Removed {removed_oversized} oversized bboxes")
|
||
else:
|
||
print(f" No oversized bboxes found")
|
||
|
||
raw_data = detect_overlapping_bboxes(raw_data, iou_threshold=0.7)
|
||
raw_data = fix_handedness_swaps_by_trajectory(raw_data, position_threshold=200)
|
||
raw_data = fix_handedness_inconsistencies(raw_data, original_data, context_window=5)
|
||
# raw_data = detect_bbox_jumps(raw_data, base_iou_threshold=0.4)
|
||
raw_data = fix_handedness_swaps_frame_to_frame(raw_data, iou_threshold=0.6)
|
||
|
||
# Apply patience mechanism with interpolation when hand reappears
|
||
print("\n" + "-"*80)
|
||
print("Step 3: Applying patience mechanism with interpolation")
|
||
print("-"*80)
|
||
PATIENCE_FRAMES = 25
|
||
MAX_PATIENCE_WITHOUT_RETURN = 0 # No padding if hand never reappears
|
||
patience_applied_count = 0
|
||
interpolated_count = 0
|
||
trimmed_count = 0
|
||
|
||
for hand_id in [0, 1]:
|
||
hand_name = 'left' if hand_id == 0 else 'right'
|
||
bbox_key = f'{hand_name}_bbox'
|
||
keyp_key = f'{hand_name}_keypoints'
|
||
conf_key = f'{hand_name}_conf'
|
||
|
||
i = 0
|
||
while i < len(raw_data):
|
||
# Find start of a gap (hand missing)
|
||
if raw_data[i][bbox_key] is None:
|
||
# Look back to find last valid bbox
|
||
last_valid_idx = None
|
||
for j in range(i - 1, -1, -1):
|
||
if raw_data[j][bbox_key] is not None:
|
||
last_valid_idx = j
|
||
break
|
||
|
||
if last_valid_idx is None:
|
||
i += 1
|
||
continue
|
||
|
||
# Find end of gap (hand reappears or sequence ends)
|
||
gap_start = i
|
||
gap_end = None
|
||
for j in range(i, min(i + PATIENCE_FRAMES, len(raw_data))):
|
||
if raw_data[j][bbox_key] is not None:
|
||
gap_end = j
|
||
break
|
||
|
||
if gap_end is not None:
|
||
# Hand reappeared within patience threshold - CHECK DISTANCE FIRST
|
||
gap_length = gap_end - gap_start
|
||
start_bbox = raw_data[last_valid_idx][bbox_key]
|
||
end_bbox = raw_data[gap_end][bbox_key]
|
||
|
||
# Compute center distance between start and end bbox
|
||
start_cx = (start_bbox[0] + start_bbox[2]) / 2
|
||
start_cy = (start_bbox[1] + start_bbox[3]) / 2
|
||
end_cx = (end_bbox[0] + end_bbox[2]) / 2
|
||
end_cy = (end_bbox[1] + end_bbox[3]) / 2
|
||
center_distance = np.sqrt((end_cx - start_cx)**2 + (end_cy - start_cy)**2)
|
||
|
||
# Compute bbox width (use average of start and end)
|
||
start_width = start_bbox[2] - start_bbox[0]
|
||
end_width = end_bbox[2] - end_bbox[0]
|
||
avg_width = (start_width + end_width) / 2
|
||
|
||
# Only interpolate if bboxes are close enough (within 2× bbox width)
|
||
max_distance = avg_width
|
||
|
||
if center_distance <= max_distance:
|
||
# Bboxes are close - INTERPOLATE
|
||
print(f" {hand_name.capitalize()}: Interpolating frames {gap_start}-{gap_end-1} "
|
||
f"(gap={gap_length}, from frame {last_valid_idx} to {gap_end}, "
|
||
f"distance={center_distance:.1f}px < {max_distance:.1f}px)")
|
||
|
||
# Total distance includes the step from last_valid to gap_start
|
||
# and from gap_end-1 to gap_end
|
||
total_steps = gap_end - last_valid_idx
|
||
|
||
for j in range(gap_start, gap_end):
|
||
# Linear interpolation: evenly distribute motion across all steps
|
||
alpha = (j - last_valid_idx) / total_steps
|
||
interp_bbox = (1 - alpha) * start_bbox + alpha * end_bbox
|
||
raw_data[j][bbox_key] = interp_bbox
|
||
|
||
# Create dummy keypoints
|
||
dummy_keypoints = np.zeros((21, 3))
|
||
dummy_keypoints[:, 2] = 0.5
|
||
raw_data[j][keyp_key] = dummy_keypoints
|
||
raw_data[j][conf_key] = 0.5
|
||
interpolated_count += 1
|
||
|
||
i = gap_end
|
||
else:
|
||
# Bboxes are too far apart - DON'T interpolate, treat as separate instances
|
||
print(f" {hand_name.capitalize()}: Skipping interpolation for frames {gap_start}-{gap_end-1} "
|
||
f"(distance={center_distance:.1f}px > {max_distance:.1f}px - likely different hand instances)")
|
||
|
||
# Move to the next detection without filling the gap
|
||
i = gap_end
|
||
else:
|
||
# Hand never reappeared within patience threshold
|
||
# Apply limited patience (only MAX_PATIENCE_WITHOUT_RETURN frames), then skip the rest
|
||
last_bbox = raw_data[last_valid_idx][bbox_key]
|
||
frames_to_fill = min(MAX_PATIENCE_WITHOUT_RETURN, len(raw_data) - gap_start)
|
||
|
||
if frames_to_fill > 0:
|
||
print(f" {hand_name.capitalize()}: Applying static patience for frames {gap_start}-{gap_start + frames_to_fill - 1} "
|
||
f"(hand never reappeared, using last bbox from frame {last_valid_idx})")
|
||
|
||
for j in range(gap_start, gap_start + frames_to_fill):
|
||
raw_data[j][bbox_key] = last_bbox.copy()
|
||
|
||
# Create dummy keypoints
|
||
dummy_keypoints = np.zeros((21, 3))
|
||
dummy_keypoints[:, 2] = 0.5
|
||
raw_data[j][keyp_key] = dummy_keypoints
|
||
raw_data[j][conf_key] = 0.5
|
||
patience_applied_count += 1
|
||
|
||
# Skip ahead to find where hand reappears (or end of sequence)
|
||
next_valid_idx = None
|
||
for j in range(gap_start + frames_to_fill, len(raw_data)):
|
||
if raw_data[j][bbox_key] is not None:
|
||
next_valid_idx = j
|
||
break
|
||
|
||
# Jump to next valid detection or end of sequence
|
||
i = next_valid_idx if next_valid_idx is not None else len(raw_data)
|
||
else:
|
||
i += 1
|
||
|
||
if interpolated_count > 0:
|
||
print(f" Interpolated {interpolated_count} frames where hand reappeared within patience threshold")
|
||
if patience_applied_count > 0:
|
||
print(f" Applied static patience to {patience_applied_count} frames where hand never reappeared")
|
||
if interpolated_count == 0 and patience_applied_count == 0:
|
||
print(f" No patience needed")
|
||
|
||
# Remove spurious short motions (brief appearances after long absence)
|
||
# This runs AFTER patience to filter out hallucinations that patience might have extended
|
||
print("\n" + "-"*80)
|
||
print("Step 4: Removing spurious short motions")
|
||
print("-"*80)
|
||
MIN_MOTION_DURATION = 30 # Motion must last at least this many frames
|
||
MIN_ABSENCE_BEFORE = 30 # Must be absent for this long before
|
||
MIN_ABSENCE_AFTER = 30 # Must be absent for this long after
|
||
|
||
for hand_name in ['left', 'right']:
|
||
bbox_key = f'{hand_name}_bbox'
|
||
keyp_key = f'{hand_name}_keypoints'
|
||
conf_key = f'{hand_name}_conf'
|
||
|
||
# Find all continuous segments where hand is present
|
||
segments = []
|
||
start_idx = None
|
||
|
||
for i, frame_data in enumerate(raw_data):
|
||
if frame_data[bbox_key] is not None:
|
||
if start_idx is None:
|
||
start_idx = i # Start of a new segment
|
||
else:
|
||
if start_idx is not None:
|
||
# End of a segment
|
||
segments.append((start_idx, i - 1))
|
||
start_idx = None
|
||
|
||
# Don't forget the last segment if it extends to the end
|
||
if start_idx is not None:
|
||
segments.append((start_idx, len(raw_data) - 1))
|
||
|
||
# Check each segment
|
||
removed_segments = []
|
||
for seg_start, seg_end in segments:
|
||
seg_duration = seg_end - seg_start + 1
|
||
|
||
# Check absence before
|
||
absence_before = seg_start # frames before this segment
|
||
|
||
# Check absence after
|
||
absence_after = len(raw_data) - 1 - seg_end # frames after this segment
|
||
|
||
# If this is a short motion after long absence and before long absence, remove it
|
||
if (seg_duration < MIN_MOTION_DURATION and
|
||
absence_before >= MIN_ABSENCE_BEFORE and
|
||
absence_after >= MIN_ABSENCE_AFTER):
|
||
|
||
print(f" {hand_name.capitalize()}: Removing spurious short motion at frames {seg_start}-{seg_end} "
|
||
f"(duration={seg_duration}, absence_before={absence_before}, absence_after={absence_after})")
|
||
|
||
# Remove this segment
|
||
for i in range(seg_start, seg_end + 1):
|
||
raw_data[i][bbox_key] = None
|
||
raw_data[i][keyp_key] = None
|
||
raw_data[i][conf_key] = 0.0
|
||
|
||
removed_segments.append((seg_start, seg_end))
|
||
|
||
if len(removed_segments) > 0:
|
||
print(f" {hand_name.capitalize()}: Removed {len(removed_segments)} spurious short motion segments")
|
||
else:
|
||
print(f" {hand_name.capitalize()}: No spurious short motions found")
|
||
|
||
# raw_data = interpolate_missing_bboxes(raw_data, max_gap=10)
|
||
# # Bbox smoothing disabled - it was causing drift (e.g., frames 326->330)
|
||
# raw_data = smooth_bbox_sequence(raw_data, window_size=5)
|
||
|
||
# Re-check for overlaps after interpolation (interpolation can create new overlaps)
|
||
print("\n" + "-"*80)
|
||
print("Step 4: Re-checking overlaps after interpolation")
|
||
print("-"*80)
|
||
raw_data = detect_overlapping_bboxes(raw_data, iou_threshold=0.7)
|
||
|
||
# Restore hands that were removed due to overlap, but ONLY if the "winner" hand was
|
||
# subsequently invalidated by other checks (handedness consistency).
|
||
# Logic: If LEFT was removed because RIGHT had higher confidence, but then RIGHT itself
|
||
# was removed by handedness check, we should restore LEFT since it was the original detection.
|
||
print("\n" + "-"*80)
|
||
print("Step 5: Restoring hands removed by overlap if winner was invalidated")
|
||
print("-"*80)
|
||
restored_count = 0
|
||
for i, (frame_data, orig) in enumerate(zip(raw_data, original_data)):
|
||
for hand_name in ['left', 'right']:
|
||
bbox_key = f'{hand_name}_bbox'
|
||
keyp_key = f'{hand_name}_keypoints'
|
||
conf_key = f'{hand_name}_conf'
|
||
removal_flag = f'{hand_name}_removed_due_to_overlap_with'
|
||
|
||
# Check if this hand was removed due to overlap
|
||
if frame_data.get(removal_flag) is not None:
|
||
other_hand = frame_data[removal_flag] # The hand that "won" the overlap check
|
||
other_bbox_key = f'{other_hand}_bbox'
|
||
other_handedness_flag = f'{other_hand}_removed_by_handedness_check'
|
||
|
||
# Check if the "winner" hand was subsequently removed by handedness check
|
||
if frame_data.get(other_handedness_flag, False):
|
||
print(f" Frame {i}: Restoring {hand_name} hand (winner '{other_hand}' was invalidated)")
|
||
# Restore original keypoints and confidence
|
||
if orig[keyp_key] is not None:
|
||
frame_data[bbox_key] = orig[bbox_key].copy()
|
||
frame_data[keyp_key] = orig[keyp_key].copy()
|
||
frame_data[conf_key] = orig[conf_key]
|
||
else:
|
||
# Create dummy keypoints if original didn't have them
|
||
dummy_keypoints = np.zeros((21, 3))
|
||
dummy_keypoints[:, 2] = 0.5
|
||
frame_data[keyp_key] = dummy_keypoints
|
||
frame_data[conf_key] = 0.5
|
||
# Clear the removal flag since we've restored it
|
||
frame_data[removal_flag] = None
|
||
restored_count += 1
|
||
|
||
raw_data = fix_handedness_swaps_frame_to_frame(raw_data, iou_threshold=0.6)
|
||
|
||
if restored_count > 0:
|
||
print(f" Restored {restored_count} hand instances")
|
||
else:
|
||
print(f" No hands needed restoration")
|
||
|
||
# Visualize before/after comparison
|
||
dash_length = 10
|
||
if vis_dir is not None:
|
||
for i, frame_data in enumerate(raw_data):
|
||
img_path = frame_data['img_path']
|
||
img_cv2 = cv2.imread(img_path)
|
||
img_vis = img_cv2.copy()
|
||
|
||
# Draw original bboxes (dashed)
|
||
orig = original_data[i]
|
||
if orig['left_bbox'] is not None:
|
||
x1, y1, x2, y2 = orig['left_bbox'].astype(int)
|
||
# Draw dashed rectangle
|
||
|
||
for j in range(x1, x2, dash_length * 2):
|
||
cv2.line(img_vis, (j, y1), (min(j + dash_length, x2), y1), (173, 216, 230), 2)
|
||
cv2.line(img_vis, (j, y2), (min(j + dash_length, x2), y2), (173, 216, 230), 2)
|
||
for j in range(y1, y2, dash_length * 2):
|
||
cv2.line(img_vis, (x1, j), (x1, min(j + dash_length, y2)), (173, 216, 230), 2)
|
||
cv2.line(img_vis, (x2, j), (x2, min(j + dash_length, y2)), (173, 216, 230), 2)
|
||
cv2.putText(img_vis, "L_before", (x1, y1-30), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (173, 216, 230), 2)
|
||
|
||
if orig['right_bbox'] is not None:
|
||
x1, y1, x2, y2 = orig['right_bbox'].astype(int)
|
||
for j in range(x1, x2, dash_length * 2):
|
||
cv2.line(img_vis, (j, y1), (min(j + dash_length, x2), y1), (144, 238, 144), 2)
|
||
cv2.line(img_vis, (j, y2), (min(j + dash_length, x2), y2), (144, 238, 144), 2)
|
||
for j in range(y1, y2, dash_length * 2):
|
||
cv2.line(img_vis, (x1, j), (x1, min(j + dash_length, y2)), (144, 238, 144), 2)
|
||
cv2.line(img_vis, (x2, j), (x2, min(j + dash_length, y2)), (144, 238, 144), 2)
|
||
cv2.putText(img_vis, "R_before", (x1, y1-30), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (144, 238, 144), 2)
|
||
|
||
# Draw cleaned bboxes (solid) and keypoints
|
||
if frame_data['left_bbox'] is not None:
|
||
x1, y1, x2, y2 = frame_data['left_bbox'].astype(int)
|
||
cv2.rectangle(img_vis, (x1, y1), (x2, y2), (255, 0, 0), 3) # Blue for left
|
||
cv2.putText(img_vis, "L_after", (x1, y1-10), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 0, 0), 2)
|
||
|
||
# Draw left hand keypoints in blue
|
||
if frame_data['left_keypoints'] is not None:
|
||
for kp in frame_data['left_keypoints']:
|
||
x, y, conf = kp
|
||
if conf > 0.1: # Only draw confident keypoints
|
||
cv2.circle(img_vis, (int(x), int(y)), 3, (255, 0, 0), -1) # Blue circles
|
||
|
||
if frame_data['right_bbox'] is not None:
|
||
x1, y1, x2, y2 = frame_data['right_bbox'].astype(int)
|
||
cv2.rectangle(img_vis, (x1, y1), (x2, y2), (0, 255, 0), 3) # Green for right
|
||
cv2.putText(img_vis, "R_after", (x1, y1-10), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)
|
||
|
||
# Draw right hand keypoints in green
|
||
if frame_data['right_keypoints'] is not None:
|
||
for kp in frame_data['right_keypoints']:
|
||
x, y, conf = kp
|
||
if conf > 0.1: # Only draw confident keypoints
|
||
cv2.circle(img_vis, (int(x), int(y)), 3, (0, 255, 0), -1) # Green circles
|
||
|
||
img_fn = os.path.splitext(os.path.basename(img_path))[0]
|
||
cv2.imwrite(os.path.join(pass2_dir, f'{img_fn}_cleaned.jpg'), img_vis)
|
||
|
||
# Create video from Pass 2 visualizations
|
||
if vis_dir is not None:
|
||
pass2_dir = os.path.join(vis_dir, 'pass2_cleaned_bboxes')
|
||
video_path = os.path.join(vis_dir, 'pass2_cleaned_bboxes.mp4')
|
||
create_video_from_images(pass2_dir, video_path, fps=30)
|
||
|
||
# Plot bbox trajectories to visualize handedness consistency
|
||
plot_bbox_trajectories(raw_data, vis_dir, filename='bbox_trajectories_after_cleaning.png')
|
||
|
||
print("\n" + "="*80)
|
||
print("PASS 2 Complete: Bbox sequences cleaned")
|
||
print("="*80)
|
||
|
||
return raw_data
|
||
|
||
|
||
# ============================================================================
|
||
# PASS 3: Run HaMeR on cleaned bboxes
|
||
# ============================================================================
|
||
|
||
def render_openpose(img, hand_keypoints):
|
||
"""Render hand keypoints (simplified version)."""
|
||
# Implementation from original code (lines 718-730)
|
||
return img # Placeholder
|
||
|
||
|
||
def convert_crop_coords_to_orig_img(bbox, keypoints, crop_size):
|
||
"""Convert cropped coordinates to original image coordinates."""
|
||
cx, cy, h = bbox[:, 0], bbox[:, 1], bbox[:, 2]
|
||
keypoints *= h[..., None, None] / crop_size
|
||
keypoints[:,:,0] = (cx - h/2)[..., None] + keypoints[:,:,0]
|
||
keypoints[:,:,1] = (cy - h/2)[..., None] + keypoints[:,:,1]
|
||
return keypoints
|
||
|
||
|
||
def run_hamer_on_cleaned_bboxes(raw_data, model, model_cfg, renderer, args):
|
||
"""Pass 3: Run HaMeR on cleaned bboxes."""
|
||
print("\n" + "="*80)
|
||
print("PASS 3: Running HaMeR on cleaned bboxes")
|
||
print("="*80)
|
||
|
||
device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
|
||
results_dict = {}
|
||
|
||
for frame_data in tqdm(raw_data, desc="Pass 3: Running HaMeR"):
|
||
img_path = frame_data['img_path']
|
||
frame_idx = frame_data['frame_idx']
|
||
img_cv2 = cv2.imread(img_path)
|
||
img_fn = os.path.splitext(os.path.basename(img_path))[0]
|
||
|
||
# Prepare bboxes
|
||
bboxes = []
|
||
is_right = []
|
||
vit_keypoints_list = []
|
||
|
||
if frame_data['left_bbox'] is not None:
|
||
bboxes.append(frame_data['left_bbox'])
|
||
is_right.append(0)
|
||
vit_keypoints_list.append(frame_data['left_keypoints'])
|
||
|
||
if frame_data['right_bbox'] is not None:
|
||
bboxes.append(frame_data['right_bbox'])
|
||
is_right.append(1)
|
||
vit_keypoints_list.append(frame_data['right_keypoints'])
|
||
|
||
if len(bboxes) == 0:
|
||
results_dict[img_path] = {
|
||
'mano': [],
|
||
'cam_trans': [],
|
||
'tracked_ids': [],
|
||
'tracked_time': [],
|
||
'extra_data': [],
|
||
'tid': np.array([]),
|
||
'shot': 0
|
||
}
|
||
continue
|
||
|
||
boxes = np.array(bboxes)
|
||
right = np.array(is_right)
|
||
|
||
# Validate and fix keypoints (create dummy keypoints if None)
|
||
fixed_keypoints_list = []
|
||
for i, kp in enumerate(vit_keypoints_list):
|
||
if kp is None:
|
||
# Create dummy keypoints with 0.5 confidence
|
||
dummy_keypoints = np.zeros((21, 3))
|
||
dummy_keypoints[:, 2] = 0.5
|
||
fixed_keypoints_list.append(dummy_keypoints)
|
||
elif not isinstance(kp, np.ndarray):
|
||
raise ValueError(f"Frame {frame_idx} ({img_fn}): keypoints are not numpy array for hand {i}, got {type(kp)}")
|
||
elif kp.shape != (21, 3):
|
||
raise ValueError(f"Frame {frame_idx} ({img_fn}): keypoints have wrong shape {kp.shape} for hand {i}, expected (21, 3)")
|
||
else:
|
||
fixed_keypoints_list.append(kp)
|
||
|
||
vit_keypoints = np.stack(fixed_keypoints_list, axis=0)
|
||
|
||
# Create dataset
|
||
dataset = ViTDetDataset(model_cfg, img_cv2, boxes, right, vit_keypoints, rescale_factor=args.rescale_factor)
|
||
dataloader = torch.utils.data.DataLoader(dataset, batch_size=8, shuffle=False, num_workers=0)
|
||
|
||
# # Save HaMeR input patches for debugging
|
||
# if args.render and args.res_folder is not None:
|
||
# hamer_input_dir = os.path.join(os.path.dirname(args.res_folder), f'hamer_input_patches_{model_cfg.EXTRA.FOCAL_LENGTH}')
|
||
# os.makedirs(hamer_input_dir, exist_ok=True)
|
||
|
||
# for hand_idx, (bbox, is_right_val) in enumerate(zip(boxes, right)):
|
||
# # Get the cropped patch from dataset
|
||
# sample = dataset[hand_idx]
|
||
# img_patch = sample['img_patch'].astype(np.uint8) # This is a tensor (C, H, W) normalized
|
||
|
||
# # Denormalize and convert back to image
|
||
# print('img_patch_bgr: ', img_patch.shape)
|
||
# img_patch_bgr = cv2.cvtColor(img_patch, cv2.COLOR_RGB2BGR)
|
||
|
||
# hand_name = 'right' if is_right_val == 1 else 'left'
|
||
# patch_filename = f'{img_fn}_{hand_name}.jpg'
|
||
# cv2.imwrite(os.path.join(hamer_input_dir, patch_filename), img_patch_bgr)
|
||
|
||
# Run HaMeR
|
||
all_verts = []
|
||
all_cam_t = []
|
||
all_right = []
|
||
all_mano_params = []
|
||
all_pred_2d = []
|
||
all_bboxes = []
|
||
|
||
for batch in dataloader:
|
||
batch = recursive_to(batch, device)
|
||
with torch.no_grad():
|
||
out = model(batch)
|
||
|
||
multiplier = (2 * batch['right'] - 1)
|
||
pred_cam = out['pred_cam']
|
||
pred_cam[:, 1] = multiplier * pred_cam[:, 1]
|
||
|
||
box_center = batch["box_center"].float()
|
||
box_size = batch["box_size"].float()
|
||
img_size = batch["img_size"].float()
|
||
scaled_focal_length = model_cfg.EXTRA.FOCAL_LENGTH / model_cfg.MODEL.IMAGE_SIZE * img_size.max()
|
||
|
||
# Compute camera translation (SLAHMR-style)
|
||
batch_size = batch['img'].shape[0]
|
||
pred_cam_t_full = torch.zeros(batch_size, 3, device=pred_cam.device)
|
||
|
||
for n in range(batch_size):
|
||
cam = pred_cam[n]
|
||
H, W = img_size[n, 1], img_size[n, 0]
|
||
focal = scaled_focal_length[n] if scaled_focal_length.ndim > 0 else scaled_focal_length
|
||
cx, cy = box_center[n]
|
||
scale = box_size[n]
|
||
|
||
tz = 2 * focal / (scale * cam[0] + 1e-6)
|
||
tx = cam[1] + tz / focal * (cx - W / 2)
|
||
ty = cam[2] + tz / focal * (cy - H / 2)
|
||
|
||
pred_cam_t_full[n] = torch.tensor([tx, ty, tz], device=pred_cam.device)
|
||
|
||
pred_cam_t_full = pred_cam_t_full.detach().cpu().numpy()
|
||
|
||
for n in range(batch_size):
|
||
verts = out['pred_vertices'][n].detach().cpu().numpy()
|
||
pred_joints = out['pred_keypoints_2d'][n].detach().cpu().numpy()
|
||
is_right_val = int(batch['right'][n].cpu().numpy())
|
||
verts[:, 0] = (2 * is_right_val - 1) * verts[:, 0]
|
||
pred_joints[:, 0] = (2 * is_right_val - 1) * pred_joints[:, 0]
|
||
cam_t = pred_cam_t_full[n]
|
||
|
||
all_verts.append(verts)
|
||
all_cam_t.append(cam_t)
|
||
all_right.append(is_right_val)
|
||
all_pred_2d.append(pred_joints)
|
||
all_bboxes.append(batch['bbox'][n].detach().cpu().numpy())
|
||
|
||
mano_params = out['pred_mano_params'][n]
|
||
mano_params['is_right'] = is_right_val
|
||
all_mano_params.append(mano_params)
|
||
|
||
# Convert 2D keypoints to original image coordinates
|
||
if len(all_pred_2d) > 0:
|
||
all_pred_2d_np = np.stack(all_pred_2d)
|
||
all_bboxes_np = np.stack(all_bboxes)
|
||
|
||
# Add confidence column
|
||
v = np.ones((all_pred_2d_np.shape[0], all_pred_2d_np.shape[1], 1))
|
||
all_pred_2d_np = np.concatenate((all_pred_2d_np, v), axis=-1)
|
||
|
||
# Convert from crop coords to original image coords
|
||
all_pred_2d_np = model_cfg.MODEL.IMAGE_SIZE * (all_pred_2d_np + 0.5)
|
||
all_pred_2d_np = convert_crop_coords_to_orig_img(bbox=all_bboxes_np, keypoints=all_pred_2d_np, crop_size=model_cfg.MODEL.IMAGE_SIZE)
|
||
all_pred_2d_np[:, :, -1] = 1 # Set all confidences to 1
|
||
|
||
extra_data = [all_pred_2d_np[i].tolist() for i in range(len(all_pred_2d_np))]
|
||
else:
|
||
extra_data = []
|
||
|
||
# Store results
|
||
results_dict[img_path] = {
|
||
'mano': all_mano_params,
|
||
'cam_trans': all_cam_t,
|
||
'tracked_ids': all_right,
|
||
'tracked_time': [0] * len(all_right),
|
||
'extra_data': extra_data,
|
||
'tid': np.array(all_right),
|
||
'shot': 0
|
||
}
|
||
|
||
# Render hand meshes if requested
|
||
if args.render and len(all_verts) > 0 and renderer is not None:
|
||
LIGHT_BLUE = (0.65098039, 0.74117647, 0.85882353)
|
||
misc_args = dict(
|
||
mesh_base_color=LIGHT_BLUE,
|
||
scene_bg_color=(1, 1, 1),
|
||
focal_length=scaled_focal_length,
|
||
)
|
||
|
||
cam_view, _ = renderer.render_rgba_multiple(
|
||
all_verts,
|
||
cam_t=all_cam_t,
|
||
render_res=img_size[0],
|
||
is_right=all_right,
|
||
**misc_args
|
||
)
|
||
|
||
# Overlay on original image
|
||
input_img = img_cv2.astype(np.float32)[:, :, ::-1] / 255.0
|
||
input_img = np.concatenate([input_img, np.ones_like(input_img[:, :, :1])], axis=2)
|
||
input_img_overlay = input_img[:, :, :3] * (1 - cam_view[:, :, 3:]) + cam_view[:, :, :3] * cam_view[:, :, 3:]
|
||
|
||
# Save rendered result
|
||
render_path = os.path.join(os.path.dirname(args.res_folder), f'render_all_{model_cfg.EXTRA.FOCAL_LENGTH}')
|
||
os.makedirs(render_path, exist_ok=True)
|
||
cv2.imwrite(os.path.join(render_path, f'{img_fn}.jpg'), 255 * input_img_overlay[:, :, ::-1])
|
||
|
||
# Create video from Pass 3 rendered results
|
||
render_path = os.path.join(os.path.dirname(args.res_folder), f'render_all_{model_cfg.EXTRA.FOCAL_LENGTH}')
|
||
if os.path.exists(render_path):
|
||
video_path = os.path.join(os.path.dirname(args.res_folder), f'render_all_{model_cfg.EXTRA.FOCAL_LENGTH}.mp4')
|
||
create_video_from_images(render_path, video_path, fps=30)
|
||
|
||
print("\n" + "="*80)
|
||
print("PASS 3 Complete: HaMeR reconstruction finished")
|
||
print("="*80)
|
||
|
||
return results_dict
|
||
|
||
|
||
# ============================================================================
|
||
# Main
|
||
# ============================================================================
|
||
|
||
def main():
|
||
parser = argparse.ArgumentParser(description='HaMeR with 3-pass architecture using YOLO hand detector')
|
||
parser.add_argument('--checkpoint', type=str, default=None, help='Path to pretrained model checkpoint')
|
||
parser.add_argument('--img_folder', type=str, default='images', help='Folder with input images')
|
||
parser.add_argument('--out_folder', type=str, default='out_demo', help='Output folder to save rendered results')
|
||
parser.add_argument('--res_folder', type=str, help='Output folder to save rendered results')
|
||
parser.add_argument('--side_view', dest='side_view', action='store_true', default=False, help='If set, render side view also')
|
||
parser.add_argument('--full_frame', dest='full_frame', action='store_true', default=True, help='If set, render all people together also')
|
||
parser.add_argument('--save_mesh', dest='save_mesh', action='store_true', default=False, help='If set, save meshes to disk also')
|
||
parser.add_argument('--batch_size', type=int, default=1, help='Batch size for inference/fitting')
|
||
parser.add_argument('--rescale_factor', type=float, default=1.3, help='Factor for padding the bbox')
|
||
parser.add_argument('--file_type', nargs='+', default=['*.jpg', '*.png'], help='List of file extensions to consider')
|
||
parser.add_argument('--conf', type=float, default=2.0, help='Factor for padding the bbox')
|
||
parser.add_argument('--type', type=str, default='EgoDexter', help='Path to pretrained model checkpoint')
|
||
parser.add_argument('--render', dest='render', action='store_true', default=False, help='If set, render side view also')
|
||
parser.add_argument('--yolo_model', type=str, default='./pretrained_models/detector.pt',
|
||
help='Path to YOLO hand detector model (like WiLoR detector.pt)')
|
||
|
||
args = parser.parse_args()
|
||
|
||
# Setup
|
||
device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
|
||
img_folder = Path(args.img_folder)
|
||
|
||
# Get all images matching file_type patterns
|
||
img_paths = sorted([img for end in args.file_type for img in img_folder.glob(end)])
|
||
|
||
print(f"Found {len(img_paths)} images in {img_folder}")
|
||
|
||
# Load models
|
||
print("\nLoading models...")
|
||
CACHE_DIR_HAMER = args.checkpoint
|
||
download_models(CACHE_DIR_HAMER)
|
||
model, model_cfg = load_hamer(args.checkpoint)
|
||
model = model.to(device)
|
||
model.eval()
|
||
|
||
# Load YOLO hand detector (like WiLoR)
|
||
print(f"\nLoading YOLO hand detector: {args.yolo_model}")
|
||
yolo_detector = YOLO(args.yolo_model)
|
||
yolo_detector.to(device)
|
||
|
||
# Load renderer
|
||
renderer = Renderer(model_cfg, faces=model.mano.faces)
|
||
|
||
# Create visualization directory
|
||
vis_dir = None
|
||
if args.render and args.res_folder is not None:
|
||
vis_dir = os.path.join(os.path.dirname(args.res_folder), f'bbox_vis_{model_cfg.EXTRA.FOCAL_LENGTH}')
|
||
os.makedirs(vis_dir, exist_ok=True)
|
||
print(f"\nBbox visualizations will be saved to: {vis_dir}")
|
||
|
||
# PASS 1: Extract raw bboxes using YOLO
|
||
raw_data = extract_raw_bboxes(img_paths, yolo_detector, vis_dir=vis_dir)
|
||
|
||
# PASS 2: Clean bbox sequences
|
||
cleaned_data = clean_bbox_sequences(raw_data, vis_dir=vis_dir)
|
||
|
||
# PASS 3: Run HaMeR
|
||
results = run_hamer_on_cleaned_bboxes(cleaned_data, model, model_cfg, renderer, args)
|
||
|
||
# Save results
|
||
output_path = Path(args.res_folder)
|
||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||
with open(output_path, 'wb') as f:
|
||
pickle.dump(results, f)
|
||
|
||
print(f"\nResults saved to {output_path}")
|
||
|
||
if __name__ == '__main__':
|
||
main()
|