Spaces:
Runtime error
Runtime error
| # TRI-VIDAR - Copyright 2022 Toyota Research Institute. All rights reserved. | |
| import flow_vis | |
| import numpy as np | |
| import torch | |
| from matplotlib.cm import get_cmap | |
| from vidar.utils.decorators import iterate1 | |
| from vidar.utils.depth import depth2inv | |
| from vidar.utils.types import is_tensor, is_list | |
| def flow_to_color(flow_uv, clip_flow=None): | |
| """ | |
| Calculate color from optical flow | |
| Parameters | |
| ---------- | |
| flow_uv : np.Array | |
| Optical flow [H,W,2] | |
| clip_flow : Float | |
| Clipping value for optical flow | |
| Returns | |
| ------- | |
| colors : np.array | |
| Optical flow colormap [H,W,3] | |
| """ | |
| # Clip if requested | |
| if clip_flow is not None: | |
| flow_uv = np.clip(flow_uv, -clip_flow, clip_flow) | |
| # Get optical flow channels | |
| u = flow_uv[:, :, 0] | |
| v = flow_uv[:, :, 1] | |
| # Calculate maximum radian | |
| rad_max = np.sqrt(2) * clip_flow if clip_flow is not None else \ | |
| np.max(np.sqrt(np.square(u) + np.square(v))) | |
| # Normalize optical flow channels | |
| epsilon = 1e-5 | |
| u = u / (rad_max + epsilon) | |
| v = v / (rad_max + epsilon) | |
| # Return colormap [0,1] | |
| return flow_vis.flow_uv_to_colors(u, v, convert_to_bgr=False) / 255 | |
| def viz_inv_depth(inv_depth, normalizer=None, percentile=95, | |
| colormap='plasma', filter_zeros=False): | |
| """ | |
| Converts an inverse depth map to a colormap for visualization. | |
| Parameters | |
| ---------- | |
| inv_depth : torch.Tensor | |
| Inverse depth map to be converted [B,1,H,W] | |
| normalizer : Float | |
| Value for inverse depth map normalization | |
| percentile : Float | |
| Percentile value for automatic normalization | |
| colormap : String | |
| Colormap to be used | |
| filter_zeros : Bool | |
| If True, do not consider zero values during normalization | |
| Returns | |
| ------- | |
| colormap : np.Array [H,W,3] | |
| Colormap generated from the inverse depth map | |
| """ | |
| if is_list(inv_depth): | |
| return [viz_inv_depth( | |
| inv[0], normalizer, percentile, colormap, filter_zeros) | |
| for inv in inv_depth] | |
| # If a tensor is provided, convert to numpy | |
| if is_tensor(inv_depth): | |
| # If it has a batch size, use first one | |
| if len(inv_depth.shape) == 4: | |
| inv_depth = inv_depth[0] | |
| # Squeeze if depth channel exists | |
| if len(inv_depth.shape) == 3: | |
| inv_depth = inv_depth.squeeze(0) | |
| inv_depth = inv_depth.detach().cpu().numpy() | |
| cm = get_cmap(colormap) | |
| if normalizer is None: | |
| if (inv_depth > 0).sum() == 0: | |
| normalizer = 1.0 | |
| else: | |
| normalizer = np.percentile( | |
| inv_depth[inv_depth > 0] if filter_zeros else inv_depth, percentile) | |
| inv_depth = inv_depth / (normalizer + 1e-6) | |
| colormap = cm(np.clip(inv_depth, 0., 1.0))[:, :, :3] | |
| colormap[inv_depth == 0] = 0 | |
| return colormap | |
| def viz_depth(depth, *args, **kwargs): | |
| """Same as viz_inv_depth, but takes depth as input instead""" | |
| return viz_inv_depth(depth2inv(depth), *args, **kwargs) | |
| def viz_normals(normals): | |
| """ | |
| Converts normals map to a colormap for visualization. | |
| Parameters | |
| ---------- | |
| normals : torch.Tensor | |
| Inverse depth map to be converted [B,3,H,W] | |
| Returns | |
| ------- | |
| colormap : np.Array | |
| Colormap generated from the normals map [H,W,3] | |
| """ | |
| # If a tensor is provided, convert to numpy | |
| if is_tensor(normals): | |
| normals = normals.permute(1, 2, 0).detach().cpu().numpy() | |
| return (normals + 1) / 2 | |
| def viz_optical_flow(optflow, clip_value=100.): | |
| """ | |
| Returns a colorized version of an optical flow map | |
| Parameters | |
| ---------- | |
| optflow : torch.Tensor | |
| Optical flow to be colorized (NOT in batch) [2,H,W] | |
| clip_value : Float | |
| Optical flow clip value for visualization | |
| Returns | |
| ------- | |
| colorized : np.Array | |
| Colorized version of the input optical flow [H,W,3] | |
| """ | |
| # If a tensor is provided, convert to numpy | |
| if is_list(optflow): | |
| return [viz_optical_flow(opt[0]) for opt in optflow] | |
| if is_tensor(optflow): | |
| if len(optflow.shape) == 4: | |
| optflow = optflow[0] | |
| optflow = optflow.permute(1, 2, 0).detach().cpu().numpy() | |
| # Return colorized optical flow | |
| return flow_to_color(optflow, clip_flow=clip_value) | |
| def viz_photo(photo, colormap='viridis', normalize=False): | |
| """ | |
| Returns a colorized version of the photometric loss | |
| Parameters | |
| ---------- | |
| photo : torch.Tensor | |
| Per-pixel photometric error | |
| colormap : String | |
| Which colormap to use | |
| normalize : Bool | |
| Whether the photometric error should be normalized between [0,1] | |
| Returns | |
| ------- | |
| colorized : np.Array | |
| Colorized version of the photometric error [H,W,3] | |
| """ | |
| if is_tensor(photo): | |
| if len(photo.shape) == 4: | |
| photo = photo[0] | |
| if len(photo.shape) == 3: | |
| photo = photo.squeeze(0) | |
| photo = photo.detach().cpu().numpy() | |
| cm = get_cmap(colormap) | |
| if normalize: | |
| photo -= photo.min() | |
| photo /= photo.max() | |
| colormap = cm(np.clip(photo, 0., 1.0))[:, :, :3] | |
| colormap[photo == 0] = 0 | |
| return colormap | |
| def viz_semantic(semantic, ontology): | |
| """ | |
| Returns a colorized version of a semantic map | |
| Parameters | |
| ---------- | |
| semantic : torch.Tensor | |
| Semantic map to be colorized [B,1,H,W] | |
| ontology : Dict | |
| Dictionary mapping between class and color | |
| Returns | |
| ------- | |
| colorized : np.Array | |
| Colorized version of the semantic map [H,W,3] | |
| """ | |
| # If it is a tensor, cast to numpy | |
| if is_tensor(semantic): | |
| if semantic.dim() == 3: | |
| semantic = semantic.squeeze(0) | |
| semantic = semantic.detach().cpu().numpy() | |
| # Create and populate color map | |
| color = np.zeros((semantic.shape[0], semantic.shape[1], 3)) | |
| for key in ontology.keys(): | |
| key_color = np.array(ontology[key]['color']) | |
| if is_tensor(key_color): | |
| key_color = key_color.detach().cpu().numpy() | |
| color[semantic == int(key)] = key_color / 255. | |
| # Return colored semantic map | |
| return color | |
| def viz_camera(camera): | |
| """ | |
| Returns a colorized version of a camera viewing rays | |
| Parameters | |
| ---------- | |
| camera : Camera of torch.Tensor | |
| Input camera or viewing rays | |
| Returns | |
| ------- | |
| colorized : np.Array | |
| Colorized version of the camera viewing rays [H,W,3] | |
| """ | |
| if is_tensor(camera): | |
| # If it's a tensor, reshape it | |
| rays = camera[-3:].permute(1, 2, 0).detach().cpu().numpy() | |
| else: | |
| # If it's a camera, get viewing rays | |
| rays = camera.no_translation().get_viewdirs(normalize=True, flatten=False, to_world=True) | |
| rays = rays[0].permute(1, 2, 0).detach().cpu().numpy() | |
| return (rays + 1) / 2 | |