Spaces:
Paused
Paused
| """Colorization helpers: depth -> jet, normals -> RGB, canonical -> RGB.""" | |
| from __future__ import annotations | |
| import numpy as np | |
| try: | |
| import cv2 | |
| _HAS_CV2 = True | |
| except Exception: # pragma: no cover | |
| _HAS_CV2 = False | |
| def _robust_minmax(values: np.ndarray, lo: float = 2.0, hi: float = 98.0): | |
| finite = values[np.isfinite(values)] | |
| if finite.size == 0: | |
| return 0.0, 1.0 | |
| a, b = np.percentile(finite, [lo, hi]) | |
| if b <= a: | |
| b = a + 1e-6 | |
| return float(a), float(b) | |
| def depth_to_jet(depth: np.ndarray, valid: np.ndarray | None = None, | |
| vmin: float | None = None, vmax: float | None = None, | |
| bg=(255, 255, 255)) -> np.ndarray: | |
| """Colorize a (H, W) depth map with a JET colormap; invalid -> background.""" | |
| depth = np.asarray(depth, dtype=np.float32) | |
| if valid is None: | |
| valid = np.isfinite(depth) & (depth > 0) | |
| if vmin is None or vmax is None: | |
| vmin, vmax = _robust_minmax(depth[valid]) if valid.any() else (0.0, 1.0) | |
| norm = np.clip((depth - vmin) / (vmax - vmin), 0, 1) | |
| u8 = (norm * 255).astype(np.uint8) | |
| if _HAS_CV2: | |
| rgb = cv2.applyColorMap(u8, cv2.COLORMAP_JET)[:, :, ::-1] # BGR->RGB | |
| else: # simple fallback | |
| rgb = np.stack([u8, np.zeros_like(u8), 255 - u8], axis=-1) | |
| rgb = rgb.copy() | |
| rgb[~valid] = np.array(bg, dtype=np.uint8) | |
| return rgb | |
| def depth_to_jet_colors(depth_values: np.ndarray, vmin: float, vmax: float) -> np.ndarray: | |
| """JET colors (N, 3) uint8 for a flat array of depth values.""" | |
| depth_values = np.asarray(depth_values, dtype=np.float32) | |
| norm = np.clip((depth_values - vmin) / (vmax - vmin), 0, 1) | |
| u8 = (norm * 255).astype(np.uint8) | |
| if _HAS_CV2: | |
| lut = cv2.applyColorMap(np.arange(256, dtype=np.uint8)[:, None], | |
| cv2.COLORMAP_JET)[:, 0, ::-1] | |
| return lut[u8] | |
| return np.stack([u8, np.zeros_like(u8), 255 - u8], axis=-1) | |
| def normals_to_rgb(normals: np.ndarray) -> np.ndarray: | |
| """Standard normal-map RGB from OUTWARD camera-space normals (OpenCV frame). | |
| Displays in the convention used by Sapiens and other normal papers: | |
| facing camera = blue, +X right = red, +Y up = green. Input normals are the | |
| toward-camera (outward) normals from ``pointmap_to_normals``. | |
| """ | |
| n = np.asarray(normals, dtype=np.float32) | |
| disp = np.stack([n[..., 0], -n[..., 1], -n[..., 2]], axis=-1) | |
| rgb = (disp + 1.0) * 0.5 * 255.0 | |
| return np.clip(rgb, 0, 255).astype(np.uint8) | |
| def canonical_to_rgb(canonical: np.ndarray, valid: np.ndarray | None = None, | |
| lo=2.0, hi=98.0, ranges=None, bg=(255, 255, 255)): | |
| """Map canonical XYZ coordinates to RGB via per-axis percentile stretch. | |
| Returns ``(rgb, ranges)`` where ``ranges`` is the list of ``(min, max)`` per | |
| axis, so a consistent mapping can be reused across frames. | |
| """ | |
| canonical = np.asarray(canonical, dtype=np.float32) | |
| flat = canonical.reshape(-1, 3) | |
| if valid is not None: | |
| sel = flat[valid.reshape(-1)] | |
| else: | |
| sel = flat | |
| if ranges is None: | |
| ranges = [] | |
| for c in range(3): | |
| vals = sel[:, c] | |
| vals = vals[np.isfinite(vals)] | |
| if vals.size == 0: | |
| ranges.append((0.0, 1.0)) | |
| else: | |
| a, b = np.percentile(vals, [lo, hi]) | |
| if b <= a: | |
| b = a + 1e-6 | |
| ranges.append((float(a), float(b))) | |
| out = np.zeros_like(flat) | |
| for c in range(3): | |
| a, b = ranges[c] | |
| out[:, c] = np.clip((flat[:, c] - a) / (b - a), 0, 1) | |
| rgb = (out * 255).astype(np.uint8).reshape(canonical.shape) | |
| if valid is not None: | |
| rgb = rgb.copy() | |
| rgb[~valid] = np.array(bg, dtype=np.uint8) | |
| return rgb, ranges | |
| def canonical_colors(canonical_values: np.ndarray, ranges) -> np.ndarray: | |
| """RGB colors (N, 3) uint8 for flat canonical coords given fixed ranges.""" | |
| canonical_values = np.asarray(canonical_values, dtype=np.float32) | |
| out = np.zeros_like(canonical_values) | |
| for c in range(3): | |
| a, b = ranges[c] | |
| out[:, c] = np.clip((canonical_values[:, c] - a) / (b - a), 0, 1) | |
| return (out * 255).astype(np.uint8) | |
| def hsv_palette(n: int) -> np.ndarray: | |
| """A palette of ``n`` distinct bright RGB colors (uint8) via the HSV wheel.""" | |
| import colorsys | |
| cols = [] | |
| for i in range(n): | |
| r, g, b = colorsys.hsv_to_rgb(i / max(n, 1), 1.0, 1.0) | |
| cols.append([int(r * 255), int(g * 255), int(b * 255)]) | |
| return np.array(cols, dtype=np.uint8) | |