utils/image_utils.py

"""
utils/image_utils.py
────────────────────
PIL ↔ numpy ↔ torch conversion helpers, plus depth-map visualisation
and PLY export utilities shared across pipeline stages.
"""

from __future__ import annotations

import io
import logging
from pathlib import Path
from typing import Union

import numpy as np
import torch
from PIL import Image

logger = logging.getLogger(__name__)

# ── Type aliases ──────────────────────────────────────────────────────────────
ArrayLike = Union[np.ndarray, torch.Tensor, Image.Image]


# ── Basic conversions ─────────────────────────────────────────────────────────

def to_pil(img: ArrayLike, mode: str = "RGB") -> Image.Image:
    """Convert any array-like to a PIL Image."""
    if isinstance(img, Image.Image):
        return img.convert(mode)
    if isinstance(img, torch.Tensor):
        img = img.detach().cpu().numpy()
    if img.dtype != np.uint8:
        # Assume [0,1] float → [0,255] uint8
        img = np.clip(img, 0.0, 1.0)
        img = (img * 255).astype(np.uint8)
    if img.ndim == 2:
        return Image.fromarray(img, mode="L").convert(mode)
    if img.ndim == 3 and img.shape[0] in (1, 3, 4):
        # CHW → HWC
        img = np.transpose(img, (1, 2, 0))
    return Image.fromarray(img).convert(mode)


def to_numpy(img: ArrayLike) -> np.ndarray:
    """Return HWC uint8 numpy array."""
    if isinstance(img, np.ndarray):
        return img
    if isinstance(img, Image.Image):
        return np.array(img)
    if isinstance(img, torch.Tensor):
        t = img.detach().cpu()
        if t.ndim == 4:
            t = t.squeeze(0)
        if t.shape[0] in (1, 3, 4):
            t = t.permute(1, 2, 0)
        arr = t.numpy()
        if arr.dtype != np.uint8:
            arr = np.clip(arr, 0.0, 1.0)
            arr = (arr * 255).astype(np.uint8)
        return arr
    raise TypeError(f"Unsupported type: {type(img)}")


def resize_image(img: Image.Image, max_side: int = 1024) -> Image.Image:
    """Resize image so longest side ≤ max_side, preserving aspect ratio."""
    w, h = img.size
    if max(w, h) <= max_side:
        return img
    scale = max_side / max(w, h)
    new_w, new_h = int(w * scale), int(h * scale)
    # Make dimensions divisible by 8 (required by most diffusion models)
    new_w = (new_w // 8) * 8
    new_h = (new_h // 8) * 8
    return img.resize((new_w, new_h), Image.LANCZOS)


# ── Depth map utilities ───────────────────────────────────────────────────────

def normalise_depth(depth: np.ndarray) -> np.ndarray:
    """Normalise a raw depth array to [0, 1] float32."""
    d = depth.astype(np.float32)
    dmin, dmax = d.min(), d.max()
    if dmax - dmin < 1e-8:
        return np.zeros_like(d)
    return (d - dmin) / (dmax - dmin)


def depth_to_colormap(depth: np.ndarray, colormap: str = "inferno") -> Image.Image:
    """Convert a float depth array to a false-colour PIL image for display."""
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.cm as cm

    norm = normalise_depth(depth)
    # get_cmap() was removed in Matplotlib 3.9; use matplotlib.colormaps instead
    try:
        cmap = matplotlib.colormaps[colormap]
    except (AttributeError, KeyError):
        cmap = cm.get_cmap(colormap)   # fallback for Matplotlib < 3.7
    coloured = (cmap(norm)[:, :, :3] * 255).astype(np.uint8)
    return Image.fromarray(coloured)


def depth_to_uint16(depth: np.ndarray) -> Image.Image:
    """
    Return a normalised depth map as an 8-bit greyscale PIL image.

    Gradio's gr.Image (type='pil') only renders mode 'L' or 'RGB' correctly —
    mode 'I' (32-bit int) displays as blank white, and mode 'I;16' crashes
    fromarray().  For the Gradio display panel we only need visual fidelity,
    so 8-bit 'L' is sufficient and always renders correctly.

    For a true lossless 16-bit PNG download, call save_depth_uint16() instead.
    """
    norm = normalise_depth(depth)
    uint8 = (norm * 255).astype(np.uint8)
    return Image.fromarray(uint8, mode="L")


# ── PLY / point cloud I/O ────────────────────────────────────────────────────

def rgbd_to_pointcloud(
    rgb: np.ndarray,
    depth: np.ndarray,
    fx: float = 500.0,
    fy: float = 500.0,
    cx: float | None = None,
    cy: float | None = None,
    depth_scale: float = 1.0,
    max_depth: float = 10.0,
    fg_mask: np.ndarray | None = None,
) -> tuple[np.ndarray, np.ndarray]:
    """
    Back-project an RGBD pair into a coloured point cloud.

    Parameters
    ----------
    rgb          : HxWx3 uint8
    depth        : HxW float32  (normalised [0,1])
    fx, fy       : focal lengths in pixels
    cx, cy       : principal point  (defaults to image centre)
    depth_scale  : multiply normalised depth by this to get pseudo-metres
    max_depth    : clip depth values beyond this distance
    fg_mask      : optional HxW float32 in [0,1] (or bool) matte from the
                   background-removal stage. When supplied, this is the
                   foreground/background split (thresholded at 0.5) instead
                   of the near-black brightness heuristic — the matte from a
                   real segmentation model is far more reliable than
                   "is this pixel dark".

    Returns
    -------
    points : Nx3 float32   (X,Y,Z in camera space)
    colors : Nx3 uint8     (R,G,B)
    """
    h, w = depth.shape[:2]
    cx = cx if cx is not None else w / 2.0
    cy = cy if cy is not None else h / 2.0

    d = depth.astype(np.float32)

    # ── Outlier removal before back-projection ────────────────────────────────
    # Monocular depth maps are noisy at sky/background regions. Clip to the
    # 2nd-95th percentile so extreme depth values don't produce streaking fans.
    # Sky pixels in MiDaS typically map to the highest normalised depth values.
    d_lo, d_hi = np.percentile(d, 2), np.percentile(d, 95)
    d = np.clip(d, d_lo, d_hi)
    d_range = d_hi - d_lo
    if d_range > 1e-6:
        d = (d - d_lo) / d_range   # re-normalise to [0,1] after clipping

    d = d * depth_scale

    # Build pixel grids
    u = np.arange(w, dtype=np.float32)
    v = np.arange(h, dtype=np.float32)
    uu, vv = np.meshgrid(u, v)

    # Back-project
    x = (uu - cx) * d / fx
    y = (vv - cy) * d / fy
    z = d

    # Foreground/background split. Prefer the segmentation matte from the
    # background-removal stage (accurate, subject-aware); fall back to the
    # brightness heuristic only when no matte is available (e.g. background
    # removal was disabled for this run).
    if fg_mask is not None:
        fg = fg_mask
        if fg.shape[:2] != (h, w):
            fg_pil = Image.fromarray(fg.astype(np.float32), mode="F").resize((w, h), Image.BILINEAR)
            fg = np.array(fg_pil, dtype=np.float32)
        keep_mask = fg > 0.5
        if not keep_mask.any():
            # The matte found no foreground at all (blank/ambiguous image, or a
            # segmentation miss). Silently returning an empty point cloud would
            # be a worse outcome than just not masking — fall back to keeping
            # everything so the reconstruction stage still has something to work with.
            logger.warning("Foreground matte is empty (no pixels above threshold) — skipping matte-based masking for this frame.")
            keep_mask = np.ones((h, w), dtype=bool)
    else:
        brightness = rgb.mean(axis=2)
        keep_mask = brightness >= 8  # reject near-black — typically featureless sky/background

    mask = (z > 0.01) & (z < max_depth) & keep_mask
    points = np.stack([x[mask], y[mask], z[mask]], axis=-1)
    colors = rgb[mask]

    return points.astype(np.float32), colors.astype(np.uint8)


def save_ply(path: Union[str, Path], points: np.ndarray, colors: np.ndarray) -> None:
    """Write a coloured point cloud to a binary PLY file."""
    from plyfile import PlyData, PlyElement

    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)

    n = len(points)
    vertex = np.zeros(
        n,
        dtype=[
            ("x", "f4"), ("y", "f4"), ("z", "f4"),
            ("red", "u1"), ("green", "u1"), ("blue", "u1"),
        ],
    )
    vertex["x"] = points[:, 0]
    vertex["y"] = points[:, 1]
    vertex["z"] = points[:, 2]
    vertex["red"]   = colors[:, 0]
    vertex["green"] = colors[:, 1]
    vertex["blue"]  = colors[:, 2]

    el = PlyElement.describe(vertex, "vertex")
    PlyData([el], byte_order="<").write(str(path))
    logger.info("Saved PLY  →  %s  (%d points)", path, n)


def load_ply(path: Union[str, Path]) -> tuple[np.ndarray, np.ndarray]:
    """Load a coloured PLY and return (points Nx3, colors Nx3)."""
    from plyfile import PlyData

    plydata = PlyData.read(str(path))
    v = plydata["vertex"]
    points = np.stack([v["x"], v["y"], v["z"]], axis=-1).astype(np.float32)
    colors = np.stack([v["red"], v["green"], v["blue"]], axis=-1).astype(np.uint8)
    return points, colors


# ── File helpers ──────────────────────────────────────────────────────────────

def pil_to_bytes(img: Image.Image, fmt: str = "PNG") -> bytes:
    buf = io.BytesIO()
    img.save(buf, format=fmt)
    return buf.getvalue()


def bytes_to_pil(data: bytes) -> Image.Image:
    return Image.open(io.BytesIO(data))