Spaces:
Running
on
Zero
Running
on
Zero
| import torch | |
| from matplotlib import cm | |
| def robust_min_max(tensor, quantile=0.99): | |
| T, H, W = tensor.shape | |
| min_vals = [] | |
| max_vals = [] | |
| for i in range(T): | |
| min_vals.append(torch.quantile(tensor[i], q=1-quantile, interpolation='nearest').item()) | |
| max_vals.append(torch.quantile(tensor[i], q=quantile, interpolation='nearest').item()) | |
| return min(min_vals), max(max_vals) | |
| class ColorMapper: | |
| def __init__(self, colormap: str = "inferno"): | |
| self.colormap = torch.tensor(cm.get_cmap(colormap).colors) | |
| def apply(self, image: torch.Tensor, v_min=None, v_max=None): | |
| # assert len(image.shape) == 2 | |
| if v_min is None: | |
| v_min = image.min() | |
| if v_max is None: | |
| v_max = image.max() | |
| image = (image - v_min) / (v_max - v_min) | |
| image = (image * 255).long() | |
| colormap = self.colormap.to(image.device) | |
| image = colormap[image] | |
| return image | |
| def color_video_disp(disp): | |
| visualizer = ColorMapper() | |
| disp_img = visualizer.apply(disp, v_min=0, v_max=1) | |
| return disp_img | |
| def pmap_to_disp(point_maps, valid_masks): | |
| disp_map = 1.0 / (point_maps[..., 2] + 1e-4) | |
| min_disparity, max_disparity = robust_min_max(disp_map) | |
| disp_map = torch.clamp((disp_map - min_disparity) / (max_disparity - min_disparity+1e-4), 0, 1) | |
| disp_map = color_video_disp(disp_map) | |
| disp_map[~valid_masks] = 0 | |
| return disp_map | |
| # imageio.mimsave(os.path.join(args.save_dir, os.path.basename(args.data[:-4])+'_disp.mp4'), disp, fps=24, quality=9, macro_block_size=1) | |