app/visualize.py

"""PNG renderers for the Kakera response payload.

All public renderers take NumPy arrays and return raw PNG bytes. The frontend
receives base64-encoded strings (see ``to_b64``) ready to drop into ``<img
src>``.

Color philosophy follows ``specs/03_FRONTEND.md``: paper-white background,
aizome indigo for primary signal, vermilion as the sparing accent (saddle
channel, instance contours), near-black for grayscale endpoints.

Implementation uses only NumPy + OpenCV. No matplotlib / Pillow on purpose
(keeps the Docker image lean).
"""

from __future__ import annotations

import base64

import cv2
import numpy as np


# Project palette - RGB tuples. Converted to BGR inside ``_encode_rgb``.
PAPER_WHITE: tuple[int, int, int] = (245, 240, 232)
INDIGO:      tuple[int, int, int] = (38, 52, 110)
VERMILION:   tuple[int, int, int] = (213, 74, 56)
NEAR_BLACK:  tuple[int, int, int] = (23, 25, 31)

# Pure-color tints for the R / G / B channel renderers.
_PURE_R: tuple[int, int, int] = (255, 0, 0)
_PURE_G: tuple[int, int, int] = (0, 255, 0)
_PURE_B: tuple[int, int, int] = (0, 0, 255)


# ---------------------------------------------------------------------------
# Low-level helpers
# ---------------------------------------------------------------------------

def _encode_rgb(rgb_u8: np.ndarray) -> bytes:
    """Encode an HWC RGB uint8 array as PNG bytes."""
    bgr = cv2.cvtColor(rgb_u8, cv2.COLOR_RGB2BGR)
    ok, buf = cv2.imencode(".png", bgr)
    if not ok:
        raise RuntimeError("cv2.imencode failed for RGB image.")
    return buf.tobytes()


def _encode_rgba(rgba_u8: np.ndarray) -> bytes:
    """Encode an HWC RGBA uint8 array as PNG bytes (preserves alpha)."""
    bgra = cv2.cvtColor(rgba_u8, cv2.COLOR_RGBA2BGRA)
    ok, buf = cv2.imencode(".png", bgra)
    if not ok:
        raise RuntimeError("cv2.imencode failed for RGBA image.")
    return buf.tobytes()


def _clip01(arr: np.ndarray) -> np.ndarray:
    """Clamp to [0, 1] as float32 (no copy when already inside)."""
    a = arr.astype(np.float32, copy=False)
    return np.clip(a, 0.0, 1.0)


def _ramp(value_2d: np.ndarray, lo_rgb: tuple[int, int, int],
          hi_rgb: tuple[int, int, int]) -> np.ndarray:
    """Linear blend between two RGB endpoints driven by ``value_2d in [0,1]``.

    Returns an (H, W, 3) uint8 RGB image.
    """
    v = _clip01(value_2d)[..., None]
    lo = np.array(lo_rgb, dtype=np.float32)
    hi = np.array(hi_rgb, dtype=np.float32)
    out = lo + v * (hi - lo)
    return np.clip(out, 0.0, 255.0).astype(np.uint8)


def _diverging(value_2d: np.ndarray,
               lo_rgb: tuple[int, int, int],
               mid_rgb: tuple[int, int, int],
               hi_rgb: tuple[int, int, int]) -> np.ndarray:
    """Three-stop diverging ramp pivoted at value 0.5."""
    v = _clip01(value_2d)
    low_mask = v < 0.5
    t = np.where(low_mask, v * 2.0, (v - 0.5) * 2.0)[..., None]
    lo  = np.array(lo_rgb,  dtype=np.float32)
    mid = np.array(mid_rgb, dtype=np.float32)
    hi  = np.array(hi_rgb,  dtype=np.float32)
    low_part  = lo  + t * (mid - lo)
    high_part = mid + t * (hi  - mid)
    out = np.where(low_mask[..., None], low_part, high_part)
    return np.clip(out, 0.0, 255.0).astype(np.uint8)


def _tinted_grayscale(value_2d: np.ndarray,
                      pure_rgb: tuple[int, int, int]) -> np.ndarray:
    """Multiply a [0,1] map by a pure RGB color; 0 -> black, 1 -> pure color."""
    v = _clip01(value_2d)[..., None]
    pure = np.array(pure_rgb, dtype=np.float32)
    return np.clip(v * pure, 0.0, 255.0).astype(np.uint8)


# ---------------------------------------------------------------------------
# Public renderers
# ---------------------------------------------------------------------------

def render_rgb(rgb_u8: np.ndarray) -> bytes:
    """Encode an HWC uint8 RGB array as PNG bytes."""
    if rgb_u8.dtype != np.uint8 or rgb_u8.ndim != 3 or rgb_u8.shape[2] != 3:
        raise ValueError(
            f"render_rgb expects (H, W, 3) uint8; got shape={rgb_u8.shape} dtype={rgb_u8.dtype}."
        )
    return _encode_rgb(rgb_u8)


def render_channel(channel_2d: np.ndarray, *, idx: int) -> bytes:
    """Render one of the 9 input-tensor channels as a colored PNG.

    See ``CHANNEL_NAMES`` in ``app/inference.py`` for the index-to-name map.
    Coloring rules mirror ``specs/02_BACKEND.md``.
    """
    if channel_2d.ndim != 2:
        raise ValueError(f"render_channel expects 2D input; got shape {channel_2d.shape}.")

    if idx == 0:
        img = _tinted_grayscale(channel_2d, _PURE_R)
    elif idx == 1:
        img = _tinted_grayscale(channel_2d, _PURE_G)
    elif idx == 2:
        img = _tinted_grayscale(channel_2d, _PURE_B)
    elif idx == 3:
        img = _ramp(channel_2d, PAPER_WHITE, NEAR_BLACK)
    elif idx == 4:
        img = _ramp(channel_2d, PAPER_WHITE, INDIGO)
    elif idx == 5:
        img = _ramp(channel_2d, PAPER_WHITE, VERMILION)
    elif idx == 6:
        img = _diverging(channel_2d, VERMILION, PAPER_WHITE, INDIGO)
    elif idx == 7:
        img = _ramp(channel_2d, PAPER_WHITE, INDIGO)
    elif idx == 8:
        img = _ramp(channel_2d, PAPER_WHITE, INDIGO)
    else:
        raise ValueError(f"render_channel: idx must be in 0..8; got {idx}.")
    return _encode_rgb(img)


def render_heatmap(prob_2d: np.ndarray) -> bytes:
    """Render a probability map (paper-white -> indigo)."""
    if prob_2d.ndim != 2:
        raise ValueError(f"render_heatmap expects 2D input; got shape {prob_2d.shape}.")
    return _encode_rgb(_ramp(prob_2d, PAPER_WHITE, INDIGO))


def render_binary(binary_2d: np.ndarray) -> bytes:
    """Render a binary mask as paper-white background with indigo foreground."""
    if binary_2d.ndim != 2:
        raise ValueError(f"render_binary expects 2D input; got shape {binary_2d.shape}.")
    mask = (binary_2d > 0).astype(np.float32)
    return _encode_rgb(_ramp(mask, PAPER_WHITE, INDIGO))


def render_instances(label_2d: np.ndarray) -> bytes:
    """Render a labeled instance map with one random color per instance.

    Background (label 0) is paper-white. Colors are deterministic across
    requests (``np.random.default_rng(42)``).
    """
    if label_2d.ndim != 2:
        raise ValueError(
            f"render_instances expects 2D input; got shape {label_2d.shape}."
        )
    h, w = label_2d.shape
    img = np.full((h, w, 3), PAPER_WHITE, dtype=np.uint8)
    if label_2d.max() <= 0:
        return _encode_rgb(img)

    max_label = int(label_2d.max())
    rng = np.random.default_rng(42)
    # One color per id 1..max_label (id 0 stays paper-white).
    palette = rng.integers(40, 230, size=(max_label + 1, 3), dtype=np.uint8)
    palette[0] = PAPER_WHITE

    img[:] = palette[label_2d]
    return _encode_rgb(img)


def render_overlay(
    rgb_u8: np.ndarray,
    label_2d: np.ndarray,
    *,
    opacity: float,
) -> tuple[bytes, bytes]:
    """Composite vermilion instance contours on top of the H&E image.

    Returns:
        ``(overlay_png, contours_rgba_png)``:
            ``overlay_png`` is the RGB composite at the given opacity, ready
            for a static ``<img>`` tag.
            ``contours_rgba_png`` is the contours layer alone with a
            transparent background; the frontend re-blends it client-side
            as the user moves the opacity slider.
    """
    if rgb_u8.dtype != np.uint8 or rgb_u8.ndim != 3 or rgb_u8.shape[2] != 3:
        raise ValueError(
            f"render_overlay expects (H, W, 3) uint8 RGB; got {rgb_u8.shape}/{rgb_u8.dtype}."
        )
    if label_2d.ndim != 2 or label_2d.shape != rgb_u8.shape[:2]:
        raise ValueError(
            f"render_overlay: label shape {label_2d.shape} must match "
            f"RGB spatial dims {rgb_u8.shape[:2]}."
        )
    opacity = float(max(0.0, min(1.0, opacity)))

    h, w = label_2d.shape
    contours_rgba = np.zeros((h, w, 4), dtype=np.uint8)

    label_i32 = label_2d.astype(np.int32, copy=False)
    unique_ids = np.unique(label_i32)
    unique_ids = unique_ids[unique_ids > 0]

    for iid in unique_ids:
        binary = (label_i32 == iid).astype(np.uint8)
        contours, _ = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
        if not contours:
            continue
        cv2.drawContours(contours_rgba, contours, -1,
                         (VERMILION[0], VERMILION[1], VERMILION[2], 255), 1)

    # Composite: only stroke pixels are blended.
    overlay = rgb_u8.copy()
    stroke_mask = contours_rgba[..., 3] > 0
    if stroke_mask.any() and opacity > 0.0:
        stroke_rgb = contours_rgba[..., :3].astype(np.float32)
        base_rgb = overlay.astype(np.float32)
        blended = (1.0 - opacity) * base_rgb + opacity * stroke_rgb
        overlay[stroke_mask] = np.clip(blended[stroke_mask], 0, 255).astype(np.uint8)

    return _encode_rgb(overlay), _encode_rgba(contours_rgba)


def to_b64(png_bytes: bytes) -> str:
    """Wrap raw PNG bytes as a ``data:image/png;base64,...`` string."""
    return "data:image/png;base64," + base64.b64encode(png_bytes).decode("ascii")