Spaces:
Running
on
Zero
Running
on
Zero
File size: 2,930 Bytes
a51c6d2 |
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 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 |
from typing import *
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.99)
disp = (disp - min_disp) / (max_disp - min_disp)
colored = np.nan_to_num(matplotlib.colormaps[cmap](1.0 - disp)[..., :3], 0)
colored = np.ascontiguousarray((colored.clip(0, 1) * 255).astype(np.uint8))
return colored
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)[..., :3], 0)
colored = np.ascontiguousarray((colored.clip(0, 1) * 255).astype(np.uint8))
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)[..., :3], 0)
colored = np.ascontiguousarray((colored.clip(0, 1) * 255).astype(np.uint8))
return colored
def colorize_segmentation(segmentation: np.ndarray, cmap: str = 'Set1') -> np.ndarray:
colored = matplotlib.colormaps[cmap]((segmentation % 20) / 20)[..., :3]
colored = np.ascontiguousarray((colored.clip(0, 1) * 255).astype(np.uint8))
return colored
def colorize_normal(normal: np.ndarray, mask: np.ndarray = None) -> np.ndarray:
if mask is not None:
normal = np.where(mask[..., None], normal, 0)
normal = normal * [0.5, -0.5, -0.5] + 0.5
normal = (normal.clip(0, 1) * 255).astype(np.uint8)
return normal
def colorize_error_map(error_map: np.ndarray, mask: np.ndarray = None, cmap: str = 'plasma', value_range: Tuple[float, float] = None):
vmin, vmax = value_range if value_range is not None else (np.nanmin(error_map), np.nanmax(error_map))
cmap = matplotlib.colormaps[cmap]
colorized_error_map = cmap(((error_map - vmin) / (vmax - vmin)).clip(0, 1))[..., :3]
if mask is not None:
colorized_error_map = np.where(mask[..., None], colorized_error_map, 0)
colorized_error_map = np.ascontiguousarray((colorized_error_map.clip(0, 1) * 255).astype(np.uint8))
return colorized_error_map
|