LEGENDFTW/image-segmentation-explorer

1"""segmentation.py — Thresholding, K-means, and Watershed implementations."""2 3import numpy as np4import cv25from PIL import Image6 7 8METHOD_DESCRIPTIONS = {9    "Thresholding": (10        "Converts the image to greyscale, then labels every pixel as "11        "foreground (above threshold) or background (below threshold). "12        "Simple and fast. Works best when objects are clearly brighter or darker than the background."13    ),14    "K-means": (15        "Groups pixels into K clusters based on colour similarity. "16        "Each cluster gets a single representative colour. "17        "Good for colour-based segmentation; does not use spatial information."18    ),19    "Watershed": (20        "Treats pixel intensity as a topographic surface and 'floods' it from marked seed points. "21        "Boundaries form where flooding from different seeds meets. "22        "Excellent at separating touching or overlapping objects."23    ),24}25 26 27def segment_threshold(image: np.ndarray, threshold: int, mode: str) -> tuple:28    """29    Returns (segmented_rgb, mask) where mask is a binary uint8 array.30    mode: 'Binary', 'Otsu', 'Adaptive'31    """32    gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)33 34    if mode == "Otsu":35        _, mask = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)36    elif mode == "Adaptive":37        mask = cv2.adaptiveThreshold(38            gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C,39            cv2.THRESH_BINARY, 11, 240        )41    else:  # Binary42        _, mask = cv2.threshold(gray, threshold, 255, cv2.THRESH_BINARY)43 44    # Colour the result: foreground = original, background = dark grey45    result = image.copy()46    result[mask == 0] = [40, 40, 40]47    return result, mask48 49 50def segment_kmeans(image: np.ndarray, k: int) -> tuple:51    """52    Returns (segmented_rgb, label_map).53    Each pixel is replaced by its cluster's mean colour.54    """55    h, w = image.shape[:2]56    pixels = image.reshape(-1, 3).astype(np.float32)57 58    criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 20, 1.0)59    _, labels, centers = cv2.kmeans(60        pixels, k, None, criteria, 5, cv2.KMEANS_RANDOM_CENTERS61    )62    centers = np.uint8(centers)63    segmented = centers[labels.flatten()].reshape(h, w, 3)64    label_map = labels.reshape(h, w).astype(np.uint8)65    return segmented, label_map66 67 68def segment_watershed(image: np.ndarray, min_distance: int) -> tuple:69    """70    Returns (segmented_rgb, label_map).71    Uses distance transform + local maxima as seeds.72    """73    from skimage.segmentation import watershed74    from skimage.feature import peak_local_max75    from scipy import ndimage as ndi76 77    gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)78 79    # Threshold to get foreground mask80    _, thresh = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)81 82    # Distance transform83    dist = ndi.distance_transform_edt(thresh)84 85    # Find peaks (seeds)86    coords = peak_local_max(dist, min_distance=max(1, min_distance),87                            labels=thresh)88    mask_peaks = np.zeros(dist.shape, dtype=bool)89    mask_peaks[tuple(coords.T)] = True90    markers, _ = ndi.label(mask_peaks)91 92    # Watershed93    labels = watershed(-dist, markers, mask=thresh)94 95    # Colour each region with a distinct colour96    result = np.zeros_like(image)97    palette = _make_palette(labels.max() + 1)98    for region_id in range(1, labels.max() + 1):99        result[labels == region_id] = palette[region_id % len(palette)]100    result[labels == 0] = [30, 30, 30]101 102    return result, labels.astype(np.int32)103 104 105def compute_segment_metrics(label_map: np.ndarray, method: str) -> dict:106    """Return basic segmentation metrics."""107    if method == "Thresholding":108        fg = int(np.sum(label_map > 0))109        bg = int(np.sum(label_map == 0))110        total = label_map.size111        return {112            "num_segments": 2,113            "foreground_pct": round(fg / total * 100, 1),114            "background_pct": round(bg / total * 100, 1),115        }116    elif method == "K-means":117        k = int(label_map.max()) + 1118        sizes = {f"Cluster {i}": int(np.sum(label_map == i)) for i in range(k)}119        return {120            "num_segments": k,121            "cluster_sizes": sizes,122        }123    else:  # Watershed124        ids = [i for i in np.unique(label_map) if i > 0]125        return {126            "num_segments": len(ids),127            "avg_region_size": int(np.mean([np.sum(label_map == i) for i in ids])) if ids else 0,128        }129 130 131def _make_palette(n: int) -> list:132    """Generate n visually distinct colours."""133    palette = [134        [231, 76,  60],  [46, 204, 113], [52, 152, 219],135        [155, 89, 182], [241,196,  15], [230,126,  34],136        [26, 188, 156], [236, 64, 122], [100,181,246],137        [174,213,129], [255,167,  38], [171, 71, 188],138    ]139    while len(palette) < n:140        palette += palette141    return palette142