Save image tensor as numpy and optionally gif and png Args: imgs: (b, h, w, 3) save_gif: whether to save gif save_png: whether to save pngs of each images output_dir: output folder overwrite:
(
imgs: torch.Tensor,
save_npy: bool,
save_gif: bool,
save_png: bool,
output_dir: str,
overwrite: bool = False,
gif_fps: int = 10,
)
| 80 | |
| 81 | |
| 82 | def save_imgs( |
| 83 | imgs: torch.Tensor, |
| 84 | save_npy: bool, |
| 85 | save_gif: bool, |
| 86 | save_png: bool, |
| 87 | output_dir: str, |
| 88 | overwrite: bool = False, |
| 89 | gif_fps: int = 10, |
| 90 | ): |
| 91 | """ |
| 92 | Save image tensor as numpy and optionally gif and png |
| 93 | Args: |
| 94 | imgs: |
| 95 | (b, h, w, 3) |
| 96 | save_gif: |
| 97 | whether to save gif |
| 98 | save_png: |
| 99 | whether to save pngs of each images |
| 100 | output_dir: |
| 101 | output folder |
| 102 | overwrite: |
| 103 | whether to overwrite the old content |
| 104 | """ |
| 105 | if os.path.exists(output_dir) and not overwrite: |
| 106 | raise RuntimeError |
| 107 | os.makedirs(output_dir, exist_ok=True) |
| 108 | |
| 109 | imgs = imgs.detach().cpu().numpy() |
| 110 | |
| 111 | # save raw npy |
| 112 | if save_npy: |
| 113 | filename = os.path.join(output_dir, 'imgs.npy') |
| 114 | np.save(filename, imgs) |
| 115 | |
| 116 | # gif |
| 117 | if save_gif: |
| 118 | filename = os.path.join(output_dir, 'imgs.gif') |
| 119 | render.create_gif( |
| 120 | images=imgs, |
| 121 | filename=filename, |
| 122 | fps=gif_fps, |
| 123 | ) |
| 124 | |
| 125 | # pngs |
| 126 | if save_png: |
| 127 | subdir = os.path.join(output_dir, 'images') |
| 128 | os.makedirs(subdir, exist_ok=True) |
| 129 | for i in range(imgs.shape[0]): |
| 130 | filename = os.path.join(subdir, f'{i}.png') |
| 131 | imageio.imwrite(filename, (imgs[i] * 255.).astype(np.uint8)) |
| 132 | |
| 133 | |
| 134 | def compute_mse( |