import json
import os
import numpy as np
import imageio

if __name__ == '__main__':
    hpt_dataset_dir = 'data/language_table'
    metadata = json.load(open(os.path.join(hpt_dataset_dir, "metadata.json")))
    image_width = metadata["w"]
    image_height = metadata["h"]
    num_frames = metadata["num_images"]
    frames = np.fromfile(os.path.join(hpt_dataset_dir, "video.bin"), dtype=np.uint8
        ).reshape(num_frames, image_height, image_width, 3)
    
    # random sample N frames
    N = 30
    sample_indices = np.random.choice(num_frames, N, replace=False)
    sample_frames = frames[sample_indices]
    
    for i, f in enumerate(sample_frames):
        imageio.imwrite(f"sim/assets/langtable_prompt/frame_{i:02d}.png", f)