(video_batch: torch.Tensor, save_path: str, fps: int = 5, args=None, key=None, ind=None, foldername=None)
| 59 | |
| 60 | |
| 61 | def save_video_as_grid_and_mp4(video_batch: torch.Tensor, save_path: str, fps: int = 5, args=None, key=None, ind=None, foldername=None): |
| 62 | os.makedirs(save_path, exist_ok=True) |
| 63 | |
| 64 | for i, vid in enumerate(video_batch): |
| 65 | gif_frames = [] |
| 66 | for frame in vid: |
| 67 | frame = rearrange(frame, "c h w -> h w c") |
| 68 | |
| 69 | frame = (255.0 * frame).cpu().numpy().astype(np.uint8) |
| 70 | gif_frames.append(frame) |
| 71 | |
| 72 | file_name = f"{i:06d}" |
| 73 | |
| 74 | if ind is not None: |
| 75 | file_name = f"{ind}_{file_name}" |
| 76 | |
| 77 | if foldername is not None: |
| 78 | file_name = f"folder-{foldername}_{file_name}" |
| 79 | |
| 80 | if key is not None: |
| 81 | file_name = f"{key}_{file_name}" |
| 82 | |
| 83 | file_name = file_name + ".mp4" |
| 84 | |
| 85 | now_save_path = os.path.join(save_path, file_name) |
| 86 | with imageio.get_writer(now_save_path, fps=fps) as writer: |
| 87 | for frame in gif_frames: |
| 88 | writer.append_data(frame) |
| 89 | |
| 90 | def save_text(text: str, save_path: str): |
| 91 | txt_save_path = os.path.join(save_path, "text.txt") |
no test coverage detected