"""Segmentation-map CNN for glaucoma grading. Trains a CNN on combined disc/cup segmentation maps — pixel values 0 = background, 1 = disc (rim only), 2 = cup — instead of raw RGB fundus images, forcing the model to learn purely from optic nerve head geometry (CDR, rim width, cup location, etc.). Two segmentation sources are supported: gt – rasterise expert contour/mask annotations directly (pure NumPy/PIL, no CUDA — safe in DataLoader worker processes) unet – run a trained UNetSegmenter on the raw fundus image Usage (import from training script): from v3.classes.seg_cnn import SegMapRecord, SegMapDataset, SegCNN, seg_map_to_tensor """ from __future__ import annotations from dataclasses import dataclass from pathlib import Path from typing import List, Optional, Tuple import numpy as np import torch import torch.nn as nn from PIL import Image, ImageDraw from PIL.Image import Resampling from torch.utils.data import Dataset from torchvision import models, transforms from tqdm import tqdm # --------------------------------------------------------------------------- # Data record # --------------------------------------------------------------------------- @dataclass class SegMapRecord: """One labelled eye sample for the seg-map CNN.""" sample_id: str image_path: Path # original fundus image (used by unet mode) annotation_disc: Path # contour (.txt) or mask (.bmp/.png) annotation_cup: Path annotation_type_disc: str # "contour" or "mask" annotation_type_cup: str patient_id: int # for group-CV: keep both eyes of a patient together eye: str # "OD" or "OS" label: int # 0 = Normal, 1 = Glaucoma # --------------------------------------------------------------------------- # Seg-map utilities # --------------------------------------------------------------------------- def _combine_masks(disc_mask: np.ndarray, cup_mask: np.ndarray) -> np.ndarray: """ Combine binary disc and cup masks into a 3-class label map. Returns a uint8 array with values: 0 — background 1 — optic disc rim (disc but not cup) 2 — optic cup """ disc = (disc_mask > 0).astype(np.uint8) cup = (cup_mask > 0).astype(np.uint8) cup = (cup & disc) # structural prior: cup must be inside disc seg = disc + cup # 0, 1 (rim), or 2 (cup) return seg.astype(np.uint8) def crop_to_disc(seg_map: np.ndarray) -> np.ndarray: """ Crop a seg map tightly to the disc bounding box. The disc is anywhere seg_map > 0 (i.e. rim or cup). Returns the original array unchanged if no disc is found. """ rows = np.any(seg_map > 0, axis=1) cols = np.any(seg_map > 0, axis=0) if not rows.any(): return seg_map r0, r1 = int(np.argmax(rows)), int(len(rows) - 1 - np.argmax(rows[::-1])) c0, c1 = int(np.argmax(cols)), int(len(cols) - 1 - np.argmax(cols[::-1])) return seg_map[r0:r1 + 1, c0:c1 + 1] def seg_map_to_tensor( seg_map: np.ndarray, channels: int, target_size: int, ) -> torch.Tensor: """ Convert an (H, W) seg map with values {0, 1, 2} to a float tensor. channels=1 → (1, H, W) float in [0, 1] (values 0/0.5/1.0) channels=3 → (3, H, W) one-hot binary channels [bg, disc_rim, cup] """ pil = Image.fromarray(seg_map.astype(np.uint8), mode="L") pil = pil.resize((target_size, target_size), Image.NEAREST) seg = np.array(pil, dtype=np.uint8) if channels == 1: arr = seg.astype(np.float32) / 2.0 # {0, 0.5, 1.0} return torch.from_numpy(arr).unsqueeze(0) if channels == 3: bg = (seg == 0).astype(np.float32) disc_rim = (seg == 1).astype(np.float32) cup = (seg == 2).astype(np.float32) return torch.from_numpy(np.stack([bg, disc_rim, cup], axis=0)) raise ValueError(f"channels must be 1 or 3, got {channels}") # --------------------------------------------------------------------------- # GT mask loading (pure NumPy / PIL — no CUDA, safe in DataLoader workers) # --------------------------------------------------------------------------- def _load_contour(path: Path) -> np.ndarray: """Load x,y contour pairs from a whitespace- or comma-delimited text file.""" for delimiter in (",", None): try: arr = np.loadtxt(str(path), delimiter=delimiter, comments="#", dtype=np.float32) if arr.size > 0: break except Exception: arr = np.zeros((0, 2), dtype=np.float32) if arr.size == 0 or arr.ndim == 1: return np.zeros((0, 2), dtype=np.float32) if arr.shape[1] < 2: return np.zeros((0, 2), dtype=np.float32) return arr[:, :2] def _contour_to_mask( coords: np.ndarray, image_size: Tuple[int, int], target_size: int ) -> np.ndarray: """ Rasterise a polygon defined by (x, y) coords into a binary mask. image_size is the (width, height) of the original fundus image — the coordinate space the contour was annotated in. The mask is drawn at that resolution then resized to target_size, matching UNetSegmenter's behaviour and avoiding off-canvas clipping. """ if coords is None or len(coords) < 3: return np.zeros((target_size, target_size), dtype=np.uint8) points = [tuple(map(float, pt)) for pt in coords] img = Image.new("L", image_size, 0) ImageDraw.Draw(img).polygon(points, outline=1, fill=1) img = img.resize((target_size, target_size), Resampling.NEAREST) return (np.array(img, dtype=np.uint8) > 0).astype(np.uint8) def _extract_masks_from_image( mask_path: Path, target_size: int ) -> Tuple[np.ndarray, np.ndarray]: """ Extract disc and cup binary masks from a segmentation image file. Handles both grayscale label images (e.g. REFUGE .bmp) and RGB colour-coded masks. Returns (disc_mask, cup_mask) both at target_size × target_size. """ raw = Image.open(mask_path) arr = np.array(raw) if arr.ndim == 2: # Grayscale: identify background from edge statistics edges = np.concatenate([arr[0], arr[-1], arr[:, 0], arr[:, -1]]) bg_val = int(np.argmax(np.bincount(edges.astype(np.int64).clip(0, 255), minlength=256))) disc_arr = (arr != bg_val).astype(np.uint8) vals = np.unique(arr) non_bg = vals[vals != bg_val] cup_arr: np.ndarray if non_bg.size > 1: cup_val = int(non_bg.min()) cup_arr = (arr == cup_val).astype(np.uint8) else: cup_arr = np.zeros_like(disc_arr, dtype=np.uint8) else: img_rgb = raw.convert("RGB") arr = np.array(img_rgb) h, w, c = arr.shape edges_rgb = np.concatenate( [arr[0], arr[-1], arr[:, 0], arr[:, -1]], axis=0 ) edge_colors, edge_counts = np.unique(edges_rgb.reshape(-1, c), axis=0, return_counts=True) bg_color = edge_colors[int(np.argmax(edge_counts))] colors, counts = np.unique(arr.reshape(-1, c), axis=0, return_counts=True) not_bg = np.any(colors != bg_color.reshape(1, -1), axis=1) colors, counts = colors[not_bg], counts[not_bg] disc_arr = np.zeros((h, w), dtype=np.uint8) cup_arr = np.zeros((h, w), dtype=np.uint8) if colors.shape[0] >= 1: order = np.argsort(-counts) disc_color = colors[order[0]] disc_arr[np.all(arr == disc_color, axis=-1)] = 1 if colors.shape[0] >= 2: cup_color = colors[order[1]] cup_arr[np.all(arr == cup_color, axis=-1)] = 1 # Resize to target_size with nearest-neighbour to preserve binary values def _resize(m: np.ndarray) -> np.ndarray: pil = Image.fromarray((m > 0).astype(np.uint8) * 255) pil = pil.resize((target_size, target_size), Resampling.NEAREST) return (np.array(pil) > 0).astype(np.uint8) return _resize(disc_arr), _resize(cup_arr) def load_gt_masks(rec: "SegMapRecord", target_size: int) -> Tuple[np.ndarray, np.ndarray]: """ Load GT disc + cup masks for one record. Handles annotation_type "contour" (x,y text file) and "mask" (image file). Returns (disc_mask, cup_mask) as uint8 arrays of shape (target_size, target_size). """ disc_mask: Optional[np.ndarray] = None cup_mask: Optional[np.ndarray] = None # Get original image size so contour coordinates are drawn in the right space with Image.open(rec.image_path) as _img: image_size = _img.size # (width, height) # ---- Disc ---- if rec.annotation_type_disc == "mask": disc_mask, cup_from_disc = _extract_masks_from_image(rec.annotation_disc, target_size) if cup_from_disc.any(): cup_mask = cup_from_disc else: # contour coords = _load_contour(rec.annotation_disc) disc_mask = _contour_to_mask(coords, image_size, target_size) # ---- Cup ---- if cup_mask is None: if rec.annotation_type_cup == "mask": _, cup_from_cup = _extract_masks_from_image(rec.annotation_cup, target_size) cup_mask = cup_from_cup else: # contour coords = _load_contour(rec.annotation_cup) cup_mask = _contour_to_mask(coords, image_size, target_size) if disc_mask is None: disc_mask = np.zeros((target_size, target_size), dtype=np.uint8) if cup_mask is None: cup_mask = np.zeros((target_size, target_size), dtype=np.uint8) # Structural prior: cup must lie within disc cup_mask = (cup_mask > 0) & (disc_mask > 0) return disc_mask.astype(np.uint8), cup_mask.astype(np.uint8) # --------------------------------------------------------------------------- # U-Net fine-tuning dataset # --------------------------------------------------------------------------- class UNetFineTuneDataset(Dataset): """ Loads (image_tensor, mask_tensor) pairs for fine-tuning the U-Net on PAPILA GT annotations. Uses the same preprocessing as UNetSegmenter so the fine-tuned weights are compatible with inference. """ def __init__( self, records: List[SegMapRecord], target_size: int = 512, normalize: str = "per_image", ) -> None: self.records = records self.target_size = target_size self.normalize = normalize self.to_tensor = transforms.ToTensor() def __len__(self) -> int: return len(self.records) def _normalize(self, tensor: torch.Tensor) -> torch.Tensor: if self.normalize == "per_image": mean = tensor.mean(dim=(1, 2), keepdim=True) std = tensor.std(dim=(1, 2), keepdim=True).clamp(min=1e-6) return (tensor - mean) / std if self.normalize == "imagenet": mean = torch.tensor([0.485, 0.456, 0.406]).view(-1, 1, 1) std = torch.tensor([0.229, 0.224, 0.225]).view(-1, 1, 1) return (tensor - mean) / std return tensor def __getitem__(self, idx: int): rec = self.records[idx] image = Image.open(rec.image_path).convert("RGB") image = image.resize((self.target_size, self.target_size), Resampling.BILINEAR) img_tensor = self._normalize(self.to_tensor(image)) disc_mask, cup_mask = load_gt_masks(rec, self.target_size) mask_tensor = torch.from_numpy( np.stack([disc_mask, cup_mask], axis=0).astype(np.float32) ) return img_tensor, mask_tensor # --------------------------------------------------------------------------- # U-Net precomputation (run once per full record list, not per fold) # --------------------------------------------------------------------------- def precompute_unet_seg_maps( records: List["SegMapRecord"], segmenter, threshold: float = 0.5, ) -> List[np.ndarray]: """ Run the U-Net on every record and return a list of combined seg maps. Call this once before the CV loop and pass the results to each fold's SegMapDataset via precomputed_seg_maps, so the U-Net isn't re-run per fold. """ to_tensor = transforms.ToTensor() seg_maps = [] for rec in tqdm(records, desc="U-Net inference", unit="img", leave=False): image = Image.open(rec.image_path).convert("RGB") resized = segmenter.preprocess_image(image) tensor = segmenter._normalize_tensor( to_tensor(resized).to(segmenter.device) ).unsqueeze(0) with torch.no_grad(): logits = segmenter.model(tensor) probs = torch.sigmoid(logits)[0].cpu().numpy() disc = (probs[0] > threshold).astype(np.uint8) cup = (probs[1] > threshold).astype(np.uint8) cup = (cup & disc) seg_maps.append(_combine_masks(disc, cup.astype(np.uint8))) return seg_maps # --------------------------------------------------------------------------- # Dataset # --------------------------------------------------------------------------- class SegMapDataset(Dataset): """ PyTorch Dataset that yields (seg_tensor, label) pairs. Parameters ---------- records : list of SegMapRecord target_size : CNN input spatial size (images are resized to this) channels : 1 = single-channel label map; 3 = one-hot three channels augment : apply random flips + rotation (for training set) unet_segmenter : if provided, use U-Net predictions instead of GT masks; must be a loaded UNetSegmenter with model weights set unet_threshold : threshold for U-Net logit → binary mask seg_target_size: resolution at which GT masks are rasterised (or U-Net output size). Default 512 matches UNetSegmenter default. crop_to_disc : crop the seg map tightly to the disc bounding box before resizing to target_size (default True — eliminates the background zeros that make up most of the full image) """ def __init__( self, records: List[SegMapRecord], target_size: int = 224, channels: int = 3, augment: bool = False, unet_segmenter=None, unet_threshold: float = 0.5, seg_target_size: int = 512, crop_to_disc: bool = True, precomputed_seg_maps: Optional[List[np.ndarray]] = None, ) -> None: self.records = records self.target_size = target_size self.channels = channels self.augment = augment self.seg_target_size = seg_target_size self.crop_to_disc = crop_to_disc if precomputed_seg_maps is not None: self._seg_maps = precomputed_seg_maps elif unet_segmenter is not None: self._seg_maps = precompute_unet_seg_maps( records, unet_segmenter, unet_threshold ) else: self._seg_maps = None # ------------------------------------------------------------------ def __len__(self) -> int: return len(self.records) # ------------------------------------------------------------------ def _augment(self, seg_map: np.ndarray) -> np.ndarray: """Random flips + 90° rotations (label-safe since NEAREST resize).""" if np.random.rand() < 0.5: seg_map = np.fliplr(seg_map) if np.random.rand() < 0.5: seg_map = np.flipud(seg_map) k = np.random.randint(0, 4) if k: seg_map = np.rot90(seg_map, k=k) return np.ascontiguousarray(seg_map) # ------------------------------------------------------------------ def __getitem__(self, idx: int): rec = self.records[idx] if self._seg_maps is not None: seg_map = self._seg_maps[idx] else: disc_mask, cup_mask = load_gt_masks(rec, self.seg_target_size) seg_map = _combine_masks(disc_mask, cup_mask) if self.crop_to_disc: seg_map = crop_to_disc(seg_map) if self.augment: seg_map = self._augment(seg_map) tensor = seg_map_to_tensor(seg_map, self.channels, self.target_size) return tensor, rec.label # --------------------------------------------------------------------------- # Model # --------------------------------------------------------------------------- class SegCNN(nn.Module): """ Pretrained CNN backbone adapted for segmentation-map input. Parameters ---------- num_classes : output classes (2 for binary glaucoma grading) backbone : "resnet18" | "resnet50" | "efficientnet_b0" pretrained : initialise with ImageNet weights (recommended even for non-RGB input — transfer generalises across domains) in_channels : 1 (single label map) or 3 (one-hot channels) dropout : dropout rate before the final classifier head """ def __init__( self, num_classes: int = 2, backbone: str = "resnet18", pretrained: bool = True, in_channels: int = 3, dropout: float = 0.3, ) -> None: super().__init__() weights_arg = "DEFAULT" if pretrained else None if backbone == "resnet18": base = models.resnet18(weights=weights_arg) feat_dim = base.fc.in_features base.fc = nn.Identity() elif backbone == "resnet50": base = models.resnet50(weights=weights_arg) feat_dim = base.fc.in_features base.fc = nn.Identity() elif backbone == "efficientnet_b0": base = models.efficientnet_b0(weights=weights_arg) feat_dim = base.classifier[1].in_features base.classifier = nn.Identity() else: raise ValueError(f"Unknown backbone: {backbone!r}") # Adapt first conv layer if in_channels ≠ 3 if in_channels != 3: first_conv = self._find_first_conv(base) new_conv = nn.Conv2d( in_channels, first_conv.out_channels, kernel_size=first_conv.kernel_size, stride=first_conv.stride, padding=first_conv.padding, bias=first_conv.bias is not None, ) if pretrained: # Average pretrained RGB weights across channel dim with torch.no_grad(): new_conv.weight.copy_( first_conv.weight.mean(dim=1, keepdim=True).expand_as(new_conv.weight) ) self._replace_first_conv(base, new_conv) self.backbone = base self.head = nn.Sequential( nn.Dropout(p=dropout), nn.Linear(feat_dim, num_classes), ) # ------------------------------------------------------------------ @staticmethod def _find_first_conv(module: nn.Module) -> nn.Conv2d: for m in module.modules(): if isinstance(m, nn.Conv2d): return m raise RuntimeError("No Conv2d found in backbone") @staticmethod def _replace_first_conv(module: nn.Module, new_conv: nn.Conv2d) -> None: """Replace the first Conv2d in-place (handles resnet and efficientnet).""" for name, child in module.named_children(): if isinstance(child, nn.Conv2d): setattr(module, name, new_conv) return try: SegCNN._replace_first_conv(child, new_conv) return except RuntimeError: pass raise RuntimeError("Could not replace first Conv2d") # ------------------------------------------------------------------ def forward(self, x: torch.Tensor) -> torch.Tensor: feats = self.backbone(x) if feats.dim() > 2: feats = feats.flatten(1) return self.head(feats)