436 lines
15 KiB
Python
436 lines
15 KiB
Python
|
import argparse
|
||
|
|
import json
|
||
|
|
import os
|
||
|
|
import sys
|
||
|
|
import time
|
||
|
|
import urllib.request
|
||
|
|
from pathlib import Path
|
||
|
|
|
||
|
|
import cv2
|
||
|
|
import numpy as np
|
||
|
|
import pygame
|
||
|
|
from pygame.locals import *
|
||
|
|
|
||
|
|
# Define colors
|
||
|
|
BLACK = (0, 0, 0)
|
||
|
|
WHITE = (255, 255, 255)
|
||
|
|
RED = (255, 0, 0)
|
||
|
|
GREEN = (0, 255, 0)
|
||
|
|
BLUE = (0, 0, 255)
|
||
|
|
YELLOW = (255, 255, 0)
|
||
|
|
CYAN = (0, 255, 255)
|
||
|
|
MAGENTA = (255, 0, 255)
|
||
|
|
|
||
|
|
# Define keypoint colors (custom palette)
|
||
|
|
KEYPOINT_COLORS = [
|
||
|
|
(255, 0, 0), # nose (red)
|
||
|
|
(255, 85, 0), # left_eye (orange-red)
|
||
|
|
(255, 170, 0), # right_eye (orange)
|
||
|
|
(255, 255, 0), # left_ear (yellow)
|
||
|
|
(170, 255, 0), # right_ear (yellow-green)
|
||
|
|
(85, 255, 0), # left_shoulder (green-yellow)
|
||
|
|
(0, 255, 0), # right_shoulder (green)
|
||
|
|
(0, 255, 85), # left_elbow (green-cyan)
|
||
|
|
(0, 255, 170), # right_elbow (cyan-green)
|
||
|
|
(0, 255, 255), # left_wrist (cyan)
|
||
|
|
(0, 170, 255), # right_wrist (cyan-blue)
|
||
|
|
(0, 85, 255), # left_hip (blue-cyan)
|
||
|
|
(0, 0, 255), # right_hip (blue)
|
||
|
|
(85, 0, 255), # left_knee (blue-purple)
|
||
|
|
(170, 0, 255), # right_knee (purple-blue)
|
||
|
|
(255, 0, 255), # left_ankle (magenta)
|
||
|
|
(255, 0, 170) # right_ankle (magenta-pink)
|
||
|
|
]
|
||
|
|
|
||
|
|
# Person ID colors
|
||
|
|
PERSON_COLORS = [
|
||
|
|
(255, 0, 0), # red
|
||
|
|
(0, 255, 0), # green
|
||
|
|
(0, 0, 255), # blue
|
||
|
|
(255, 255, 0), # yellow
|
||
|
|
(255, 0, 255), # magenta
|
||
|
|
(0, 255, 255), # cyan
|
||
|
|
(255, 128, 0), # orange
|
||
|
|
(128, 0, 255), # purple
|
||
|
|
(0, 255, 128), # mint
|
||
|
|
(255, 255, 255) # white
|
||
|
|
]
|
||
|
|
|
||
|
|
def download_video(url, output_dir="downloaded_videos"):
|
||
|
|
"""Download a video from a URL and return the local file path"""
|
||
|
|
os.makedirs(output_dir, exist_ok=True)
|
||
|
|
video_name = os.path.basename(url).split("?")[0]
|
||
|
|
if not video_name or "." not in video_name:
|
||
|
|
video_name = f"video_{int(time.time())}.mp4"
|
||
|
|
|
||
|
|
output_path = os.path.join(output_dir, video_name)
|
||
|
|
if os.path.exists(output_path):
|
||
|
|
print(f"✅ Video already downloaded: {output_path}")
|
||
|
|
return output_path
|
||
|
|
|
||
|
|
print(f"⬇️ Downloading video from {url} to {output_path}...")
|
||
|
|
urllib.request.urlretrieve(url, output_path)
|
||
|
|
print(f"✅ Video downloaded successfully to {output_path}")
|
||
|
|
return output_path
|
||
|
|
|
||
|
|
def load_pose_data(json_file):
|
||
|
|
"""Load pose data from a JSON file"""
|
||
|
|
print(f"📂 Loading pose data from {json_file}...")
|
||
|
|
with open(json_file, 'r') as f:
|
||
|
|
data = json.load(f)
|
||
|
|
|
||
|
|
# Extract metadata
|
||
|
|
width = data.get('w', 1280)
|
||
|
|
height = data.get('h', 720)
|
||
|
|
fps = data.get('fps', 30)
|
||
|
|
total_frames = data.get('frames', 0)
|
||
|
|
|
||
|
|
# Get frame sampling and precision from metadata if available
|
||
|
|
metadata = data.get('meta', {})
|
||
|
|
frame_sampling = metadata.get('frame_sampling', 1)
|
||
|
|
precision = metadata.get('precision', 3)
|
||
|
|
|
||
|
|
# Extract connections
|
||
|
|
connections = []
|
||
|
|
for conn in data.get('connections', []):
|
||
|
|
start = conn.get('s', 0)
|
||
|
|
end = conn.get('e', 0)
|
||
|
|
connections.append((start, end))
|
||
|
|
|
||
|
|
# Extract keypoint names
|
||
|
|
keypoint_names = data.get('keypoints', [])
|
||
|
|
|
||
|
|
# Extract frame data
|
||
|
|
frames = data.get('data', [])
|
||
|
|
|
||
|
|
print(f"✅ Loaded {len(frames)} frames of pose data")
|
||
|
|
print(f"📊 Video: {width}x{height}@{fps}fps, {total_frames} total frames")
|
||
|
|
print(f"🔍 Frame sampling: {frame_sampling}, Precision: {precision}")
|
||
|
|
|
||
|
|
return {
|
||
|
|
'width': width,
|
||
|
|
'height': height,
|
||
|
|
'fps': fps,
|
||
|
|
'total_frames': total_frames,
|
||
|
|
'frame_sampling': frame_sampling,
|
||
|
|
'precision': precision,
|
||
|
|
'connections': connections,
|
||
|
|
'keypoint_names': keypoint_names,
|
||
|
|
'frames': frames
|
||
|
|
}
|
||
|
|
|
||
|
|
def create_pygame_window(width, height, title="Pose Viewer"):
|
||
|
|
"""Create a PyGame window"""
|
||
|
|
pygame.init()
|
||
|
|
window = pygame.display.set_mode((width, height))
|
||
|
|
pygame.display.set_caption(title)
|
||
|
|
return window
|
||
|
|
|
||
|
|
def draw_pose(frame, pose_data, frame_idx, original_width, original_height, prev_frame_idx=None):
|
||
|
|
"""Draw pose data on a given frame"""
|
||
|
|
# Clone the frame to avoid modifying the original
|
||
|
|
pose_frame = frame.copy()
|
||
|
|
|
||
|
|
# Find the closest pose frame to the current video frame (should be first and only in temp_pose_data)
|
||
|
|
if pose_data['frames']:
|
||
|
|
closest_frame = pose_data['frames'][0]
|
||
|
|
|
||
|
|
connections = pose_data['connections']
|
||
|
|
|
||
|
|
# Draw each person
|
||
|
|
for person_idx, person in enumerate(closest_frame['p']):
|
||
|
|
person_id = person['id']
|
||
|
|
person_color = PERSON_COLORS[person_id % len(PERSON_COLORS)]
|
||
|
|
|
||
|
|
# Get keypoints
|
||
|
|
keypoints = person['k']
|
||
|
|
|
||
|
|
# Create a dictionary to store keypoints by index
|
||
|
|
kp_dict = {}
|
||
|
|
for kp in keypoints:
|
||
|
|
kp_dict[kp[0]] = (
|
||
|
|
int(kp[1] * original_width),
|
||
|
|
int(kp[2] * original_height),
|
||
|
|
kp[3]
|
||
|
|
)
|
||
|
|
|
||
|
|
# Draw connections
|
||
|
|
for conn in connections:
|
||
|
|
if conn[0] in kp_dict and conn[1] in kp_dict:
|
||
|
|
start_point = kp_dict[conn[0]][:2]
|
||
|
|
end_point = kp_dict[conn[1]][:2]
|
||
|
|
|
||
|
|
# Use average confidence to determine line thickness
|
||
|
|
avg_conf = (kp_dict[conn[0]][2] + kp_dict[conn[1]][2]) / 2
|
||
|
|
thickness = int(avg_conf * 3) + 1
|
||
|
|
|
||
|
|
cv2.line(pose_frame, start_point, end_point, person_color, thickness)
|
||
|
|
|
||
|
|
# Draw keypoints
|
||
|
|
for kp_idx, (x, y, conf) in kp_dict.items():
|
||
|
|
# Circle size based on confidence
|
||
|
|
radius = int(conf * 5) + 2
|
||
|
|
cv2.circle(pose_frame, (x, y), radius, KEYPOINT_COLORS[kp_idx % len(KEYPOINT_COLORS)], -1)
|
||
|
|
|
||
|
|
# Draw person ID
|
||
|
|
bbox = person['b']
|
||
|
|
x, y = int(bbox[0]), int(bbox[1])
|
||
|
|
cv2.putText(
|
||
|
|
pose_frame,
|
||
|
|
f"ID: {person_id}",
|
||
|
|
(x, y - 10),
|
||
|
|
cv2.FONT_HERSHEY_SIMPLEX,
|
||
|
|
0.7,
|
||
|
|
person_color,
|
||
|
|
2
|
||
|
|
)
|
||
|
|
|
||
|
|
return pose_frame
|
||
|
|
|
||
|
|
def draw_ui_controls(surface, width, height, playing, current_frame, total_frames):
|
||
|
|
"""Draw UI controls on the PyGame surface"""
|
||
|
|
# Background for controls
|
||
|
|
control_height = 50
|
||
|
|
control_surface = pygame.Surface((width, control_height))
|
||
|
|
control_surface.fill(BLACK)
|
||
|
|
|
||
|
|
# Draw play/pause button
|
||
|
|
button_width = 80
|
||
|
|
button_height = 30
|
||
|
|
button_x = 20
|
||
|
|
button_y = (control_height - button_height) // 2
|
||
|
|
|
||
|
|
pygame.draw.rect(control_surface, BLUE, (button_x, button_y, button_width, button_height))
|
||
|
|
|
||
|
|
font = pygame.font.SysFont(None, 24)
|
||
|
|
text = font.render("Pause" if playing else "Play", True, WHITE)
|
||
|
|
text_rect = text.get_rect(center=(button_x + button_width//2, button_y + button_height//2))
|
||
|
|
control_surface.blit(text, text_rect)
|
||
|
|
|
||
|
|
# Draw stop button
|
||
|
|
stop_button_x = button_x + button_width + 20
|
||
|
|
pygame.draw.rect(control_surface, RED, (stop_button_x, button_y, button_width, button_height))
|
||
|
|
|
||
|
|
stop_text = font.render("Stop", True, WHITE)
|
||
|
|
stop_text_rect = stop_text.get_rect(center=(stop_button_x + button_width//2, button_y + button_height//2))
|
||
|
|
control_surface.blit(stop_text, stop_text_rect)
|
||
|
|
|
||
|
|
# Draw seek bar
|
||
|
|
seekbar_x = stop_button_x + button_width + 40
|
||
|
|
seekbar_y = button_y + button_height // 2
|
||
|
|
seekbar_width = width - seekbar_x - 40
|
||
|
|
seekbar_height = 10
|
||
|
|
|
||
|
|
# Background bar
|
||
|
|
pygame.draw.rect(control_surface, (100, 100, 100),
|
||
|
|
(seekbar_x, seekbar_y - seekbar_height//2, seekbar_width, seekbar_height))
|
||
|
|
|
||
|
|
# Progress bar
|
||
|
|
progress = current_frame / total_frames if total_frames > 0 else 0
|
||
|
|
progress_width = int(seekbar_width * progress)
|
||
|
|
pygame.draw.rect(control_surface, GREEN,
|
||
|
|
(seekbar_x, seekbar_y - seekbar_height//2, progress_width, seekbar_height))
|
||
|
|
|
||
|
|
# Display current time / total time
|
||
|
|
time_text = font.render(f"Frame: {current_frame} / {total_frames}", True, WHITE)
|
||
|
|
time_rect = time_text.get_rect(center=(seekbar_x + seekbar_width//2, seekbar_y - 20))
|
||
|
|
control_surface.blit(time_text, time_rect)
|
||
|
|
|
||
|
|
# Blit the control surface to the main surface
|
||
|
|
surface.blit(control_surface, (0, height - control_height))
|
||
|
|
|
||
|
|
# Return button regions for click handling
|
||
|
|
play_button_rect = pygame.Rect(button_x, height - control_height + button_y, button_width, button_height)
|
||
|
|
stop_button_rect = pygame.Rect(stop_button_x, height - control_height + button_y, button_width, button_height)
|
||
|
|
seekbar_rect = pygame.Rect(seekbar_x, height - control_height + seekbar_y - seekbar_height//2,
|
||
|
|
seekbar_width, seekbar_height)
|
||
|
|
|
||
|
|
return play_button_rect, stop_button_rect, seekbar_rect
|
||
|
|
|
||
|
|
def run_viewer(video_path, json_path):
|
||
|
|
"""Main function to run the pose viewer"""
|
||
|
|
# Load pose data
|
||
|
|
pose_data = load_pose_data(json_path)
|
||
|
|
|
||
|
|
# Open video
|
||
|
|
cap = cv2.VideoCapture(video_path)
|
||
|
|
if not cap.isOpened():
|
||
|
|
print(f"❌ Error: Could not open video {video_path}")
|
||
|
|
return
|
||
|
|
|
||
|
|
# Get video properties
|
||
|
|
video_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
|
||
|
|
video_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
|
||
|
|
fps = cap.get(cv2.CAP_PROP_FPS)
|
||
|
|
total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
|
||
|
|
|
||
|
|
# Create window - now only showing single visualization
|
||
|
|
window_width = video_width
|
||
|
|
window_height = video_height + 50 # Additional space for controls
|
||
|
|
window = create_pygame_window(window_width, window_height, f"Pose Viewer - {os.path.basename(video_path)}")
|
||
|
|
|
||
|
|
# Setup clock
|
||
|
|
clock = pygame.time.Clock()
|
||
|
|
|
||
|
|
# State variables
|
||
|
|
playing = False
|
||
|
|
current_frame = 0
|
||
|
|
prev_frame_idx = None
|
||
|
|
last_valid_pose_frame = None
|
||
|
|
|
||
|
|
# Create a frame lookup for fast access - maps video frame number to pose data frame
|
||
|
|
frame_lookup = {}
|
||
|
|
max_pose_frame = 0
|
||
|
|
for pose_frame in pose_data['frames']:
|
||
|
|
frame_num = pose_frame['f']
|
||
|
|
frame_lookup[frame_num] = pose_frame
|
||
|
|
max_pose_frame = max(max_pose_frame, frame_num)
|
||
|
|
|
||
|
|
# Initial render
|
||
|
|
ret, frame = cap.read()
|
||
|
|
if not ret:
|
||
|
|
print("❌ Error: Could not read the first frame")
|
||
|
|
return
|
||
|
|
|
||
|
|
# Main loop
|
||
|
|
while True:
|
||
|
|
# Handle events
|
||
|
|
for event in pygame.event.get():
|
||
|
|
if event.type == QUIT:
|
||
|
|
pygame.quit()
|
||
|
|
sys.exit()
|
||
|
|
elif event.type == KEYDOWN:
|
||
|
|
if event.key == K_ESCAPE:
|
||
|
|
pygame.quit()
|
||
|
|
sys.exit()
|
||
|
|
elif event.key == K_SPACE:
|
||
|
|
playing = not playing
|
||
|
|
elif event.type == MOUSEBUTTONDOWN:
|
||
|
|
# Check if any buttons were clicked
|
||
|
|
mouse_pos = pygame.mouse.get_pos()
|
||
|
|
play_button_rect, stop_button_rect, seekbar_rect = draw_ui_controls(
|
||
|
|
window, window_width, window_height, playing, current_frame, total_frames
|
||
|
|
)
|
||
|
|
|
||
|
|
if play_button_rect.collidepoint(mouse_pos):
|
||
|
|
playing = not playing
|
||
|
|
elif stop_button_rect.collidepoint(mouse_pos):
|
||
|
|
playing = False
|
||
|
|
current_frame = 0
|
||
|
|
cap.set(cv2.CAP_PROP_POS_FRAMES, current_frame)
|
||
|
|
ret, frame = cap.read()
|
||
|
|
prev_frame_idx = None
|
||
|
|
last_valid_pose_frame = None
|
||
|
|
elif seekbar_rect.collidepoint(mouse_pos):
|
||
|
|
# Calculate position ratio
|
||
|
|
x_offset = mouse_pos[0] - seekbar_rect.x
|
||
|
|
ratio = x_offset / seekbar_rect.width
|
||
|
|
|
||
|
|
# Set frame position
|
||
|
|
current_frame = int(ratio * total_frames)
|
||
|
|
cap.set(cv2.CAP_PROP_POS_FRAMES, current_frame)
|
||
|
|
ret, frame = cap.read()
|
||
|
|
prev_frame_idx = None # Reset previous frame index after seeking
|
||
|
|
|
||
|
|
# Handle playback
|
||
|
|
if playing:
|
||
|
|
ret, frame = cap.read()
|
||
|
|
if not ret:
|
||
|
|
# End of video, loop back to start
|
||
|
|
playing = False
|
||
|
|
current_frame = 0
|
||
|
|
cap.set(cv2.CAP_PROP_POS_FRAMES, current_frame)
|
||
|
|
ret, frame = cap.read()
|
||
|
|
prev_frame_idx = None
|
||
|
|
last_valid_pose_frame = None
|
||
|
|
if not ret:
|
||
|
|
break
|
||
|
|
current_frame += 1
|
||
|
|
|
||
|
|
# Find the appropriate pose frame for the current video frame
|
||
|
|
frame_sampling = pose_data['frame_sampling']
|
||
|
|
|
||
|
|
# Try to find the exact frame in lookup
|
||
|
|
pose_frame = frame_lookup.get(current_frame)
|
||
|
|
|
||
|
|
# If not found, find the closest previous frame based on sampling
|
||
|
|
if not pose_frame:
|
||
|
|
# Calculate what the nearest pose frame should be
|
||
|
|
# This searches for the most recent pose frame
|
||
|
|
nearest_frame = current_frame
|
||
|
|
while nearest_frame > 0 and nearest_frame not in frame_lookup:
|
||
|
|
nearest_frame -= 1
|
||
|
|
|
||
|
|
if nearest_frame in frame_lookup:
|
||
|
|
pose_frame = frame_lookup[nearest_frame]
|
||
|
|
|
||
|
|
# Update the last valid pose frame if we found one
|
||
|
|
if pose_frame:
|
||
|
|
last_valid_pose_frame = pose_frame
|
||
|
|
|
||
|
|
# Draw pose on frame - use the most recent valid pose frame
|
||
|
|
if last_valid_pose_frame:
|
||
|
|
# Create a special frame dict with only the current pose for draw_pose
|
||
|
|
temp_pose_data = pose_data.copy()
|
||
|
|
temp_pose_data['frames'] = [last_valid_pose_frame]
|
||
|
|
|
||
|
|
pose_frame = draw_pose(frame, temp_pose_data, current_frame, video_width, video_height)
|
||
|
|
else:
|
||
|
|
# If no pose data found yet, just show the original frame
|
||
|
|
pose_frame = frame.copy()
|
||
|
|
|
||
|
|
prev_frame_idx = current_frame
|
||
|
|
|
||
|
|
# Convert frame from BGR to RGB for PyGame
|
||
|
|
rgb_pose_frame = cv2.cvtColor(pose_frame, cv2.COLOR_BGR2RGB)
|
||
|
|
pygame_pose_frame = pygame.surfarray.make_surface(rgb_pose_frame.swapaxes(0, 1))
|
||
|
|
|
||
|
|
# Draw frame
|
||
|
|
window.blit(pygame_pose_frame, (0, 0))
|
||
|
|
|
||
|
|
# Draw UI controls
|
||
|
|
play_button_rect, stop_button_rect, seekbar_rect = draw_ui_controls(
|
||
|
|
window, window_width, window_height, playing, current_frame, total_frames
|
||
|
|
)
|
||
|
|
|
||
|
|
# Draw metadata
|
||
|
|
font = pygame.font.SysFont(None, 20)
|
||
|
|
metadata_text = f"Frame Sampling: {pose_data['frame_sampling']}, Precision: {pose_data['precision']}"
|
||
|
|
metadata_surface = font.render(metadata_text, True, WHITE)
|
||
|
|
window.blit(metadata_surface, (10, 10))
|
||
|
|
|
||
|
|
# Update display
|
||
|
|
pygame.display.flip()
|
||
|
|
|
||
|
|
# Cap framerate
|
||
|
|
clock.tick(fps)
|
||
|
|
|
||
|
|
# Clean up
|
||
|
|
cap.release()
|
||
|
|
pygame.quit()
|
||
|
|
|
||
|
|
def main():
|
||
|
|
parser = argparse.ArgumentParser(description='Pose Viewer for JSON pose data with video')
|
||
|
|
parser.add_argument('--video', '-v', required=True, help='Video file path or URL')
|
||
|
|
parser.add_argument('--json', '-j', required=True, help='JSON pose data file path')
|
||
|
|
|
||
|
|
args = parser.parse_args()
|
||
|
|
|
||
|
|
# Handle URL input for video
|
||
|
|
video_path = args.video
|
||
|
|
if video_path.startswith('http://') or video_path.startswith('https://'):
|
||
|
|
video_path = download_video(video_path)
|
||
|
|
|
||
|
|
if not os.path.exists(video_path):
|
||
|
|
print(f"❌ Error: Video file not found: {video_path}")
|
||
|
|
return
|
||
|
|
|
||
|
|
if not os.path.exists(args.json):
|
||
|
|
print(f"❌ Error: JSON file not found: {args.json}")
|
||
|
|
return
|
||
|
|
|
||
|
|
# Run the viewer
|
||
|
|
run_viewer(video_path, args.json)
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
main()
|