(tensor,
save_file=None,
fps=30,
suffix='.mp4',
nrow=8,
normalize=True,
value_range=(-1, 1),
retry=5)
| 21 | |
| 22 | |
| 23 | def cache_video(tensor, |
| 24 | save_file=None, |
| 25 | fps=30, |
| 26 | suffix='.mp4', |
| 27 | nrow=8, |
| 28 | normalize=True, |
| 29 | value_range=(-1, 1), |
| 30 | retry=5): |
| 31 | # cache file |
| 32 | cache_file = osp.join('/tmp', rand_name( |
| 33 | suffix=suffix)) if save_file is None else save_file |
| 34 | |
| 35 | # save to cache |
| 36 | error = None |
| 37 | for _ in range(retry): |
| 38 | try: |
| 39 | # preprocess |
| 40 | tensor = tensor.clamp(min(value_range), max(value_range)) |
| 41 | tensor = torch.stack([ |
| 42 | torchvision.utils.make_grid( |
| 43 | u, nrow=nrow, normalize=normalize, value_range=value_range) |
| 44 | for u in tensor.unbind(2) |
| 45 | ], |
| 46 | dim=1).permute(1, 2, 3, 0) |
| 47 | tensor = (tensor * 255).type(torch.uint8).cpu() |
| 48 | |
| 49 | # write video |
| 50 | writer = imageio.get_writer( |
| 51 | cache_file, fps=fps, codec='libx264', quality=8) |
| 52 | for frame in tensor.numpy(): |
| 53 | writer.append_data(frame) |
| 54 | writer.close() |
| 55 | return cache_file |
| 56 | except Exception as e: |
| 57 | error = e |
| 58 | continue |
| 59 | else: |
| 60 | print(f'cache_video failed, error: {error}', flush=True) |
| 61 | return None |
| 62 | |
| 63 | |
| 64 | def cache_image(tensor, |
nothing calls this directly
no test coverage detected