307 lines
11 KiB
Python
Executable File
307 lines
11 KiB
Python
Executable File
"""Utilities for preparing REFUGE (REFUGE1/REFUGE2) datasets.
|
||
|
||
Builds a unified manifest across all provided splits (REFUGE1 train/val/test
|
||
and REFUGE2 validation/test), exposing image paths, glaucoma labels, disc/cup
|
||
masks, and fovea coordinates so downstream segmentation/classification modules
|
||
can operate without additional bookkeeping.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
from dataclasses import dataclass
|
||
from pathlib import Path
|
||
from typing import Dict, Iterable, List, Optional, Tuple
|
||
|
||
import pandas as pd
|
||
|
||
|
||
@dataclass
|
||
class RefugeSample:
|
||
"""Lightweight container describing a REFUGE sample."""
|
||
|
||
sample_id: str
|
||
dataset: str
|
||
split: str
|
||
image_path: Path
|
||
label: Optional[int]
|
||
device: Optional[str]
|
||
mask_path: Optional[Path]
|
||
fovea_coord: Optional[Tuple[float, float]]
|
||
|
||
|
||
class RefugePreprocessing:
|
||
"""Builds manifests and provides shared helpers for REFUGE workflows.
|
||
|
||
Responsibilities:
|
||
* scan the REFUGE directory structure and build a consistent manifest
|
||
(train/val/test, device vendor, ground-truth labels)
|
||
* expose convenience loaders for raw RGB frames, OD/OC masks, and
|
||
optional fovea landmarks
|
||
* compute geometric metadata (disc centres, diameters) so downstream
|
||
stages can crop ROIs lazily instead of storing pre-rendered tiles
|
||
"""
|
||
|
||
def __init__(self, root_dir: Path | str) -> None:
|
||
self.root_dir = Path(root_dir)
|
||
self._manifest = None # populated by build_manifest()
|
||
|
||
# ------------------------------------------------------------------
|
||
# Manifest handling
|
||
# ------------------------------------------------------------------
|
||
def build_manifest(self, refresh: bool = False) -> Iterable[RefugeSample]:
|
||
"""Return an iterable of :class:`RefugeSample` records.
|
||
|
||
Parameters
|
||
----------
|
||
refresh:
|
||
when True, force a rescan of the filesystem instead of reusing the
|
||
cached manifest.
|
||
|
||
Returns
|
||
-------
|
||
Iterable[RefugeSample]
|
||
A sequence containing one entry per sample in the REFUGE datasets.
|
||
|
||
Notes
|
||
-----
|
||
The actual manifest-building logic will live here: parsing the
|
||
directory structure, reading any provided CSV/Excel metadata, and
|
||
aligning masks/labels. For now, this method raises ``NotImplementedError``
|
||
so callers are reminded to hook it up before use.
|
||
"""
|
||
|
||
if self._manifest is not None and not refresh:
|
||
return self._manifest
|
||
|
||
manifest: List[RefugeSample] = []
|
||
|
||
manifest.extend(self._collect_refuge1_train())
|
||
manifest.extend(self._collect_refuge1_val())
|
||
manifest.extend(self._collect_refuge1_test())
|
||
manifest.extend(self._collect_refuge2_val())
|
||
manifest.extend(self._collect_refuge2_test())
|
||
|
||
self._manifest = manifest
|
||
return self._manifest
|
||
|
||
# ------------------------------------------------------------------
|
||
# Accessors for downstream modules
|
||
# ------------------------------------------------------------------
|
||
def load_image(self, sample: RefugeSample):
|
||
"""Return the RGB fundus image for ``sample``.
|
||
|
||
Implementors should handle color-space consistency (e.g., ensure RGB vs
|
||
BGR) and any global normalisation desired across devices.
|
||
"""
|
||
|
||
raise NotImplementedError("Image loading to be implemented")
|
||
|
||
def load_mask(self, sample: RefugeSample):
|
||
"""Return the optic disc/cup mask for ``sample`` if available."""
|
||
|
||
raise NotImplementedError("Mask loading to be implemented")
|
||
|
||
def disc_geometry(self, sample: RefugeSample) -> Dict[str, float]:
|
||
"""Compute disc centre and diameter from the mask.
|
||
|
||
The segmentation module will rely on this to crop 2.5–3× disc-diameter
|
||
ROIs at training time.
|
||
"""
|
||
|
||
raise NotImplementedError("Disc geometry helper to be implemented")
|
||
|
||
# ------------------------------------------------------------------
|
||
# Internal helpers
|
||
# ------------------------------------------------------------------
|
||
def _collect_refuge1_train(self) -> List[RefugeSample]:
|
||
base = self.root_dir / "Train" / "REFUGE1-train"
|
||
if not base.exists():
|
||
return []
|
||
|
||
fovea_path = base / "Fovea_location.xlsx"
|
||
fovea_map = self._read_fovea_table(fovea_path, img_col="ImgName")
|
||
|
||
samples: List[RefugeSample] = []
|
||
image_root = base / "Training400"
|
||
mask_root = base / "Disc_Cup_Masks"
|
||
|
||
for label_name, label_val in ("Glaucoma", 1), ("Non-Glaucoma", 0):
|
||
img_dir = image_root / label_name
|
||
mask_dir = mask_root / label_name
|
||
if not img_dir.exists():
|
||
continue
|
||
for image_path in sorted(img_dir.glob("*.jpg")):
|
||
img_name = image_path.name
|
||
mask_path = (mask_dir / image_path.with_suffix(".bmp").name)
|
||
fovea = fovea_map.get(img_name)
|
||
sample_id = f"refuge1_train_{image_path.stem}"
|
||
samples.append(
|
||
RefugeSample(
|
||
sample_id=sample_id,
|
||
dataset="refuge1",
|
||
split="train",
|
||
image_path=image_path,
|
||
label=label_val,
|
||
device=None,
|
||
mask_path=mask_path if mask_path.exists() else None,
|
||
fovea_coord=fovea,
|
||
)
|
||
)
|
||
return samples
|
||
|
||
def _collect_refuge1_val(self) -> List[RefugeSample]:
|
||
base = self.root_dir / "Train" / "REFUGE1-val"
|
||
if not base.exists():
|
||
return []
|
||
|
||
fovea_path = base / "Fovea_locations.xlsx"
|
||
df = pd.read_excel(fovea_path)
|
||
samples: List[RefugeSample] = []
|
||
image_root = base / "REFUGE-Validation400"
|
||
mask_root = base / "Disc_Cup_Masks"
|
||
|
||
for _, row in df.iterrows():
|
||
img_name = row["ImgName"]
|
||
image_path = image_root / img_name
|
||
mask_path = mask_root / Path(img_name).with_suffix(".bmp").name
|
||
fovea = self._extract_fovea(row, x_key="Fovea_X", y_key="Fovea_Y")
|
||
label = int(row.get("Glaucoma Label", 0)) if not pd.isna(row.get("Glaucoma Label", 0)) else None
|
||
sample_id = f"refuge1_val_{Path(img_name).stem}"
|
||
samples.append(
|
||
RefugeSample(
|
||
sample_id=sample_id,
|
||
dataset="refuge1",
|
||
split="val",
|
||
image_path=image_path,
|
||
label=label,
|
||
device=None,
|
||
mask_path=mask_path if mask_path.exists() else None,
|
||
fovea_coord=fovea,
|
||
)
|
||
)
|
||
return samples
|
||
|
||
def _collect_refuge1_test(self) -> List[RefugeSample]:
|
||
base = self.root_dir / "Train" / "REFUGE1-test"
|
||
if not base.exists():
|
||
return []
|
||
|
||
df = pd.read_excel(base / "Glaucoma_label_and_Fovea_location.xlsx")
|
||
image_root = base / "Test400"
|
||
mask_root = base / "Disc_Cup_Masks"
|
||
samples: List[RefugeSample] = []
|
||
|
||
for _, row in df.iterrows():
|
||
img_name = row["ImgName"]
|
||
image_path = image_root / img_name
|
||
mask_path = mask_root / Path(img_name).with_suffix(".bmp").name
|
||
fovea = self._extract_fovea(row, x_key="Fovea_X", y_key="Fovea_Y")
|
||
label = int(row.get("Label(Glaucoma=1)", 0)) if not pd.isna(row.get("Label(Glaucoma=1)", 0)) else None
|
||
sample_id = f"refuge1_test_{Path(img_name).stem}"
|
||
samples.append(
|
||
RefugeSample(
|
||
sample_id=sample_id,
|
||
dataset="refuge1",
|
||
split="test",
|
||
image_path=image_path,
|
||
label=label,
|
||
device=None,
|
||
mask_path=mask_path if mask_path.exists() else None,
|
||
fovea_coord=fovea,
|
||
)
|
||
)
|
||
return samples
|
||
|
||
def _collect_refuge2_val(self) -> List[RefugeSample]:
|
||
base = self.root_dir / "Validation"
|
||
if not base.exists():
|
||
return []
|
||
|
||
label_df = pd.read_csv(base / "glaucoma.csv")
|
||
fovea_df = pd.read_csv(base / "fovea.csv")
|
||
fovea_map = {
|
||
row["ImageName"]: (float(row["Fovea_X"]), float(row["Fovea_Y"]))
|
||
for _, row in fovea_df.iterrows()
|
||
}
|
||
samples: List[RefugeSample] = []
|
||
image_root = base / "Images"
|
||
mask_root = base / "Disc_Masks"
|
||
|
||
for _, row in label_df.iterrows():
|
||
img_name = row["FileName"]
|
||
image_path = image_root / img_name
|
||
mask_path = mask_root / Path(img_name).with_suffix(".png").name
|
||
label = row.get("Glaucoma Risk")
|
||
label = int(label) if label == label else None
|
||
sample_id = f"refuge2_val_{Path(img_name).stem}"
|
||
samples.append(
|
||
RefugeSample(
|
||
sample_id=sample_id,
|
||
dataset="refuge2",
|
||
split="val",
|
||
image_path=image_path,
|
||
label=label,
|
||
device=None,
|
||
mask_path=mask_path if mask_path.exists() else None,
|
||
fovea_coord=fovea_map.get(img_name),
|
||
)
|
||
)
|
||
return samples
|
||
|
||
def _collect_refuge2_test(self) -> List[RefugeSample]:
|
||
base = self.root_dir / "Test"
|
||
if not base.exists():
|
||
return []
|
||
|
||
label_df = pd.read_excel(base / "task1.xls", header=None, names=["ImgName", "Glaucoma"])
|
||
fovea_df = pd.read_excel(base / "fovea.xlsx")
|
||
fovea_map = {
|
||
row["ImageName"]: (float(row["Fovea_X"]), float(row["Fovea_Y"]))
|
||
for _, row in fovea_df.iterrows()
|
||
}
|
||
samples: List[RefugeSample] = []
|
||
image_root = base / "refuge2-test"
|
||
mask_root = base / "Disc_Mask"
|
||
|
||
for _, row in label_df.iterrows():
|
||
img_name = row["ImgName"]
|
||
image_path = image_root / img_name
|
||
mask_path = mask_root / Path(img_name).with_suffix(".png").name
|
||
label = row.get("Glaucoma")
|
||
label = int(label) if label == label else None
|
||
sample_id = f"refuge2_test_{Path(img_name).stem}"
|
||
samples.append(
|
||
RefugeSample(
|
||
sample_id=sample_id,
|
||
dataset="refuge2",
|
||
split="test",
|
||
image_path=image_path,
|
||
label=label,
|
||
device=None,
|
||
mask_path=mask_path if mask_path.exists() else None,
|
||
fovea_coord=fovea_map.get(img_name),
|
||
)
|
||
)
|
||
return samples
|
||
|
||
@staticmethod
|
||
def _read_fovea_table(path: Path, img_col: str) -> Dict[str, Tuple[float, float]]:
|
||
if not path.exists():
|
||
return {}
|
||
df = pd.read_excel(path)
|
||
mapping: Dict[str, Tuple[float, float]] = {}
|
||
for _, row in df.iterrows():
|
||
mapping[row[img_col]] = (
|
||
float(row.get("Fovea_X", float("nan"))),
|
||
float(row.get("Fovea_Y", float("nan"))),
|
||
)
|
||
return mapping
|
||
|
||
@staticmethod
|
||
def _extract_fovea(row: pd.Series, x_key: str, y_key: str) -> Optional[Tuple[float, float]]:
|
||
x_val = row.get(x_key)
|
||
y_val = row.get(y_key)
|
||
if pd.isna(x_val) or pd.isna(y_val):
|
||
return None
|
||
return float(x_val), float(y_val)
|