53 lines
2.2 KiB
Python
53 lines
2.2 KiB
Python
import numpy as np
|
|
import matplotlib
|
|
|
|
|
|
def colorize_depth(depth: np.ndarray, mask: np.ndarray = None, normalize: bool = True, cmap: str = 'Spectral') -> np.ndarray:
|
|
if mask is None:
|
|
depth = np.where(depth > 0, depth, np.nan)
|
|
else:
|
|
depth = np.where((depth > 0) & mask, depth, np.nan)
|
|
disp = 1 / depth
|
|
if normalize:
|
|
min_disp, max_disp = np.nanquantile(disp, 0.001), np.nanquantile(disp, 0.999)
|
|
disp = (disp - min_disp) / (max_disp - min_disp)
|
|
#colored = np.nan_to_num(matplotlib.colormaps[cmap](1.0 - disp), 0)
|
|
#colored = (colored.clip(0, 1) * 255).astype(np.uint8)[:, :, :3]
|
|
grayscale = np.nan_to_num(disp, nan=0.0)
|
|
grayscale = (grayscale * 255).astype(np.uint8)
|
|
return grayscale
|
|
|
|
|
|
def colorize_depth_affine(depth: np.ndarray, mask: np.ndarray = None, cmap: str = 'Spectral') -> np.ndarray:
|
|
if mask is not None:
|
|
depth = np.where(mask, depth, np.nan)
|
|
|
|
min_depth, max_depth = np.nanquantile(depth, 0.001), np.nanquantile(depth, 0.999)
|
|
depth = (depth - min_depth) / (max_depth - min_depth)
|
|
colored = np.nan_to_num(matplotlib.colormaps[cmap](depth), 0)
|
|
colored = (colored.clip(0, 1) * 255).astype(np.uint8)[:, :, :3]
|
|
return colored
|
|
|
|
|
|
def colorize_disparity(disparity: np.ndarray, mask: np.ndarray = None, normalize: bool = True, cmap: str = 'Spectral') -> np.ndarray:
|
|
if mask is not None:
|
|
disparity = np.where(mask, disparity, np.nan)
|
|
|
|
if normalize:
|
|
min_disp, max_disp = np.nanquantile(disparity, 0.001), np.nanquantile(disparity, 0.999)
|
|
disparity = (disparity - min_disp) / (max_disp - min_disp)
|
|
colored = np.nan_to_num(matplotlib.colormaps[cmap](1.0 - disparity), 0)
|
|
colored = (colored.clip(0, 1) * 255).astype(np.uint8)[:, :, :3]
|
|
return colored
|
|
|
|
|
|
def colorize_segmentation(segmentation: np.ndarray, cmap: str = 'Set1') -> np.ndarray:
|
|
colored = matplotlib.colormaps[cmap]((segmentation % 20) / 20)
|
|
colored = (colored.clip(0, 1) * 255).astype(np.uint8)[:, :, :3]
|
|
return colored
|
|
|
|
|
|
def colorize_normal(normal: np.ndarray) -> np.ndarray:
|
|
normal = normal * [0.5, -0.5, -0.5] + 0.5
|
|
normal = (normal.clip(0, 1) * 255).astype(np.uint8)
|
|
return normal |