Spaces:
Running on Zero
Running on Zero
| """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", | |
| ] | |
| 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"}, | |
| } | |