Spaces:
Sleeping
Sleeping
File size: 1,320 Bytes
78d2329 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 | import torch
import torch.utils.data
import numpy as np
import torchvision.utils as vutils
import cv2
from matplotlib.cm import get_cmap
import matplotlib as mpl
import matplotlib.cm as cm
# https://github.com/autonomousvision/unimatch/blob/master/utils/visualization.py
def vis_disparity(disp):
disp_vis = (disp - disp.min()) / (disp.max() - disp.min()) * 255.0
disp_vis = disp_vis.astype("uint8")
disp_vis = cv2.applyColorMap(disp_vis, cv2.COLORMAP_INFERNO)
return disp_vis
def viz_depth_tensor(disp, return_numpy=False, colormap="plasma", as_uint8=True, vmin=None, vmax=None):
# visualize inverse depth
assert isinstance(disp, torch.Tensor)
device = disp.device
disp = disp.cpu().numpy()
if vmin is None:
vmin = disp.min()
if vmax is None:
vmax = disp.max()
normalizer = mpl.colors.Normalize(vmin=vmin, vmax=vmax)
mapper = cm.ScalarMappable(norm=normalizer, cmap=colormap)
colormapped_im = mapper.to_rgba(disp)[:, :, :3]
if as_uint8:
colormapped_im = (colormapped_im * 255).astype(np.uint8)
else:
colormapped_im = (colormapped_im).astype(np.float32)
if return_numpy:
return colormapped_im
viz = torch.from_numpy(colormapped_im).permute(2, 0, 1) # [3, H, W]
viz = viz.to(device)
return viz
|