Utils.py

"""Minimal Utils module for Fast-FoundationStereo inference."""
import numpy as np
import torch

AMP_DTYPE = torch.float16


def set_logging_format(level=None):
    import logging
    FORMAT = '%(message)s'
    logging.basicConfig(level=logging.INFO, format=FORMAT, datefmt='%m-%d|%H:%M:%S')


def set_seed(random_seed):
    import random
    np.random.seed(random_seed)
    random.seed(random_seed)
    torch.manual_seed(random_seed)
    torch.cuda.manual_seed_all(random_seed)


def vis_disparity(disp, min_val=None, max_val=None, invalid_thres=np.inf, color_map=None, cmap=None, other_output=None):
    """Visualize disparity as a color-mapped image.

    Args:
        disp: (H, W) numpy array of disparity values.
        min_val: optional minimum value for normalization.
        max_val: optional maximum value for normalization.
        invalid_thres: disparities >= this are invalid.
        color_map: OpenCV colormap constant.
        cmap: optional custom colormap function.
        other_output: optional dict to store min/max values.

    Returns:
        (H, W, 3) uint8 RGB image.
    """
    import cv2
    disp = disp.copy()
    H, W = disp.shape[:2]
    invalid_mask = disp >= invalid_thres
    if (invalid_mask == 0).sum() == 0:
        return np.zeros((H, W, 3), dtype=np.uint8)
    if min_val is None:
        min_val = disp[invalid_mask == 0].min()
    if max_val is None:
        max_val = disp[invalid_mask == 0].max()
    if other_output is not None:
        other_output['min_val'] = min_val
        other_output['max_val'] = max_val
    vis = ((disp - min_val) / (max_val - min_val)).clip(0, 1) * 255
    if cmap is None:
        vis = cv2.applyColorMap(vis.clip(0, 255).astype(np.uint8),
                                color_map if color_map is not None else cv2.COLORMAP_TURBO)[..., ::-1]
    else:
        vis = cmap(vis.astype(np.uint8))[..., :3] * 255
    if invalid_mask.any():
        vis[invalid_mask] = 0
    return vis.astype(np.uint8)