Generate videos for Bench2Drive data

create_bench2drive_videos.py

import os, tarfile
from subprocess import check_call, DEVNULL
from tempfile import mkdtemp, TemporaryDirectory

def create_video(tar_path, out_path):
    with TemporaryDirectory() as tmpd:
        print("Reading tar:", os.path.basename(tar_path))
        tar = tarfile.open(tar_path, "r:gz")
        tar_members = tar.getmembers()
        front_camera_imgs = [member for member in tar_members if member.name.split('/')[-2] == "rgb_front"]
        top_down_imgs = [member for member in tar_members if member.name.split('/')[-2] == "rgb_top_down"]

        print("Extracting...")
        tar.extractall(tmpd, members=front_camera_imgs + top_down_imgs)
        
        print("Encoding...")
        scene_name = tar_path.rsplit('/', 1)[1].split('.', 1)[0]
        front_camera_path = front_camera_imgs[0].name.rsplit('/', 1)[0]
        top_down_path = top_down_imgs[0].name.rsplit('/', 1)[0]

        common_args = [
            "-c:v", "hevc",
            "-pix_fmt", "yuv420p",
            "-crf", "35", # quality 0(high) - 51(low)
        ]
        check_call([
            "ffmpeg", "-y", "-r", "25", # fps, base=10
        ] + [
            "-i", os.path.join(tmpd, front_camera_path, "%05d.jpg")
        ] + common_args + [
            os.path.join(out_path, scene_name + ".mp4")
        ], stdout=DEVNULL, stderr=DEVNULL)
        check_call([
            "ffmpeg", "-y", "-r", "25", # fps, base=10
        ] + [
            "-i", os.path.join(tmpd, top_down_path, "%05d.jpg")
        ] + common_args + [
            os.path.join(out_path, scene_name + ".topdown.mp4")
        ], stdout=DEVNULL, stderr=DEVNULL)

        print("Completed.")

def create_video_batch(tar_dir):
    for file in os.listdir(tar_dir):
        if file.endswith(".tar.gz"):
            create_video(os.path.join(tar_dir, file), tar_dir)

create_video_batch("/home/jacobz/Documents")