bdanko/how2sign-landmarks-front-raw-parquet

import cv2
import numpy as np
import tempfile
import os
from datasets import load_dataset
from IPython.display import Video, display

# 1. Initialize Streams
print("Initializing streams...")
landmark_stream = load_dataset("bdanko/how2sign-landmarks-front-raw-parquet", split="train", streaming=True)
rgb_stream = load_dataset("bdanko/how2sign-rgb-front-clips", split="train", streaming=True)

# 2. Get Landmark Sample
landmark_sample = next(iter(landmark_stream))
target_id = landmark_sample['video_id']
landmarks = np.frombuffer(landmark_sample['features'], dtype=np.float32).reshape(landmark_sample['shape'])

# 3. Find matching RGB clip
print(f"Searching for RGB clip: {target_id}...")
rgb_sample = None
for sample in rgb_stream:
    if sample['__key__'] == target_id:
        rgb_sample = sample
        break

if rgb_sample is None:
    print(f"Error: Could not find RGB clip for {target_id}")
else:
    # 4. Handle Video Data
    video_data = rgb_sample['mp4']
    with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tf:
        if isinstance(video_data, bytes): tf.write(video_data)
        elif isinstance(video_data, dict) and 'bytes' in video_data: tf.write(video_data['bytes'])
        else: # Handle path
            tf.close()
            os.remove(tf.name)
            tf.name = video_data if isinstance(video_data, str) else video_data['path']
        video_path = tf.name

    # 5. Initialize Video Writer
    cap = cv2.VideoCapture(video_path)
    w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    fps = cap.get(cv2.CAP_PROP_FPS) or 25.0
    
    # We'll save to 'output_raw.mp4' first
    output_path = 'overlay_result.mp4'
    fourcc = cv2.VideoWriter_fourcc(*'MP4V') 
    out = cv2.VideoWriter(output_path, fourcc, fps, (w, h))

    print(f"Rendering {len(landmarks)} frames to {output_path}...")
    
    for frame_idx in range(len(landmarks)):
        ret, frame = cap.read()
        if not ret: break
        
        curr_lms = landmarks[frame_idx]
        for i, lm in enumerate(curr_lms):
            x, y = int(lm[0] * w), int(lm[1] * h)
            # Pose=Blue, Face=White, Hands=Green
            color = (255, 0, 0) if i < 33 else (255, 255, 255) if i < 501 else (0, 255, 0)
            cv2.circle(frame, (x, y), 2, color, -1)
        
        out.write(frame)
        if frame_idx % 50 == 0:
            print(f"Progress: {frame_idx}/{len(landmarks)} frames...")

    cap.release()
    out.release()
    
    # 6. Display the result
    print(f"\nProcessing complete! File saved as {output_path}")
    # Note: 'MP4V' is sometimes not playable directly in all browsers. 
    # In Colab, you can use the Video display or download the file from the sidebar.
    display(Video(output_path, embed=True, width=640))