postprocessing.py

"""
postprocessing.py
------------------
Turns a raw U-Net probability map into a clean, thin, single-pixel-wide
boundary line -- the same visual style as the paper's hand-drawn
"Ground truth" panel (see reference figure) -- instead of the thick,
speckled blob you get from a bare `prob > threshold` mask.

Also provides the accuracy metrics needed to actually *know* how close a
prediction is to a reference mask (IoU, Dice, boundary-tolerant F1,
pixel accuracy). There is no code-only way to guarantee "95% accuracy" --
that number depends entirely on how good your training data is -- but
this module is what lets you measure and drive that number honestly
during real training, instead of eyeballing a figure.
"""

import numpy as np
import cv2
from scipy import ndimage as ndi
from skimage.morphology import skeletonize, remove_small_objects, disk
try:
    from skimage.morphology import binary_closing
except ImportError:  # pragma: no cover
    from skimage.morphology import closing as binary_closing


def clean_boundary_mask(prob_map, threshold=0.5, min_size=25,
                         close_radius=1, skeleton=True, line_thickness=1):
    """
    Convert a raw probability map (float32 in [0,1]) into a clean boundary
    mask that visually matches a hand-annotated ground-truth line drawing:
    thresholded -> small-blob removal -> morphological closing (bridges
    1-2px gaps in the ridge) -> skeletonize to 1px -> optional re-thicken.

    Parameters
    ----------
    prob_map      : float32 array in [0,1], output of sliding_window_predict
    threshold     : probability cutoff
    min_size      : remove connected components smaller than this many px
                    (kills isolated speckle noise, the #1 reason U-Net
                    overlays look "blobby" instead of clean lines)
    close_radius   : binary closing radius, bridges small gaps in a ridge
                     before skeletonizing (0 to disable)
    skeleton       : if True, thin to a 1px-wide centerline
    line_thickness : final line width in px (re-dilate after skeletonizing)

    Returns
    -------
    uint8 {0,255} mask, thin clean boundary lines
    """
    binary = prob_map > threshold

    if min_size > 0:
        binary = remove_small_objects(binary, min_size=min_size)

    if close_radius > 0:
        binary = binary_closing(binary, disk(close_radius))

    if skeleton:
        binary = skeletonize(binary)

    mask = (binary * 255).astype(np.uint8)

    if line_thickness > 1:
        kernel = np.ones((line_thickness, line_thickness), np.uint8)
        mask = cv2.dilate(mask, kernel)

    return mask


def draw_yellow_boundary(color_img, boundary_mask, color=(255, 210, 0), thickness=1):
    """
    Draw a boundary mask as a thin yellow line over the original color
    image -- matching the "Ground truth" panel style in the reference
    figure, rather than a solid translucent blob fill.
    """
    overlay = color_img.copy()
    mask = boundary_mask
    if thickness > 1:
        kernel = np.ones((thickness, thickness), np.uint8)
        mask = cv2.dilate(mask, kernel)
    ys, xs = np.where(mask > 0)
    overlay[ys, xs] = color
    return overlay


# ─────────────────────────────────────────────────────────────────────────
# Accuracy metrics
# ─────────────────────────────────────────────────────────────────────────
def _tolerant_match(pred, gt, tol_px=2):
    """
    Boundary detection is inherently a "close enough" problem: a predicted
    line one pixel off from the hand-drawn line is still a correct
    detection, not an error. We dilate each mask by `tol_px` before
    computing precision/recall so near-miss pixels count as matches --
    standard practice for boundary/edge evaluation (c.f. BSDS boundary F1).
    """
    if tol_px > 0:
        k = np.ones((2 * tol_px + 1, 2 * tol_px + 1), np.uint8)
        gt_dil = cv2.dilate(gt.astype(np.uint8), k) > 0
        pred_dil = cv2.dilate(pred.astype(np.uint8), k) > 0
    else:
        gt_dil = gt.astype(bool)
        pred_dil = pred.astype(bool)
    return pred_dil, gt_dil


def compute_metrics(pred_mask, gt_mask, tol_px=2):
    """
    Compare a predicted boundary mask against a reference (real ground
    truth or pseudo-label) mask.

    Returns a dict with:
      iou            : intersection-over-union on raw (non-dilated) masks
      dice           : Dice / F1 on raw masks
      pixel_accuracy : fraction of all pixels (boundary + background) that
                        match -- usually near-meaningless here because
                        boundary pixels are a small minority of the image,
                        so a model predicting "no boundary" can still score
                        95%+ on this metric alone. Reported for reference,
                        not as the headline number.
      boundary_precision, boundary_recall, boundary_f1 : computed with a
                        `tol_px`-pixel tolerance band, which is the metric
                        that actually reflects "did we find the same
                        boundary lines" for this kind of task
    """
    pred = pred_mask > 0
    gt = gt_mask > 0

    inter = np.logical_and(pred, gt).sum()
    union = np.logical_or(pred, gt).sum()
    iou = inter / union if union > 0 else 1.0
    dice = 2 * inter / (pred.sum() + gt.sum()) if (pred.sum() + gt.sum()) > 0 else 1.0
    pixel_acc = (pred == gt).sum() / pred.size

    pred_dil, gt_dil = _tolerant_match(pred, gt, tol_px=tol_px)
    tp_prec = np.logical_and(pred, gt_dil).sum()   # predicted px that land near a GT px
    tp_rec = np.logical_and(gt, pred_dil).sum()    # GT px that have a predicted px nearby
    precision = tp_prec / pred.sum() if pred.sum() > 0 else 1.0
    recall = tp_rec / gt.sum() if gt.sum() > 0 else 1.0
    f1 = 2 * precision * recall / (precision + recall) if (precision + recall) > 0 else 0.0

    return {
        "iou": float(iou),
        "dice": float(dice),
        "pixel_accuracy": float(pixel_acc),
        "boundary_precision": float(precision),
        "boundary_recall": float(recall),
        "boundary_f1": float(f1),
    }