from PIL import Image
import imageio.v2 as imageio
import numpy as np
import os
import re

frame_folder = "./output"
output_mp4 = "output.mp4"
fps = 4

def frame_number(filename):
    m = re.search(r"frame_(\d+)\.png", filename)
    return int(m.group(1))

files = sorted(
    [
        f for f in os.listdir(frame_folder)
        if f.startswith("frame_") and f.endswith(".png")
    ],
    key=frame_number,
)

with imageio.get_writer(
    output_mp4,
    fps=fps,
    codec="libx264",
    pixelformat="yuv420p",
) as writer:
    for f in files:
        path = os.path.join(frame_folder, f)

        img = Image.open(path).convert("RGB")
        frame = np.asarray(img)

        writer.append_data(frame)

print(f"Saved {output_mp4}")