src/utils/utils.py

"""Preview helpers consumed by ``app.py`` to turn raw pager outputs into
uint8 images for the Gradio UI.
"""

from __future__ import annotations

from typing import Tuple

import numpy as np
from scipy.ndimage import median_filter


def prepare_image_for_logging(image: np.ndarray, apply_median: bool = False) -> np.ndarray:
    """Per-sample min/max stretch to ``uint8``, optionally median-filtered."""
    if apply_median:
        image = median_filter(image.astype(np.float32), size=3)
    image = (image - image.min()) / (image.max() - image.min() + 1e-8)
    return (image * 255).astype("uint8")


def prepare_depth_for_logging(
    pager,
    depth,
    sky_mask,
    orig_size: Tuple[int, int],
    cmap,
    log_scale=None,
    sky_mask_threshold: float = 0.3,
    sky_mask_softness: float = 0.15,
    sky_mask_open_kernel: int = 0,
):
    """Return ``(raw_depth, viz_depth)``.

    ``raw_depth`` is a float32 metric-depth ERP array suitable for saving to
    ``.npz`` and feeding the evaluation pipeline. ``viz_depth`` is a uint8
    Spectral-coloured ERP image for previews.
    """
    raw_image, log_image = pager.process_depth_output(
        depth,
        orig_size=orig_size,
        sky_mask=sky_mask,
        log_scale=log_scale,
        sky_mask_threshold=sky_mask_threshold,
        sky_mask_softness=sky_mask_softness,
        sky_mask_open_kernel=sky_mask_open_kernel,
    )
    raw_image = raw_image.detach().cpu().numpy()
    log_image = log_image.detach().cpu().numpy()
    viz = prepare_image_for_logging(log_image, apply_median=True)
    viz = cmap(viz / 255.0)
    viz = (viz[..., :3] * 255)[0].astype(np.uint8)
    return raw_image, np.transpose(viz, (2, 0, 1))


def prepare_normals_for_logging(
    pager,
    normals,
    sky_mask,
    orig_size: Tuple[int, int],
    cmap=None,
    sky_mask_threshold: float = 0.3,
    sky_mask_softness: float = 0.15,
    sky_mask_open_kernel: int = 0,
):
    """Return ``(raw_normals, viz_normals)``.

    ``raw_normals`` is a float32 unit-normal ERP tensor in ``[-1, 1]``,
    shape ``(3, H, W)`` — what the evaluation script consumes directly.
    ``viz_normals`` is a uint8 RGB image (per-sample min/max rescaled).
    """
    raw_normals = pager.process_normals_output(
        normals,
        orig_size=orig_size,
        sky_mask=sky_mask,
        sky_mask_threshold=sky_mask_threshold,
        sky_mask_softness=sky_mask_softness,
        sky_mask_open_kernel=sky_mask_open_kernel,
    ).detach().cpu().numpy()
    viz = prepare_image_for_logging(raw_normals)
    return raw_normals, viz