rpotter6298
139 lines
4.4 KiB
Python
139 lines
4.4 KiB
Python
"""Build manifest for U-Net segmenter combining REFUGE and Papila annotations."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import random
|
|
from pathlib import Path
|
|
from typing import Optional
|
|
|
|
import sys
|
|
|
|
import pandas as pd
|
|
|
|
ROOT = Path(__file__).resolve().parents[1]
|
|
sys.path.append(str(ROOT))
|
|
|
|
from classes.refuge_preprocessing import RefugePreprocessing
|
|
|
|
REFUGE_ROOT = Path("REFUGE")
|
|
PAPILA_IMAGES = Path("FundusImages")
|
|
PAPILA_CONTOURS = Path("Papila/ExpertsSegmentations/Contours")
|
|
DEFAULT_OUTPUT = Path("Papila/analysis_data/unet_manifest.csv")
|
|
|
|
|
|
def pick_contour(base: str, kind: str) -> Optional[Path]:
|
|
"""Return contour path for Papila image (disc/cup)."""
|
|
candidates = [
|
|
PAPILA_CONTOURS / f"{base}_{kind}_exp2.txt",
|
|
PAPILA_CONTOURS / f"{base}_{kind}_exp1.txt",
|
|
]
|
|
for path in candidates:
|
|
if path.exists():
|
|
return path
|
|
return None
|
|
|
|
|
|
def collect_refuge() -> pd.DataFrame:
|
|
pre = RefugePreprocessing(REFUGE_ROOT)
|
|
samples = []
|
|
for sample in pre.build_manifest(refresh=True):
|
|
if sample.mask_path is None:
|
|
continue
|
|
split = sample.split
|
|
if split == "test":
|
|
split = "holdout"
|
|
samples.append(
|
|
{
|
|
"sample_id": sample.sample_id,
|
|
"dataset": "refuge",
|
|
"image_path": sample.image_path.resolve(),
|
|
"annotation_disc": sample.mask_path.resolve(),
|
|
"annotation_cup": sample.mask_path.resolve(),
|
|
"annotation_type_disc": "mask",
|
|
"annotation_type_cup": "mask",
|
|
"split": split,
|
|
}
|
|
)
|
|
return pd.DataFrame(samples)
|
|
|
|
|
|
def collect_papila() -> pd.DataFrame:
|
|
samples = []
|
|
if not PAPILA_IMAGES.exists():
|
|
return pd.DataFrame(samples)
|
|
for img_path in sorted(PAPILA_IMAGES.glob("RET*")):
|
|
base = img_path.stem
|
|
disc = pick_contour(base, "disc")
|
|
cup = pick_contour(base, "cup")
|
|
if disc is None or cup is None:
|
|
continue
|
|
samples.append(
|
|
{
|
|
"sample_id": f"papila_{base}",
|
|
"dataset": "papila",
|
|
"image_path": img_path.resolve(),
|
|
"annotation_disc": disc.resolve(),
|
|
"annotation_cup": cup.resolve(),
|
|
"annotation_type_disc": "contour",
|
|
"annotation_type_cup": "contour",
|
|
}
|
|
)
|
|
return pd.DataFrame(samples)
|
|
|
|
|
|
def assign_splits(df: pd.DataFrame, holdout_ratio: float, seed: int) -> pd.DataFrame:
|
|
rng = random.Random(seed)
|
|
df = df.copy()
|
|
if "split" not in df.columns:
|
|
df["split"] = None
|
|
for dataset, group in df.groupby("dataset"):
|
|
indices = list(group.index)
|
|
|
|
# Preserve provided splits (e.g., REFUGE train/val/test); only populate
|
|
# missing entries with "train" so downstream code has a default.
|
|
split_series = df.loc[indices, "split"]
|
|
missing = split_series.isna() | (split_series.astype(str).str.strip() == "")
|
|
if missing.any():
|
|
df.loc[missing[missing].index, "split"] = "train"
|
|
split_series = df.loc[indices, "split"]
|
|
|
|
if dataset != "papila":
|
|
continue
|
|
|
|
if holdout_ratio <= 0:
|
|
continue
|
|
|
|
desired_holdout = max(1, int(len(indices) * holdout_ratio))
|
|
split_series = df.loc[indices, "split"]
|
|
current_holdout_mask = split_series == "holdout"
|
|
current_holdout = int(current_holdout_mask.sum())
|
|
remaining = desired_holdout - current_holdout
|
|
if remaining <= 0:
|
|
continue
|
|
|
|
candidate_indices = list(split_series[split_series == "train"].index)
|
|
rng.shuffle(candidate_indices)
|
|
selected = candidate_indices[:remaining]
|
|
df.loc[selected, "split"] = "holdout"
|
|
return df
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(description="Build U-Net manifest")
|
|
parser.add_argument("--holdout", type=float, default=0.05)
|
|
parser.add_argument("--seed", type=int, default=42)
|
|
parser.add_argument("--output", type=Path, default=DEFAULT_OUTPUT)
|
|
args = parser.parse_args()
|
|
|
|
refuge_df = collect_refuge()
|
|
papila_df = collect_papila()
|
|
combined = pd.concat([refuge_df, papila_df], ignore_index=True)
|
|
combined = assign_splits(combined, holdout_ratio=args.holdout, seed=args.seed)
|
|
combined.to_csv(args.output, index=False)
|
|
print(f"Manifest saved to {args.output} with {len(combined)} entries")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|