Replace imageio with OpenCV for video saving in generate.py; updated default frame count to 100.

This commit is contained in:
Prince Canuma
2026-01-12 16:12:41 +01:00
parent 666e1f2e0c
commit 7eac6ae7de

View File

@@ -276,8 +276,13 @@ def generate_video(
output_path.parent.mkdir(parents=True, exist_ok=True) output_path.parent.mkdir(parents=True, exist_ok=True)
try: try:
import imageio import cv2
imageio.mimwrite(str(output_path), video_np, fps=fps, codec='libx264') height, width = video_np.shape[1], video_np.shape[2]
fourcc = cv2.VideoWriter_fourcc(*'avc1')
out = cv2.VideoWriter(str(output_path), fourcc, fps, (width, height))
for frame in video_np:
out.write(cv2.cvtColor(frame, cv2.COLOR_RGB2BGR))
out.release()
print(f"Saved video to {output_path}") print(f"Saved video to {output_path}")
except Exception as e: except Exception as e:
print(f"Could not save video: {e}") print(f"Could not save video: {e}")
@@ -328,8 +333,8 @@ Examples:
parser.add_argument( parser.add_argument(
"--num-frames", "-n", "--num-frames", "-n",
type=int, type=int,
default=33, default=100,
help="Number of frames (default: 33, must be 1 + 8*k)" help="Number of frames (default: 100)"
) )
parser.add_argument( parser.add_argument(
"--seed", "-s", "--seed", "-s",