stenosis.py

# stenosis.py
import numpy as np
import cv2
from skimage.morphology import skeletonize
from scipy.ndimage import distance_transform_edt

def detect_stenosis(mask: np.ndarray):

    mask = (mask > 0).astype(np.uint8)
    skeleton = skeletonize(mask).astype(np.uint8)
    dist = distance_transform_edt(mask)

    coords = np.column_stack(np.where(skeleton > 0))
    if len(coords) < 30:
        return 0.0, None

    radii = np.array([dist[y, x] for y, x in coords])

    # smooth radius (VERY IMPORTANT)
    radii = np.convolve(radii, np.ones(7)/7, mode="same")

    ref_radius = np.percentile(radii, 85)
    min_radius = radii.min()

    severity = float((1 - min_radius / ref_radius) * 100)

    # pick worst region center
    worst = np.where(radii < ref_radius * 0.6)[0]
    if len(worst) == 0:
        return severity, None

    idx = worst[len(worst)//2]
    y, x = coords[idx]

    # force visible bounding box
    box = max(int(ref_radius * 2.5), 20)
    
    h, w = mask.shape
    x1 = max(0, x - box)
    y1 = max(0, y - box)
    x2 = min(w, x + box)
    y2 = min(h, y + box)

    return severity, (x1, y1, x2, y2)