enhance.py

"""Visual enhancement modes for the displayed cutout (apoyo al revisor).

Same transforms and parameters as the static review site (docs/app.js), ported
to numpy/skimage. These are cosmetic, applied to the 8-bit RGB only — they do
NOT touch the image the model sees (that stays the raw Lupton composite).
"""

from __future__ import annotations

import numpy as np
from skimage.exposure import equalize_adapthist
from skimage.filters import unsharp_mask

MODES = ["Original", "Percentil", "Unsharp", "Asinh", "CLAHE", "Estructuras"]


def percentile_stretch(rgb: np.ndarray, p_lo: float = 1.0, p_hi: float = 99.5) -> np.ndarray:
    out = rgb.astype(np.float32)
    for c in range(3):
        lo, hi = np.percentile(out[..., c], [p_lo, p_hi])
        rng = (hi - lo) or 1.0
        out[..., c] = np.clip((out[..., c] - lo) / rng * 255.0, 0, 255)
    return out.astype(np.uint8)


def unsharp(rgb: np.ndarray, radius: float = 6.0, amount: float = 1.5) -> np.ndarray:
    f = rgb.astype(np.float32) / 255.0
    out = unsharp_mask(f, radius=radius, amount=amount, channel_axis=-1)
    return (np.clip(out, 0, 1) * 255).astype(np.uint8)


def asinh_stretch(rgb: np.ndarray, a: float = 1.5) -> np.ndarray:
    f = rgb.astype(np.float32) / 255.0
    out = np.arcsinh(a * f) / np.arcsinh(a)
    return (np.clip(out, 0, 1) * 255).astype(np.uint8)


def clahe(rgb: np.ndarray, clip_limit: float = 0.01) -> np.ndarray:
    f = rgb.astype(np.float32) / 255.0
    out = np.zeros_like(f)
    for c in range(3):
        out[..., c] = equalize_adapthist(f[..., c], clip_limit=clip_limit)
    return (np.clip(out, 0, 1) * 255).astype(np.uint8)


def estructuras(rgb: np.ndarray, amount: float = 1.5) -> np.ndarray:
    x = percentile_stretch(rgb, 1.0, 99.5)
    x = unsharp(x, 6.0, amount)
    x = asinh_stretch(x, amount * 0.5)
    return x


def to_gray(rgb: np.ndarray) -> np.ndarray:
    lum = (0.2126 * rgb[..., 0] + 0.7152 * rgb[..., 1] + 0.0722 * rgb[..., 2]).astype(np.uint8)
    return np.dstack([lum, lum, lum])


def apply(rgb: np.ndarray, mode: str = "Original", gray: bool = False) -> np.ndarray:
    """Apply the chosen enhancement mode (and optional grayscale) to an RGB."""
    if mode == "Percentil":
        rgb = percentile_stretch(rgb)
    elif mode == "Unsharp":
        rgb = unsharp(rgb)
    elif mode == "Asinh":
        rgb = asinh_stretch(rgb)
    elif mode == "CLAHE":
        rgb = clahe(rgb)
    elif mode == "Estructuras":
        rgb = estructuras(rgb)
    if gray:
        rgb = to_gray(rgb)
    return rgb