diff --git a/src/robofish/io/file.py b/src/robofish/io/file.py index 66d3b5ae06bd8077656ff41e6994eefa35495c67..3b3d6d3758142226fa1c305b5094429e4342c1b4 100644 --- a/src/robofish/io/file.py +++ b/src/robofish/io/file.py @@ -1015,6 +1015,7 @@ class File(h5py.File): """ if video_path is not None: + video_path = Path(video_path) try: run(["ffmpeg"], capture_output=True) except Exception as e: @@ -1307,13 +1308,12 @@ class File(h5py.File): min_view = np.max((np.max(start_pose, axis=0) - np.min(start_pose, axis=0))[:2]) self.view_size = np.max([options["view_size"], min_view + options["margin"]]) - if video_path is not None: - pbar = tqdm(range(n_frames)) + pbar = tqdm(range(n_frames)) if video_path is not None else None def update(frame): output_list = [] - if "pbar" in locals().keys(): + if pbar: pbar.update(1) pbar.refresh() @@ -1433,6 +1433,14 @@ class File(h5py.File): + swarm_center ) + if video_path is not None: + if video_path.exists(): + if input(f"Video {str(video_path)} exists. Overwrite? (y/n)") == "y": + video_path.unlink() + else: + return + print(f"saving video to {video_path}") + print(f"Preparing to render n_frames: {n_frames}") ani = animation.FuncAnimation( @@ -1446,20 +1454,8 @@ class File(h5py.File): ) if video_path is not None: - # if i % (n / 40) == 0: - # print(f"Saving frame {i} of {n} ({100*i/n:.1f}%)") - - video_path = Path(video_path) - if video_path.exists(): - y = input(f"Video {str(video_path)} exists. Overwrite? (y/n)") - if y == "y": - video_path.unlink() - - if not video_path.exists(): - print(f"saving video to {video_path}") - - writervideo = animation.FFMpegWriter(fps=self.frequency) - ani.save(video_path, writer=writervideo, dpi=options["dpi"]) + writervideo = animation.FFMpegWriter(fps=self.frequency) + ani.save(video_path, writer=writervideo, dpi=options["dpi"]) plt.close() else: plt.show()