stain_preprocessing.py

"""
H&E stain preprocessing for histopathology patches.

Two complementary approaches (toggle via configure_stain):

1. **macenko** — Macenko-style stain normalization: estimate per-image stain
   vectors in optical density (OD) space, map stain concentrations to match a
   reference patch, then reconstruct RGB with the reference stain matrix.

2. **hed** — Ruderman-style path via CIE HED: convert RGB → HED (skimage),
   robustly rescale H and E channels, convert back to RGB. Lighter-weight than
   full OD deconvolution but still reduces some stain intensity drift.

Train / val / test should all call apply_stain_preprocessing so behavior matches.
"""

from __future__ import annotations

from typing import Any, Dict, Optional

import numpy as np

_CFG: Dict[str, Any] = {
    "method": "none",
    "beta": 0.15,
    "alpha": 1,
    "ref_S": None,  # (2, 3) rows = H, E OD basis for reconstruction
    "ref_maxC": None,  # (2,) 99th percentile concentrations on reference
}


def get_stain_config() -> Dict[str, Any]:
    """Return a shallow copy of the active stain configuration (for logging)."""
    return {k: (v.copy() if isinstance(v, np.ndarray) else v) for k, v in _CFG.items()}


def configure_stain(
    method: str = "none",
    *,
    reference_rgb_uint8: Optional[np.ndarray] = None,
    beta: float = 0.15,
    alpha: float = 1,
) -> None:
    """
    Select preprocessing and (for macenko) fit the reference stain basis.

    Parameters
    ----------
    method:
        ``none`` | ``macenko`` | ``hed`` (case-insensitive).
    reference_rgb_uint8:
        Required for ``macenko``: one RGB patch, uint8 or float in [0, 255],
        shape (H, W, 3). Typically a single training image read with OpenCV
        (BGR→RGB) so val/test normalize toward the same lab appearance.
    beta, alpha:
        Macenko OD background threshold and angular percentile (see Macenko et
        al., ISBI 2009); ``alpha`` and ``100 - alpha`` bracket the stain angle.
    """
    m = (method or "none").lower()
    _CFG["method"] = m
    _CFG["beta"] = float(beta)
    _CFG["alpha"] = float(alpha)
    _CFG["ref_S"] = None
    _CFG["ref_maxC"] = None

    if m == "macenko":
        if reference_rgb_uint8 is None:
            raise ValueError(
                "configure_stain(method='macenko') requires reference_rgb_uint8=..."
            )
        ref = np.clip(_as_uint8_rgb(reference_rgb_uint8), 0, 255)
        _CFG["ref_S"], _CFG["ref_maxC"] = _macenko_reference_stats(
            ref, beta=_CFG["beta"], alpha=_CFG["alpha"]
        )


def _as_uint8_rgb(img: np.ndarray) -> np.ndarray:
    x = np.asarray(img, dtype=np.float64)
    if x.size == 0:
        return x.astype(np.uint8)
    if x.max() <= 1.0 + 1e-3:
        x = x * 255.0
    return np.clip(x, 0, 255).astype(np.uint8)


def _rgb_to_od(rgb_u8: np.ndarray) -> np.ndarray:
    """Natural-log optical density; rgb_u8 in [0,255]."""
    v = np.clip(rgb_u8.astype(np.float64) / 255.0, 1e-6, 1.0)
    return -np.log(v)


def _od_to_rgb(od: np.ndarray) -> np.ndarray:
    rgb = np.exp(-np.clip(od, 1e-6, 50.0))
    return np.clip(rgb * 255.0, 0, 255).astype(np.float32)


def _normalize_rows(mat: np.ndarray) -> np.ndarray:
    n = np.linalg.norm(mat, axis=1, keepdims=True) + 1e-6
    return mat / n


def _default_he_matrix() -> np.ndarray:
    """Fallback H/E OD directions if too little tissue remains after masking."""
    return _normalize_rows(
        np.array([[0.65, 0.70, 0.29], [0.07, 0.95, 0.29]], dtype=np.float64)
    )


def get_macenko_stain_matrix(
    rgb_u8: np.ndarray, beta: float = 0.15, alpha: float = 1
) -> np.ndarray:
    """
    Estimate 2×3 stain matrix (rows ≈ H then E) in OD space.

    Follows the widely used StainTools / KatherLab formulation.
    """
    od = _rgb_to_od(rgb_u8).reshape((-1, 3))
    od = od[(od > beta).any(axis=1)]
    if od.shape[0] < 64:
        return _default_he_matrix()

    _, v = np.linalg.eigh(np.cov(od, rowvar=False))
    v = v[:, [2, 1]]
    if v[0, 0] < 0:
        v[:, 0] *= -1
    if v[0, 1] < 0:
        v[:, 1] *= -1
    that = od @ v
    phi = np.arctan2(that[:, 1], that[:, 0])
    min_phi = float(np.percentile(phi, alpha))
    max_phi = float(np.percentile(phi, 100.0 - alpha))
    v1 = v @ np.array([np.cos(min_phi), np.sin(min_phi)], dtype=np.float64)
    v2 = v @ np.array([np.cos(max_phi), np.sin(max_phi)], dtype=np.float64)
    he = np.vstack([v1, v2]) if v1[0] > v2[0] else np.vstack([v2, v1])
    return _normalize_rows(he)


def _conc_od(od_flat: np.ndarray, stain_matrix: np.ndarray) -> np.ndarray:
    """Pixel concentrations (N,2) from OD (N,3) and stain matrix (2,3)."""
    return od_flat @ np.linalg.pinv(stain_matrix)


def _percentile_max_conc(conc: np.ndarray) -> np.ndarray:
    conc = np.asarray(conc, dtype=np.float64)
    conc = np.maximum(conc, 1e-6)
    pos = (conc > 1e-5).all(axis=1)
    sample = conc[pos] if np.any(pos) else conc
    return np.maximum(np.percentile(sample, 99, axis=0), 1e-6)


def _macenko_reference_stats(
    rgb_u8: np.ndarray, beta: float, alpha: float
) -> tuple[np.ndarray, np.ndarray]:
    s_ref = get_macenko_stain_matrix(rgb_u8, beta=beta, alpha=alpha)
    od = _rgb_to_od(rgb_u8).reshape((-1, 3))
    c = _conc_od(od, s_ref)
    return s_ref, _percentile_max_conc(c)


def _apply_macenko(rgb_u8: np.ndarray) -> np.ndarray:
    s_ref = _CFG["ref_S"]
    max_ref = _CFG["ref_maxC"]
    if s_ref is None or max_ref is None:
        raise RuntimeError(
            "Macenko reference is unset. Call configure_stain('macenko', reference_rgb_uint8=...)."
        )

    beta, alpha = _CFG["beta"], _CFG["alpha"]
    s_src = get_macenko_stain_matrix(rgb_u8, beta=beta, alpha=alpha)
    od = _rgb_to_od(rgb_u8).reshape((-1, 3))
    c = _conc_od(od, s_src)
    max_src = _percentile_max_conc(c)
    c_n = c / max_src * max_ref
    od_new = (c_n @ s_ref).reshape(rgb_u8.shape)
    return _od_to_rgb(od_new)


def _apply_hed(rgb_u8: np.ndarray) -> np.ndarray:
    from skimage import color

    x = np.clip(rgb_u8.astype(np.float64) / 255.0, 0.0, 1.0)
    hed = color.rgb2hed(x)
    for i in range(2):
        ch = hed[..., i]
        lo, hi = np.percentile(ch, (1.0, 99.0))
        hed[..., i] = np.clip((ch - lo) / (hi - lo + 1e-6), 0.0, 1.0)
    out = color.hed2rgb(np.clip(hed, -0.2, 1.0))
    return (np.clip(out, 0.0, 1.0) * 255.0).astype(np.float32)


def apply_stain_preprocessing(img: np.ndarray) -> np.ndarray:
    """
    Apply the configured stain transform.

    Parameters
    ----------
    img : ndarray
        HWC RGB from Keras/OpenCV after BGR→RGB. uint8 or float in [0,255]
        (or [0,1] floats — they are scaled up internally).

    Returns
    -------
    ndarray float32
        RGB in [0, 255], consistent with this project's ``Rescaling(1/255)``.
    """
    method = _CFG["method"]
    if method == "none":
        return _ensure_float255(img)

    u8 = _as_uint8_rgb(img)
    if method == "macenko":
        return _apply_macenko(u8)
    if method == "hed":
        return _apply_hed(u8)
    return _ensure_float255(img)


def _ensure_float255(img: np.ndarray) -> np.ndarray:
    x = np.asarray(img, dtype=np.float32)
    if x.size and x.max() <= 1.0 + 1e-3:
        x = x * 255.0
    return np.clip(x, 0.0, 255.0)


def keras_preprocessing_fn(img: np.ndarray) -> np.ndarray:
    """Wrapper for ``ImageDataGenerator(preprocessing_function=...)``."""
    return apply_stain_preprocessing(img)