src/vessels.py

"""
Rigorous retinal vessel analysis for gradability.

The visibility and integrity of the vascular tree is the single most reliable
proxy for whether a fundus image is gradable. We compute a multi-scale vessel
probability map and derive structural descriptors (density, mean vesselness,
skeletal length, fractal dimension, fragmentation) that degrade predictably
under blur, haze and shadow.

Two segmentation backends:
  * classical (default, no weights): CLAHE + multi-scale Frangi vesselness on the
    inverted green channel, restricted to the analysis disc, with hysteresis and
    morphological cleanup. Validated to be blur-monotonic.
  * deep (optional): register any callable  seg(rgb_uint8) -> float32 HxW in [0,1]
    via  set_deep_segmenter(). Use this to plug in a trained vessel or artery/vein
    U-Net on your host; the rest of the pipeline is identical.
"""

from __future__ import annotations
import numpy as np
import cv2
from skimage.filters import frangi
from skimage.morphology import skeletonize
from scipy import ndimage as ndi

_DEEP_SEGMENTER = None
_AV_SEGMENTER = None


def set_av_segmenter(fn):
    """Register a deep artery/vein/vessel segmenter:
    fn(rgb_uint8) -> dict(vessel, artery, vein), each float32 HxW in [0,1]."""
    global _AV_SEGMENTER
    _AV_SEGMENTER = fn


def has_av_segmenter():
    return _AV_SEGMENTER is not None


def _remove_small(binary, min_size):
    """Drop connected components smaller than min_size (deprecation-free)."""
    lbl, n = ndi.label(binary)
    if n == 0:
        return binary
    sizes = ndi.sum(np.ones_like(lbl), lbl, index=np.arange(1, n + 1))
    keep = np.zeros(n + 1, bool)
    keep[1:][sizes >= min_size] = True
    return keep[lbl]


def tophat_vessel_contrast(rgb, mask):
    """Blur-monotonic vessel-contrast energy: black-top-hat on the green channel
    isolates dark vessels; its energy inside the disc collapses under blur.
    Robust backbone of the vessel gradability score."""
    g = rgb[..., 1].astype(np.uint8)
    g = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8)).apply(g)
    ker = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (13, 13))
    bth = cv2.morphologyEx(g, cv2.MORPH_BLACKHAT, ker).astype(np.float32)
    inside = bth[mask]
    if inside.size == 0:
        return 0.0
    return float(np.sqrt(np.mean(inside ** 2)))


def set_deep_segmenter(fn):
    """Register a deep vessel segmenter: fn(rgb_uint8 HxWx3) -> prob map HxW in [0,1]."""
    global _DEEP_SEGMENTER
    _DEEP_SEGMENTER = fn


def has_deep_segmenter():
    return _DEEP_SEGMENTER is not None


# ----------------------------------------------------------------- preprocessing
def _green_clahe(rgb, mask):
    g = rgb[..., 1].astype(np.uint8)
    clahe = cv2.createCLAHE(clipLimit=2.5, tileGridSize=(8, 8))
    g = clahe.apply(g)
    # illumination flattening inside the disc
    bg = cv2.medianBlur(g, 55)
    flat = cv2.normalize(g.astype(np.float32) - bg.astype(np.float32), None, 0, 1,
                         cv2.NORM_MINMAX)
    flat[~mask] = 0
    return flat


# ----------------------------------------------------------------- segmentation
def _classical_vesselness(rgb, mask):
    flat = _green_clahe(rgb, mask)
    # multi-scale Frangi: vessels are dark ridges on inverted green
    inv = flat.max() - flat
    scales = np.arange(1.0, 5.5, 0.75)
    v = frangi(inv, sigmas=scales, black_ridges=False)
    v = cv2.normalize(v, None, 0, 1, cv2.NORM_MINMAX)
    v[~mask] = 0
    return v.astype(np.float32)


def segment_vessels(rgb, mask):
    """Return a vessel probability map in [0,1] restricted to `mask`."""
    if _AV_SEGMENTER is not None:
        try:
            av = _AV_SEGMENTER(rgb)
            p = np.asarray(av["vessel"], np.float32)
            p[~mask] = 0
            return p, "deep-rrwnet"
        except Exception:
            pass
    if _DEEP_SEGMENTER is not None:
        try:
            p = np.asarray(_DEEP_SEGMENTER(rgb)).astype(np.float32)
            if p.max() > 1.0:
                p = p / 255.0
            p[~mask] = 0
            return p, "deep"
        except Exception:
            pass  # fall back to classical
    return _classical_vesselness(rgb, mask), "classical"


def _caliber(binary):
    """Mean vessel caliber = area / skeleton length (px)."""
    n = int(binary.sum())
    if n < 20:
        return 0.0
    sk = skeletonize(binary)
    L = int(sk.sum())
    return float(n / L) if L > 0 else 0.0


def compute_avr(artery, vein, mask, thr=0.5):
    """Arteriolar-to-venular ratio from A/V probability maps."""
    a = (artery > thr) & mask
    v = (vein > thr) & mask
    ca, cv = _caliber(a), _caliber(v)
    avr = float(ca / cv) if cv > 0 else float("nan")
    return dict(avr=avr, artery_caliber=ca, vein_caliber=cv,
                artery_frac=float(a.sum()) / max(int(mask.sum()), 1),
                vein_frac=float(v.sum()) / max(int(mask.sum()), 1))


def _binarize(vprob, mask):
    """Hysteresis threshold -> clean binary vessel mask."""
    inside = vprob[mask]
    if inside.size == 0:
        return np.zeros_like(vprob, bool)
    hi = np.quantile(inside, 0.93)
    lo = np.quantile(inside, 0.80)
    strong = vprob >= hi
    weak = vprob >= lo
    # keep weak pixels connected to strong (hysteresis)
    lbl, n = ndi.label(weak)
    keep = np.zeros(n + 1, bool)
    keep[np.unique(lbl[strong])] = True
    keep[0] = False
    b = keep[lbl] & mask
    b = ndi.binary_closing(b, structure=np.ones((3, 3)))
    b = _remove_small(b, 40)
    return b


# ----------------------------------------------------------------- descriptors
def _fractal_dimension(binary):
    """Box-counting fractal dimension of the vessel skeleton."""
    Z = binary > 0
    if Z.sum() < 20:
        return 0.0
    p = min(Z.shape)
    n = 2 ** int(np.floor(np.log2(p)))
    sizes = 2 ** np.arange(int(np.log2(n)), 1, -1)
    counts = []
    for size in sizes:
        cnt = 0
        for i in range(0, Z.shape[0], size):
            for j in range(0, Z.shape[1], size):
                if Z[i:i+size, j:j+size].any():
                    cnt += 1
        counts.append(cnt)
    counts = np.array(counts, float)
    sizes = np.array(sizes, float)
    ok = counts > 0
    if ok.sum() < 2:
        return 0.0
    coeffs = np.polyfit(np.log(1.0 / sizes[ok]), np.log(counts[ok]), 1)
    return float(coeffs[0])


def analyze_vessels(rgb, fov):
    """Full vessel analysis. Returns a descriptor dict + maps."""
    from .qc_metrics import inner_disc_mask
    mask = inner_disc_mask(rgb.shape, fov, frac=0.92)
    vprob, backend = segment_vessels(rgb, mask)
    vbin = _binarize(vprob, mask)
    area = max(int(mask.sum()), 1)

    density = float(vbin.sum()) / area                       # vessel area fraction
    mean_vesselness = float(vprob[mask].mean())              # continuous strength
    skel = skeletonize(vbin)
    skel_len = float(skel.sum()) / area                      # normalised length
    fractal = _fractal_dimension(skel)                       # branching complexity
    lbl, ncomp = ndi.label(vbin)
    frag = ncomp / (skel.sum() + 1e-6) * 100                 # comps per 100 skel px
    contrast = tophat_vessel_contrast(rgb, mask)             # monotonic backbone

    out = dict(
        backend=backend, density=density, mean_vesselness=mean_vesselness,
        skeleton_length=skel_len, fractal_dimension=fractal, fragmentation=frag,
        n_components=int(ncomp), contrast=contrast,
        prob_map=vprob, binary=vbin, skeleton=skel, mask=mask,
    )
    # artery/vein maps + AVR biomarker when a deep A/V segmenter is registered
    if _AV_SEGMENTER is not None:
        try:
            av = _AV_SEGMENTER(rgb)
            out["artery"] = np.asarray(av["artery"], np.float32)
            out["vein"] = np.asarray(av["vein"], np.float32)
            out["avr"] = compute_avr(out["artery"], out["vein"], mask)
        except Exception:
            pass
    return out


def vessel_gradability_score(desc):
    """0-1 vessel gradability. The blur-monotonic top-hat contrast is the primary
    driver; structural descriptors (rigorous with a deep segmenter) refine it."""
    def curve(x, x_fail, x_pass):
        if x_pass == x_fail:
            return 0.5
        return float(np.clip((x - x_fail) / (x_pass - x_fail), 0, 1))

    s_contrast = curve(desc["contrast"], 6.0, 20.0)           # PRIMARY, monotonic
    s_fractal = curve(desc["fractal_dimension"], 1.05, 1.55)  # branching richness
    s_frag = 1.0 - curve(desc["fragmentation"], 3.0, 14.0)    # fragmentation

    if desc["backend"] == "deep":
        # trust structural evidence fully with a real segmentation
        parts = np.array([s_contrast, s_fractal, s_frag])
        score = 0.4 * parts.mean() + 0.6 * parts.min()
    else:
        # classical structural terms are approximate: contrast-dominated
        score = 0.75 * s_contrast + 0.15 * s_fractal + 0.10 * s_frag
    return float(np.clip(score, 0, 1)), dict(
        contrast=s_contrast, fractal=s_fractal, fragmentation=s_frag)