"""DRIVE vessel segmentation: dataset, model, and FOV-aware Lightning module. Deliberately separate from dataset.py/module.py (the IDRiD lesion harness), not sharing a base class with them, for reasons that are architectural, not just organisational: * DRIVE images are small and uniform (584x565) -- full images are trained on directly, no patch sampling. * Vessels are a measured ~7.5% of pixels, not IDRiD's <0.1% -- AUROC is a meaningful metric here (see metrics.pixel_auroc), where it is not for lesions. * Every DRIVE image has a field-of-view mask, and pixels outside it (the black surround around the circular retinal image) are not tissue at all -- they must be excluded from both loss and metrics, or a model that trivially gets the black border right looks better than it is. IDRiD's lesion harness has no equivalent concept. DRIVE's own public distribution ships vessel ground truth (`1st_manual/`) only for the 20 "training" images -- the 20 "test" images have no publicly released annotations (confirmed directly: `data/raw/drive/test/` contains only `images/` and `mask/`, no manual-annotation directory). All 20 labelled images are therefore used for k-fold cross-validation; there is no separate locked test split for this lesion type the way IDRiD has one for hard exudates. """ from __future__ import annotations from dataclasses import dataclass from pathlib import Path import albumentations as A import cv2 import lightning as L import numpy as np import torch import torch.nn.functional as F from albumentations.pytorch import ToTensorV2 from torch.utils.data import Dataset from drdetect.segmentation.metrics import dice_coefficient, pixel_auroc __all__ = [ "DriveVesselPair", "find_drive_vessel_pairs", "DriveVesselDataset", "VesselSegmentationModule", ] @dataclass(frozen=True) class DriveVesselPair: image_id: str image_path: Path mask_path: Path fov_path: Path def find_drive_vessel_pairs(drive_root: str | Path) -> list[DriveVesselPair]: """The 20 DRIVE "training" images -- the only ones with public vessel ground truth. Numbered 21-40 in DRIVE's own convention.""" root = Path(drive_root) / "training" image_dir, mask_dir, fov_dir = root / "images", root / "1st_manual", root / "mask" pairs = [] for image_path in sorted(image_dir.glob("*.tif")): image_id = image_path.stem.split("_")[0] mask_path = mask_dir / f"{image_id}_manual1.gif" fov_path = fov_dir / f"{image_id}_training_mask.gif" if mask_path.exists() and fov_path.exists(): pairs.append( DriveVesselPair( image_id=image_id, image_path=image_path, mask_path=mask_path, fov_path=fov_path ) ) return pairs def _pad_to_multiple(arr: np.ndarray, multiple: int = 32, *, reflect: bool) -> np.ndarray: """Pad to a SQUARE size (both axes padded to the same target, from the larger of the two), not just each axis to its own nearest multiple. DRIVE images are rectangular (584x565) and `RandomRotate90` is one of this dataset's augmentations -- a 90 degree turn swaps height and width, so a rectangular pad target produces a *different* output shape for a rotated image than an unrotated one in the same batch, and `torch.stack` cannot collate a batch of mismatched shapes. Padding both axes to one shared square target makes every rotation shape-invariant. """ h, w = arr.shape[:2] target = -(-max(h, w) // multiple) * multiple # round up to the next multiple ph, pw = target - h, target - w border = cv2.BORDER_REFLECT_101 if reflect else cv2.BORDER_CONSTANT return cv2.copyMakeBorder(arr, 0, ph, 0, pw, border, value=0) def build_vessel_transforms(train: bool) -> A.Compose: """Same reasoning as the grading and lesion pipelines: flips/rotation are label-preserving for a fundus with no canonical orientation; no blur, since vessels include capillary-scale structures a blur would erase.""" normalise = A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)) if not train: return A.Compose([normalise, ToTensorV2()], additional_targets={"fov": "mask"}) return A.Compose( [ A.HorizontalFlip(p=0.5), A.VerticalFlip(p=0.5), A.RandomRotate90(p=0.5), normalise, ToTensorV2(), ], additional_targets={"fov": "mask"}, ) class DriveVesselDataset(Dataset): """One full (padded) DRIVE image per item -- no patch sampling, unlike the IDRiD lesion dataset, since these images are already small enough to train on whole.""" def __init__(self, pairs: list[DriveVesselPair], *, train: bool = True): self.pairs = pairs self.transform = build_vessel_transforms(train) def __len__(self) -> int: return len(self.pairs) def __getitem__(self, idx: int): pair = self.pairs[idx] image = cv2.cvtColor(cv2.imread(str(pair.image_path), cv2.IMREAD_COLOR), cv2.COLOR_BGR2RGB) mask = (cv2.imread(str(pair.mask_path), cv2.IMREAD_GRAYSCALE) > 0).astype(np.uint8) fov = (cv2.imread(str(pair.fov_path), cv2.IMREAD_GRAYSCALE) > 0).astype(np.uint8) # Reflect-pad the image (real content), zero-pad mask/FOV (synthetic # border pixels are never vessel and never inside the field of view, # so this makes the padding automatically excluded everywhere the # FOV mask is already used to exclude the real black surround). image = _pad_to_multiple(image, reflect=True) mask = _pad_to_multiple(mask, reflect=False) fov = _pad_to_multiple(fov, reflect=False) out = self.transform(image=image, mask=mask, fov=fov) return out["image"], out["mask"].float().unsqueeze(0), out["fov"].float().unsqueeze(0) class VesselSegmentationModule(L.LightningModule): def __init__( self, model, *, lr: float = 1e-4, weight_decay: float = 1e-4, pos_weight: float = 1.0, dice_weight: float = 1.0, max_epochs: int = 40, ): super().__init__() self.save_hyperparameters(ignore=["model"]) self.model = model self.register_buffer("_pos_weight", torch.tensor(pos_weight)) self._val_probs: list[torch.Tensor] = [] self._val_targets: list[torch.Tensor] = [] self._val_fov: list[torch.Tensor] = [] def forward(self, x): return self.model(x) def _loss(self, logits: torch.Tensor, targets: torch.Tensor, fov: torch.Tensor) -> torch.Tensor: bce_per_pixel = F.binary_cross_entropy_with_logits( logits, targets, pos_weight=self._pos_weight, reduction="none" ) bce = (bce_per_pixel * fov).sum() / fov.sum().clamp(min=1.0) probs = torch.sigmoid(logits) * fov targets_in_fov = targets * fov intersection = (probs * targets_in_fov).sum(dim=(1, 2, 3)) union = probs.sum(dim=(1, 2, 3)) + targets_in_fov.sum(dim=(1, 2, 3)) dice_loss = 1.0 - ((2.0 * intersection + 1.0) / (union + 1.0)).mean() return bce + self.hparams.dice_weight * dice_loss def training_step(self, batch, batch_idx): x, y, fov = batch logits = self(x) loss = self._loss(logits, y, fov) if not torch.isfinite(loss): raise RuntimeError( f"Vessel segmentation loss is {loss.item()} at step {batch_idx}. Diverged." ) self.log("train/loss", loss, prog_bar=True, on_step=False, on_epoch=True) return loss def validation_step(self, batch, batch_idx): x, y, fov = batch logits = self(x) loss = self._loss(logits, y, fov) self.log("val/loss", loss, prog_bar=True, on_step=False, on_epoch=True) self._val_probs.append(torch.sigmoid(logits).detach().float().cpu()) self._val_targets.append(y.detach().cpu()) self._val_fov.append(fov.detach().cpu()) def on_validation_epoch_end(self): if not self._val_probs: return probs = torch.cat(self._val_probs).numpy() targets = torch.cat(self._val_targets).numpy() fov = torch.cat(self._val_fov).numpy().astype(bool) in_fov_probs = probs[fov] in_fov_targets = targets[fov] if in_fov_targets.sum() > 0: auroc = pixel_auroc(in_fov_targets, in_fov_probs) self.log("val/auroc", auroc, prog_bar=True) dice = dice_coefficient(in_fov_targets, in_fov_probs > 0.5) self.log("val/dice", dice, prog_bar=True) self._val_probs.clear() self._val_targets.clear() self._val_fov.clear() def configure_optimizers(self): optimiser = torch.optim.AdamW( self.model.parameters(), lr=self.hparams.lr, weight_decay=self.hparams.weight_decay ) scheduler = torch.optim.lr_scheduler.CosineAnnealingLR( optimiser, T_max=self.hparams.max_epochs ) return { "optimizer": optimiser, "lr_scheduler": {"scheduler": scheduler, "interval": "epoch"}, }