diff --git a/generate.py b/generate.py index 96f3887..6638f4b 100644 --- a/generate.py +++ b/generate.py @@ -26,7 +26,8 @@ def main(): parser.add_argument("--frames", type=int, default=25, help="frame count") parser.add_argument("--fps", type=int, default=24, help="output fps") parser.add_argument("--seed", type=int, default=None, help="random seed") - parser.add_argument("--crf", type=int, default=18, help="video quality (0-51, lower=better, 18=visually lossless)") + parser.add_argument("--crf", type=int, default=10, help="video quality (0-51, lower=better)") + parser.add_argument("--prores", action="store_true", help="use prores codec (large files, best quality)") args = parser.parse_args() @@ -88,18 +89,31 @@ def main(): frame = np.array(frame, dtype=np.uint8) frames.append(frame) - # export with high quality h264 - writer = imageio.get_writer( - args.output, - fps=args.fps, - codec='libx264', - quality=None, - pixelformat='yuv420p', - output_params=['-crf', str(args.crf), '-preset', 'slow'] - ) + # export video + if args.prores: + output_path = args.output.replace('.mp4', '.mov') if args.output.endswith('.mp4') else args.output + writer = imageio.get_writer( + output_path, + fps=args.fps, + codec='prores_ks', + pixelformat='yuv422p10le', + output_params=['-profile:v', '3'] # prores hq + ) + else: + output_path = args.output + writer = imageio.get_writer( + output_path, + fps=args.fps, + codec='libx264', + quality=None, + pixelformat='yuv420p', + output_params=['-crf', str(args.crf), '-preset', 'slow'] + ) + for frame in frames: writer.append_data(frame) writer.close() + args.output = output_path print(f"\nsaved to: {args.output}") print(f"seed: {args.seed}")