prs-eth/rollingdepth

1# Author: Bingxin Ke2# Last modified: 2024-11-253 4import concurrent.futures5from typing import Union6 7import matplotlib8import numpy as np9from tqdm import tqdm10 11 12def colorize_depth(13    depth: np.ndarray,14    min_depth: float,15    max_depth: float,16    cmap: str = "Spectral_r",17    valid_mask: Union[np.ndarray, None] = None,18) -> np.ndarray:19    assert len(depth.shape) >= 2, "Invalid dimension"20 21    if depth.ndim < 3:22        depth = depth[np.newaxis, :, :]23 24    # colorize25    cm = matplotlib.colormaps[cmap]26    depth = ((depth - min_depth) / (max_depth - min_depth)).clip(0, 1)27    img_colored_np = cm(depth, bytes=False)[:, :, :, 0:3]  # value from 0 to 128 29    if valid_mask is not None:30        valid_mask = valid_mask.squeeze()  # [H, W] or [B, H, W]31        if valid_mask.ndim < 3:32            valid_mask = valid_mask[np.newaxis, np.newaxis, :, :]33        else:34            valid_mask = valid_mask[:, np.newaxis, :, :]35        valid_mask = np.repeat(valid_mask, 3, axis=1)36        img_colored_np[~valid_mask] = 037 38    return img_colored_np39 40 41def colorize_depth_multi_thread(42    depth: np.ndarray,43    valid_mask: Union[np.ndarray, None] = None,44    chunk_size: int = 4,45    num_threads: int = 4,46    color_map: str = "Spectral",47    verbose: bool = False,48) -> np.ndarray:49    depth = depth.squeeze(1)50    assert 3 == depth.ndim51 52    n_frame = depth.shape[0]53 54    if valid_mask is None:55        valid_depth = depth56    else:57        valid_depth = depth[valid_mask]58    min_depth = valid_depth.min()59    max_depth = valid_depth.max()60 61    def process_chunk(chunk):62        chunk = colorize_depth(63            chunk, min_depth=min_depth, max_depth=max_depth, cmap=color_map64        )65        chunk = (chunk * 255).astype(np.uint8)66        return chunk67 68    # Pre-allocate the full array69    colored = np.empty((*depth.shape[:3], 3), dtype=np.uint8)70 71    with concurrent.futures.ThreadPoolExecutor(max_workers=num_threads) as executor:72        # Submit all tasks and store futures with their corresponding indices73        future_to_index = {74            executor.submit(process_chunk, depth[i : i + chunk_size]): i75            for i in range(0, n_frame, chunk_size)76        }77 78        # Process futures in the order they were submitted79        chunk_iterable = concurrent.futures.as_completed(future_to_index)80        if verbose:81            chunk_iterable = tqdm(82                chunk_iterable,83                desc=" colorizing",84                leave=False,85                total=len(future_to_index),86            )87        for future in chunk_iterable:88            index = future_to_index[future]89            start = index90            end = min(index + chunk_size, n_frame)91            result = future.result()92            colored[start:end] = result93    return colored