src/drdetect/segmentation/dataset.py

"""IDRiD lesion segmentation dataset: on-the-fly patch sampling from full-resolution images.

IDRiD images are 2848x4288 -- about 17x the pixel count of the 1024px grading
images this project trains at elsewhere, and holding even one at full
resolution through a segmentation decoder is not something this project's
measured memory budget (docs/05_PROTOTYPE_SCOPE.md) supports. The roadmap's
own plan is "patch sampling at full resolution" (docs/04_ROADMAP.md, Phase 4),
not downscaling -- a hard exudate here is 0.07% of the image by pixel count
(measured directly, IDRiD_41_EX.tif), so shrinking the image shrinks the
lesion out of existence before the model ever sees it.

Patches are drawn fresh every `__getitem__` call, not pre-extracted once,
biased toward lesion-containing regions (`lesion_patch_prob`) -- a uniformly
random 512x512 crop of a 2848x4288 image with <0.1% positive pixels would
return an all-background patch almost every time, and a model trained on
"predict background" batches learns exactly that.
"""

from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path

import albumentations as A
import cv2
import numpy as np
from albumentations.pytorch import ToTensorV2
from torch.utils.data import Dataset

__all__ = [
    "LesionImagePair",
    "find_idrid_lesion_pairs",
    "IDRiDLesionDataset",
    "build_patch_transforms",
]

LESION_FOLDERS = {
    "microaneurysms": "1. Microaneurysms",
    "haemorrhages": "2. Haemorrhages",
    "hard_exudates": "3. Hard Exudates",
    "soft_exudates": "4. Soft Exudates",
    "optic_disc": "5. Optic Disc",
}
# IDRiD's own mask filename suffixes, per lesion type -- not derivable from
# the folder name alone (e.g. "Hard Exudates" -> "_EX", not "_HardExudates").
LESION_SUFFIXES = {
    "microaneurysms": "_MA",
    "haemorrhages": "_HE",
    "hard_exudates": "_EX",
    "soft_exudates": "_SE",
    "optic_disc": "_OD",
}


@dataclass(frozen=True)
class LesionImagePair:
    image_id: str
    image_path: Path
    mask_path: Path


def find_idrid_lesion_pairs(
    idrid_root: str | Path, lesion: str, split: str
) -> list[LesionImagePair]:
    """Match IDRiD images to their mask for one lesion type and split.

    Args:
        idrid_root: e.g. "data/raw/idrid".
        lesion: a key of LESION_FOLDERS.
        split: "train" or "test".
    """
    if lesion not in LESION_FOLDERS:
        raise ValueError(f"unknown lesion {lesion!r}; expected one of {sorted(LESION_FOLDERS)}")
    root = Path(idrid_root)
    split_dir = "a. Training Set" if split == "train" else "b. Testing Set"
    image_dir = root / "A. Segmentation" / "1. Original Images" / split_dir
    mask_dir = (
        root
        / "A. Segmentation"
        / "2. All Segmentation Groundtruths"
        / split_dir
        / LESION_FOLDERS[lesion]
    )

    pairs = []
    for image_path in sorted(image_dir.glob("*.jpg")):
        mask_path = mask_dir / f"{image_path.stem}{LESION_SUFFIXES[lesion]}.tif"
        if mask_path.exists():
            pairs.append(
                LesionImagePair(
                    image_id=image_path.stem, image_path=image_path, mask_path=mask_path
                )
            )
    return pairs


def build_patch_transforms(patch_size: int, train: bool) -> A.Compose:
    """Same reasoning as drdetect.data.dataset.build_transforms: flips/rotation
    are label-preserving for a fundus with no canonical orientation, and no
    blur/noise, since it would erase small lesions the same way it would
    erase microaneurysms in the grading pipeline."""
    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()])
    return A.Compose(
        [
            A.HorizontalFlip(p=0.5),
            A.VerticalFlip(p=0.5),
            A.RandomRotate90(p=0.5),
            A.RandomBrightnessContrast(brightness_limit=0.15, contrast_limit=0.15, p=0.5),
            normalise,
            ToTensorV2(),
        ]
    )


class IDRiDLesionDataset(Dataset):
    """One lesion type, patch-sampled on the fly.

    `patches_per_image` fixes an epoch's length to something meaningful for a
    54-image training set -- without it, `len(dataset) == len(pairs)` would
    make one epoch see 54 patches total, most of a 2848x4288 image never
    sampled at all.
    """

    def __init__(
        self,
        pairs: list[LesionImagePair],
        *,
        patch_size: int = 512,
        train: bool = True,
        patches_per_image: int = 20,
        lesion_patch_prob: float = 0.8,
        seed: int = 42,
    ):
        self.pairs = pairs
        self.patch_size = patch_size
        self.train = train
        self.patches_per_image = patches_per_image
        self.lesion_patch_prob = lesion_patch_prob
        self.transform = build_patch_transforms(patch_size, train)
        self._rng = np.random.default_rng(seed)

    def __len__(self) -> int:
        return len(self.pairs) * self.patches_per_image

    def _load(self, pair: LesionImagePair) -> tuple[np.ndarray, np.ndarray]:
        image = cv2.imread(str(pair.image_path), cv2.IMREAD_COLOR)
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        mask = cv2.imread(str(pair.mask_path), cv2.IMREAD_GRAYSCALE)
        mask = (mask > 0).astype(np.uint8)
        return image, mask

    def _sample_patch(self, image: np.ndarray, mask: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
        h, w = mask.shape
        p = self.patch_size
        ys, xs = np.where(mask > 0)
        want_lesion_patch = len(ys) > 0 and self._rng.random() < self.lesion_patch_prob
        if want_lesion_patch:
            i = self._rng.integers(0, len(ys))
            cy, cx = int(ys[i]), int(xs[i])
            top = int(np.clip(cy - p // 2 + self._rng.integers(-p // 4, p // 4 + 1), 0, h - p))
            left = int(np.clip(cx - p // 2 + self._rng.integers(-p // 4, p // 4 + 1), 0, w - p))
        else:
            top = int(self._rng.integers(0, max(h - p, 1)))
            left = int(self._rng.integers(0, max(w - p, 1)))
        return image[top : top + p, left : left + p], mask[top : top + p, left : left + p]

    def __getitem__(self, idx: int):
        pair = self.pairs[idx % len(self.pairs)]
        image, mask = self._load(pair)
        patch_img, patch_mask = self._sample_patch(image, mask)
        out = self.transform(image=patch_img, mask=patch_mask)
        return out["image"], out["mask"].float().unsqueeze(0)