(depth_maps: torch.Tensor, path: str, conf_self: Optional[torch.Tensor] = None)
| 13 | |
| 14 | |
| 15 | def save_depth_maps(depth_maps: torch.Tensor, path: str, conf_self: Optional[torch.Tensor] = None): |
| 16 | min_depth = depth_maps.min() # float(torch.quantile(out, 0.01)) |
| 17 | max_depth = depth_maps.max() # float(torch.quantile(out, 0.99)) |
| 18 | |
| 19 | colored_depth = colorize_optimized( |
| 20 | depth_maps, |
| 21 | cmap_name="Spectral_r", |
| 22 | value_range=(min_depth, max_depth), |
| 23 | append_cbar=True, |
| 24 | ) |
| 25 | |
| 26 | if conf_self is not None: |
| 27 | if isinstance(conf_self, list): |
| 28 | if len(conf_self[0].shape) == 3: |
| 29 | conf_selfs = torch.cat(conf_self, dim=0) # (1, H, W) -> (N, H, W) |
| 30 | elif len(conf_self[0].shape) == 2: |
| 31 | conf_selfs = torch.stack(conf_self, dim=0) # (H, W) -> (N, H, W) |
| 32 | else: |
| 33 | conf_selfs = conf_self |
| 34 | log_conf = torch.log(conf_selfs) |
| 35 | min_conf = log_conf.min() # float(torch.quantile(out, 0.01)) |
| 36 | max_conf = log_conf.max() # float(torch.quantile(out, 0.99)) |
| 37 | colored_conf = colorize_optimized( |
| 38 | log_conf, |
| 39 | cmap_name="jet", |
| 40 | value_range=(min_conf, max_conf), |
| 41 | append_cbar=True, |
| 42 | ) |
| 43 | |
| 44 | img_paths = [f"{path}/frame_{i:04d}.png" for i in range(len(colored_depth))] |
| 45 | npy_paths = [f"{path}/frame_{i:04d}.npy" for i in range(len(depth_maps))] |
| 46 | |
| 47 | if conf_self is None: |
| 48 | to_save = (colored_depth * 255).detach().cpu().numpy().astype(np.uint8) |
| 49 | else: |
| 50 | to_save = torch.cat([colored_depth, colored_conf], dim=2) # 沿宽度方向连接 |
| 51 | to_save = (to_save * 255).detach().cpu().numpy().astype(np.uint8) |
| 52 | |
| 53 | for i, (img_path, npy_path, img_data) in enumerate(zip(img_paths, npy_paths, to_save)): |
| 54 | iio.imwrite(img_path, img_data) |
| 55 | np.save(npy_path, depth_maps[i].detach().cpu().numpy()) |
| 56 | |
| 57 | # comment this as it may fail sometimes |
| 58 | # images = [Image.open(img_path) for img_path in img_paths] |
| 59 | # images[0].save(f'{path}/_depth_maps.gif', save_all=True, append_images=images[1:], duration=100, loop=0) |
| 60 | |
| 61 | return depth_maps |
| 62 | |
| 63 | |
| 64 | def get_vertical_colorbar(h, vmin, vmax, cmap_name="jet", label=None, cbar_precision=2): |
no test coverage detected