Update generate.py

This commit is contained in:
Norbert Schmidt
2026-01-11 13:11:42 +01:00
parent bbe43a4e7d
commit 1366941d39

View File

@@ -26,7 +26,8 @@ def main():
parser.add_argument("--frames", type=int, default=25, help="frame count") 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("--fps", type=int, default=24, help="output fps")
parser.add_argument("--seed", type=int, default=None, help="random seed") 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() args = parser.parse_args()
@@ -88,18 +89,31 @@ def main():
frame = np.array(frame, dtype=np.uint8) frame = np.array(frame, dtype=np.uint8)
frames.append(frame) frames.append(frame)
# export with high quality h264 # export video
writer = imageio.get_writer( if args.prores:
args.output, output_path = args.output.replace('.mp4', '.mov') if args.output.endswith('.mp4') else args.output
fps=args.fps, writer = imageio.get_writer(
codec='libx264', output_path,
quality=None, fps=args.fps,
pixelformat='yuv420p', codec='prores_ks',
output_params=['-crf', str(args.crf), '-preset', 'slow'] 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: for frame in frames:
writer.append_data(frame) writer.append_data(frame)
writer.close() writer.close()
args.output = output_path
print(f"\nsaved to: {args.output}") print(f"\nsaved to: {args.output}")
print(f"seed: {args.seed}") print(f"seed: {args.seed}")