import mediapy
import numpy as np
from pathlib import Path
from PIL import Image

# Load video
video_path = "0804_BimanualStackPlatesOnTableFromDryingRack__riverway__sim__2025-01-07T11-30-44-05-00__000152.mp4"
video = mediapy.read_video(video_path)

# Get dimensions of each subvideo
height = video.shape[1] // 2
width = video.shape[2] // 6

# Extract and save subvideos as gifs
# for i in range(1):
subvideo = video[:, 1*height:2*height, 3*width:4*width]
output_path = Path(f"runs/debug_filtered/subvideo_0_0.gif")
output_path.parent.mkdir(parents=True, exist_ok=True)

# Convert frames to PIL images
frames = [Image.fromarray(frame) for frame in subvideo]

# Save as animated GIF
frames[0].save(
    str(output_path),
    save_all=True,
    append_images=frames[1:],
    duration=1000//15,  # Frame duration in ms (for 15 fps)
    loop=0
)