moved_repo_first_update
This commit is contained in:
rpotter6298
Executable
+140
@@ -0,0 +1,140 @@
|
||||
"""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 pandas as pd
|
||||
|
||||
import sys
|
||||
|
||||
REPO_ROOT = Path(__file__).resolve().parents[3]
|
||||
if str(REPO_ROOT) not in sys.path:
|
||||
sys.path.insert(0, str(REPO_ROOT))
|
||||
|
||||
from classes.refuge_preprocessing import RefugePreprocessing
|
||||
|
||||
REFUGE_ROOT = Path("REFUGE")
|
||||
PAPILA_IMAGES = Path("Papila/FundusImages")
|
||||
PAPILA_CONTOURS = Path("Papila/ExpertsSegmentations/Contours")
|
||||
DEFAULT_OUTPUT = Path("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)
|
||||
args.output.parent.mkdir(parents=True, exist_ok=True)
|
||||
combined.to_csv(args.output, index=False)
|
||||
print(f"Manifest saved to {args.output} with {len(combined)} entries")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,138 @@
|
||||
"""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()
|
||||
@@ -0,0 +1,976 @@
|
||||
"""REFUGE training/evaluation helper.
|
||||
|
||||
Usage examples (after activating .venv_refuge):
|
||||
|
||||
python refuge_build.py --train-seg
|
||||
python refuge_build.py --train-clf
|
||||
python refuge_build.py --eval --with-ttt
|
||||
|
||||
The script expects the REFUGE folder and writes checkpoints under
|
||||
models/refuge/segmentation and models/refuge/classifier.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import csv
|
||||
from pathlib import Path
|
||||
from typing import Dict, List, Optional, Sequence, Set, Tuple
|
||||
import shutil
|
||||
import sys
|
||||
|
||||
import torch
|
||||
import numpy as np
|
||||
from PIL import Image, ImageDraw
|
||||
from torch.utils.data import DataLoader
|
||||
from sklearn.metrics import roc_auc_score
|
||||
from tqdm import tqdm
|
||||
from torch import nn
|
||||
from torchvision import models
|
||||
|
||||
REPO_ROOT = Path(__file__).resolve().parents[3]
|
||||
if str(REPO_ROOT) not in sys.path:
|
||||
sys.path.insert(0, str(REPO_ROOT))
|
||||
|
||||
from classes.refuge_preprocessing import RefugePreprocessing, RefugeSample
|
||||
from classes.refuge_segmentation import RefugeSegmentation
|
||||
from classes.refuge_classification import (
|
||||
RefugeClassification,
|
||||
RefugeClassificationRecord,
|
||||
RefugeClassificationDataset,
|
||||
_default_image_transform,
|
||||
_geometry_from_mask,
|
||||
UNetGeometryProvider,
|
||||
)
|
||||
from classes.unet_segmenter import UNetSegmenter
|
||||
from classes.papila_builders import build_papila_clinical
|
||||
|
||||
REFUGE_ROOT = Path("REFUGE")
|
||||
SEG_CKPT = Path("models/refuge/segmentation/refuge_segmentation_best.pt")
|
||||
CLF_DIR = Path("models/refuge/classifier")
|
||||
UNET_WEIGHT_CANDIDATES = (
|
||||
Path("models/v2/refuge/segmentation/per_image/best.pt"),
|
||||
Path("models/v2/refuge/segmentation/best.pt"),
|
||||
Path("models/unet_segmenter/best.pt"),
|
||||
)
|
||||
|
||||
CLASSIFIER_BACKBONES = {
|
||||
"resnet50": models.ResNet50_Weights.DEFAULT,
|
||||
"densenet121": models.DenseNet121_Weights.DEFAULT,
|
||||
"efficientnet_b0": models.EfficientNet_B0_Weights.DEFAULT,
|
||||
"efficientnet_b7": models.EfficientNet_B7_Weights.DEFAULT,
|
||||
}
|
||||
|
||||
|
||||
def build_classifier_backbone(name: str) -> nn.Module:
|
||||
name = name.lower()
|
||||
if name not in CLASSIFIER_BACKBONES:
|
||||
raise ValueError(f"Unsupported classifier backbone '{name}'")
|
||||
|
||||
weights = CLASSIFIER_BACKBONES[name]
|
||||
|
||||
if name == "resnet50":
|
||||
model = models.resnet50(weights=weights)
|
||||
feat_dim = model.fc.in_features
|
||||
model.fc = nn.Identity()
|
||||
elif name == "densenet121":
|
||||
model = models.densenet121(weights=weights)
|
||||
feat_dim = model.classifier.in_features
|
||||
model.classifier = nn.Identity()
|
||||
elif name == "efficientnet_b0":
|
||||
model = models.efficientnet_b0(weights=weights)
|
||||
feat_dim = model.classifier[-1].in_features # type: ignore[index]
|
||||
model.classifier = nn.Identity()
|
||||
elif name == "efficientnet_b7":
|
||||
model = models.efficientnet_b7(weights=weights)
|
||||
feat_dim = model.classifier[-1].in_features # type: ignore[index]
|
||||
model.classifier = nn.Identity()
|
||||
else: # pragma: no cover
|
||||
raise ValueError(f"Unsupported classifier backbone '{name}'")
|
||||
|
||||
setattr(model, "_feature_dim", int(feat_dim))
|
||||
return model
|
||||
|
||||
|
||||
def classifier_checkpoint_dir(backbone_name: str) -> Path:
|
||||
return CLF_DIR / backbone_name
|
||||
|
||||
|
||||
def classifier_checkpoint_path(backbone_name: str) -> Path:
|
||||
return classifier_checkpoint_dir(backbone_name) / "refuge_classifier_best.pt"
|
||||
|
||||
|
||||
def resolve_unet_weights(explicit: Optional[Path]) -> Path:
|
||||
if explicit is not None:
|
||||
return explicit
|
||||
for cand in UNET_WEIGHT_CANDIDATES:
|
||||
if cand.exists():
|
||||
return cand
|
||||
return UNET_WEIGHT_CANDIDATES[0]
|
||||
|
||||
|
||||
def ensure_preprocessing() -> RefugePreprocessing:
|
||||
if not REFUGE_ROOT.exists():
|
||||
raise FileNotFoundError(f"REFUGE directory not found at {REFUGE_ROOT}")
|
||||
return RefugePreprocessing(REFUGE_ROOT)
|
||||
|
||||
|
||||
def load_allowed_ids(
|
||||
csv_path: Optional[Path], dice_threshold: float
|
||||
) -> Optional[Set[str]]:
|
||||
if csv_path is None or not csv_path.exists():
|
||||
return None
|
||||
allowed: Set[str] = set()
|
||||
with csv_path.open(newline="") as fh:
|
||||
reader = csv.DictReader(fh)
|
||||
for row in reader:
|
||||
sample_id = row.get("sample_id")
|
||||
if not sample_id or sample_id == "__mean__":
|
||||
continue
|
||||
try:
|
||||
disc = float(row.get("dice_disc", "nan"))
|
||||
cup = float(row.get("dice_cup", "nan"))
|
||||
except (TypeError, ValueError):
|
||||
continue
|
||||
if disc < dice_threshold and cup < dice_threshold:
|
||||
continue
|
||||
allowed.add(sample_id)
|
||||
return allowed
|
||||
|
||||
|
||||
def build_papila_samples(
|
||||
image_dir: Path,
|
||||
clinical_dir: Path,
|
||||
label_col: str,
|
||||
positive_labels: Sequence[str],
|
||||
allowed_ids: Optional[Set[str]],
|
||||
) -> List[RefugeSample]:
|
||||
clinical = build_papila_clinical(
|
||||
image_dir=str(image_dir),
|
||||
clinical_dir=str(clinical_dir),
|
||||
label_col=label_col,
|
||||
cat_cols=[],
|
||||
)
|
||||
positives = {lbl.lower() for lbl in positive_labels}
|
||||
samples: Dict[str, RefugeSample] = {}
|
||||
for _, row in clinical.df.iterrows():
|
||||
image_path = clinical.get_image_path(row)
|
||||
sample_id = f"papila_{Path(image_path).stem}"
|
||||
if allowed_ids is not None and sample_id not in allowed_ids:
|
||||
continue
|
||||
if sample_id in samples:
|
||||
continue
|
||||
value = row.get(label_col)
|
||||
if value is None or (isinstance(value, float) and np.isnan(value)):
|
||||
continue
|
||||
try:
|
||||
label_int = int(value)
|
||||
if label_int == 2:
|
||||
continue
|
||||
label = 1 if label_int > 0 else 0
|
||||
except (TypeError, ValueError):
|
||||
label = 1 if str(value).strip().lower() in positives else 0
|
||||
samples[sample_id] = RefugeSample(
|
||||
sample_id=sample_id,
|
||||
dataset="papila",
|
||||
split="holdout",
|
||||
image_path=Path(image_path),
|
||||
label=label,
|
||||
device=None,
|
||||
mask_path=None,
|
||||
fovea_coord=None,
|
||||
)
|
||||
return list(samples.values())
|
||||
|
||||
|
||||
def load_contour(path: Path) -> np.ndarray:
|
||||
coords = np.loadtxt(path)
|
||||
if coords.ndim == 1:
|
||||
coords = coords.reshape(-1, 2)
|
||||
return coords
|
||||
|
||||
|
||||
def contour_to_mask(coords: np.ndarray, size: Tuple[int, int]) -> np.ndarray:
|
||||
if coords is None or coords.size == 0:
|
||||
return np.zeros((size[1], size[0]), dtype=np.uint8)
|
||||
img = Image.new("L", size, 0)
|
||||
draw = ImageDraw.Draw(img)
|
||||
points = [tuple(map(float, pt)) for pt in coords]
|
||||
draw.polygon(points, outline=1, fill=1)
|
||||
return np.array(img, dtype=np.uint8)
|
||||
|
||||
|
||||
class PapilaGTGeometryProvider:
|
||||
def __init__(self, contours_dir: Path) -> None:
|
||||
self.contours_dir = contours_dir
|
||||
|
||||
def _pick(self, base: str, kind: str) -> Optional[Path]:
|
||||
for exp in ("exp2", "exp1"):
|
||||
cand = self.contours_dir / f"{base}_{kind}_{exp}.txt"
|
||||
if cand.exists():
|
||||
return cand
|
||||
return None
|
||||
|
||||
def __call__(self, sample: RefugeSample, scale: float):
|
||||
base = Path(sample.image_path).stem
|
||||
disc_path = self._pick(base, "disc")
|
||||
cup_path = self._pick(base, "cup")
|
||||
if disc_path is None or cup_path is None:
|
||||
raise RuntimeError(f"Missing ground-truth contours for {sample.sample_id}")
|
||||
|
||||
image = Image.open(sample.image_path).convert("RGB")
|
||||
disc_coords = load_contour(disc_path)
|
||||
cup_coords = load_contour(cup_path)
|
||||
disc_mask = contour_to_mask(disc_coords, image.size)
|
||||
cup_mask = contour_to_mask(cup_coords, image.size)
|
||||
cup_mask = ((cup_mask > 0) & (disc_mask > 0)).astype(np.uint8)
|
||||
geom = _geometry_from_mask(disc_mask, scale)
|
||||
return geom, disc_mask.astype(np.uint8), cup_mask.astype(np.uint8)
|
||||
|
||||
|
||||
def build_papila_records(
|
||||
args: argparse.Namespace,
|
||||
pre: RefugePreprocessing,
|
||||
checkpoint_path: Path,
|
||||
) -> Tuple[List[RefugeClassificationRecord], Optional[RefugeClassification]]:
|
||||
allowed = load_allowed_ids(
|
||||
getattr(args, "papila_metrics", None),
|
||||
getattr(args, "papila_dice_threshold", 0.01),
|
||||
)
|
||||
samples = build_papila_samples(
|
||||
args.papila_image_dir,
|
||||
args.papila_clinical_dir,
|
||||
args.papila_label_col,
|
||||
args.papila_positive_labels,
|
||||
allowed,
|
||||
)
|
||||
if not samples:
|
||||
return [], None
|
||||
|
||||
cache_dir = args.clf_cache_dir
|
||||
if cache_dir is not None and getattr(args, "papila_use_gt", False):
|
||||
cache_dir = cache_dir / "gt"
|
||||
|
||||
if getattr(args, "papila_use_gt", False):
|
||||
geometry_fn = PapilaGTGeometryProvider(args.papila_contours_dir)
|
||||
provider = geometry_fn
|
||||
else:
|
||||
seg_manifest = getattr(args, "seg_manifest", None)
|
||||
seg_weights = resolve_unet_weights(getattr(args, "seg_weights", None))
|
||||
if seg_manifest is None or seg_weights is None:
|
||||
raise SystemExit(
|
||||
"Papila evaluation without GT masks requires --seg-manifest and --seg-weights"
|
||||
)
|
||||
segmenter = UNetSegmenter(
|
||||
manifest_path=seg_manifest,
|
||||
device=args.device,
|
||||
normalize=args.seg_normalize,
|
||||
)
|
||||
seg_state = torch.load(seg_weights, map_location=args.device)
|
||||
seg_state_dict = seg_state.get("model", seg_state)
|
||||
segmenter.model.load_state_dict(seg_state_dict)
|
||||
segmenter.model.to(args.device)
|
||||
provider = UNetGeometryProvider(
|
||||
segmenter=segmenter,
|
||||
threshold=args.segmenter_threshold,
|
||||
tta=args.segmenter_tta,
|
||||
)
|
||||
geometry_fn = provider
|
||||
|
||||
papila_seg = RefugeSegmentation(pre)
|
||||
backbone = build_classifier_backbone(args.clf_backbone)
|
||||
papila_clf = RefugeClassification(
|
||||
pre,
|
||||
papila_seg,
|
||||
backbone=backbone,
|
||||
geometry_fn=provider,
|
||||
cache_dir=cache_dir,
|
||||
)
|
||||
papila_clf.crop_scale = args.crop_scale
|
||||
papila_clf.crop_size = args.crop_size
|
||||
papila_clf.eval_transform = _default_image_transform(args.crop_size)
|
||||
papila_clf.ttt_transform = papila_clf.eval_transform
|
||||
papila_state = torch.load(checkpoint_path, map_location=args.device)
|
||||
papila_clf.backbone.load_state_dict(papila_state["backbone"])
|
||||
papila_clf.classifier_head.load_state_dict(papila_state["classifier"])
|
||||
papila_clf.rotation_head.load_state_dict(papila_state["rotation"])
|
||||
papila_clf.backbone.to(args.device)
|
||||
papila_clf.classifier_head.to(args.device)
|
||||
papila_clf.rotation_head.to(args.device)
|
||||
|
||||
records = papila_clf.build_records_for_samples(
|
||||
samples, crop_scale=args.crop_scale, progress_prefix="papila"
|
||||
)
|
||||
print(f"[eval] Prepared {len(records)} PAPILA records")
|
||||
return records, papila_clf
|
||||
|
||||
|
||||
def train_segmentation(args: argparse.Namespace) -> None:
|
||||
pre = ensure_preprocessing()
|
||||
seg = RefugeSegmentation(pre)
|
||||
seg.build_datasets(
|
||||
image_size=args.seg_image_size,
|
||||
batch_size=args.seg_batch_size,
|
||||
num_workers=args.num_workers,
|
||||
)
|
||||
history = seg.train(
|
||||
epochs=args.seg_epochs,
|
||||
lr=args.seg_lr,
|
||||
weight_decay=args.seg_weight_decay,
|
||||
checkpoint_dir=SEG_CKPT.parent,
|
||||
device=args.device,
|
||||
)
|
||||
print("Segmentation training complete. Best Dice:", history.get("best_dice"))
|
||||
|
||||
|
||||
def train_unet_segmenter(args: argparse.Namespace) -> None:
|
||||
manifest_path = args.seg_manifest or Path("manifest.csv")
|
||||
mask_cache_dir = None if args.in_memory_cache else args.mask_cache_dir
|
||||
image_cache_dir = None if args.in_memory_cache else args.image_cache_dir
|
||||
if args.in_memory_cache and (args.mask_cache_dir or args.image_cache_dir):
|
||||
print("[unet-seg] in_memory_cache enabled: disk caches disabled for this run.")
|
||||
|
||||
segmenter = UNetSegmenter(
|
||||
manifest_path=manifest_path,
|
||||
device=args.device,
|
||||
target_size=args.seg_image_size,
|
||||
normalize=args.seg_normalize,
|
||||
use_stronger_aug=args.seg_strong_aug,
|
||||
train_datasets=args.seg_train_datasets,
|
||||
val_datasets=args.seg_val_datasets,
|
||||
holdout_datasets=args.seg_holdout_datasets,
|
||||
mask_cache_dir=mask_cache_dir,
|
||||
image_cache_dir=image_cache_dir,
|
||||
in_memory_cache=args.in_memory_cache,
|
||||
loader_workers=args.loader_workers,
|
||||
)
|
||||
if mask_cache_dir:
|
||||
print(f"[unet-seg] mask_cache_dir={mask_cache_dir}")
|
||||
if image_cache_dir:
|
||||
print(f"[unet-seg] image_cache_dir={image_cache_dir}")
|
||||
if args.in_memory_cache:
|
||||
print("[unet-seg] prebuilding in-memory cache")
|
||||
segmenter.prebuild_in_memory_cache(
|
||||
cache_workers=max(0, int(args.cache_workers)),
|
||||
include_train=True,
|
||||
include_val=True,
|
||||
include_holdout=False,
|
||||
)
|
||||
|
||||
segmenter.train(
|
||||
epochs=args.seg_epochs,
|
||||
batch_size=args.seg_batch_size,
|
||||
lr=args.seg_lr,
|
||||
weight_decay=args.seg_weight_decay,
|
||||
checkpoint_dir=args.seg_checkpoint_dir,
|
||||
)
|
||||
print(
|
||||
"[unet-seg] Training complete. Best checkpoint stored at",
|
||||
(args.seg_checkpoint_dir / "best.pt").resolve(),
|
||||
)
|
||||
|
||||
|
||||
def _load_segmentation(
|
||||
pre: RefugePreprocessing, args: argparse.Namespace
|
||||
) -> RefugeSegmentation:
|
||||
seg = RefugeSegmentation(pre)
|
||||
seg.build_datasets(
|
||||
image_size=args.seg_image_size,
|
||||
batch_size=args.seg_batch_size,
|
||||
num_workers=args.num_workers,
|
||||
)
|
||||
if not SEG_CKPT.exists():
|
||||
raise FileNotFoundError(f"Segmentation checkpoint missing: {SEG_CKPT}")
|
||||
state = torch.load(SEG_CKPT, map_location=args.device)
|
||||
seg.model.load_state_dict(state)
|
||||
seg.model.to(args.device)
|
||||
return seg
|
||||
|
||||
|
||||
def train_classifier(args: argparse.Namespace) -> None:
|
||||
pre = ensure_preprocessing()
|
||||
seg = _load_segmentation(pre, args)
|
||||
backbone = build_classifier_backbone(args.clf_backbone)
|
||||
|
||||
print(f"[classifier] Using backbone: {args.clf_backbone}")
|
||||
|
||||
clf = RefugeClassification(
|
||||
pre,
|
||||
seg,
|
||||
backbone=backbone,
|
||||
cache_dir=args.clf_cache_dir,
|
||||
use_all_labeled=args.clf_use_all,
|
||||
auto_val_ratio=args.clf_auto_val_ratio,
|
||||
)
|
||||
clf.build_datasets(
|
||||
crop_scale=args.crop_scale,
|
||||
crop_size=args.crop_size,
|
||||
batch_size=args.clf_batch_size,
|
||||
num_workers=args.num_workers,
|
||||
)
|
||||
|
||||
default_ckpt_path = classifier_checkpoint_path(args.clf_backbone)
|
||||
ckpt_path = args.clf_checkpoint_path or default_ckpt_path
|
||||
ckpt_dir = ckpt_path.parent
|
||||
history = clf.train(
|
||||
epochs=args.clf_epochs,
|
||||
lr=args.clf_lr,
|
||||
weight_decay=args.clf_weight_decay,
|
||||
rotation_weight=args.rotation_weight,
|
||||
checkpoint_dir=ckpt_dir,
|
||||
device=args.device,
|
||||
)
|
||||
print("Classifier training complete. Best AUC:", history.get("best_auc"))
|
||||
print(f"Checkpoint directory: {ckpt_dir}")
|
||||
saved_path = ckpt_dir / "refuge_classifier_best.pt"
|
||||
if ckpt_path != saved_path:
|
||||
ckpt_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
shutil.copy2(saved_path, ckpt_path)
|
||||
print(f"Checkpoint copied to: {ckpt_path}")
|
||||
|
||||
|
||||
def _load_classifier(
|
||||
pre: RefugePreprocessing, seg: RefugeSegmentation, args: argparse.Namespace
|
||||
) -> Tuple[RefugeClassification, Path]:
|
||||
backbone = build_classifier_backbone(args.clf_backbone)
|
||||
clf = RefugeClassification(
|
||||
pre,
|
||||
seg,
|
||||
backbone=backbone,
|
||||
cache_dir=args.clf_cache_dir,
|
||||
use_all_labeled=args.clf_use_all,
|
||||
auto_val_ratio=args.clf_auto_val_ratio,
|
||||
)
|
||||
clf.build_datasets(
|
||||
crop_scale=args.crop_scale,
|
||||
crop_size=args.crop_size,
|
||||
batch_size=args.clf_batch_size,
|
||||
num_workers=args.num_workers,
|
||||
)
|
||||
ckpt_path = args.clf_checkpoint_path or classifier_checkpoint_path(
|
||||
args.clf_backbone
|
||||
)
|
||||
if not ckpt_path.exists():
|
||||
raise FileNotFoundError(f"Classifier checkpoint missing: {ckpt_path}")
|
||||
print(f"[classifier] Loading checkpoint: {ckpt_path}")
|
||||
state = torch.load(ckpt_path, map_location=args.device)
|
||||
clf.backbone.load_state_dict(state["backbone"])
|
||||
clf.classifier_head.load_state_dict(state["classifier"])
|
||||
clf.rotation_head.load_state_dict(state["rotation"])
|
||||
clf.backbone.to(args.device)
|
||||
clf.classifier_head.to(args.device)
|
||||
clf.rotation_head.to(args.device)
|
||||
return clf, ckpt_path
|
||||
|
||||
|
||||
def _collect_records(
|
||||
pre: RefugePreprocessing,
|
||||
seg: RefugeSegmentation,
|
||||
clf: RefugeClassification,
|
||||
dataset_name: str,
|
||||
split: str,
|
||||
scale: float,
|
||||
) -> List[RefugeClassificationRecord]:
|
||||
manifest = pre.build_manifest()
|
||||
samples = [
|
||||
sample
|
||||
for sample in manifest
|
||||
if sample.dataset == dataset_name
|
||||
and sample.split == split
|
||||
and sample.label is not None
|
||||
]
|
||||
if not samples:
|
||||
return []
|
||||
print(f"[eval] Preparing {len(samples)} samples for {dataset_name.upper()} {split}")
|
||||
return clf.build_records_for_samples(
|
||||
samples, crop_scale=scale, progress_prefix=f"{dataset_name}_{split}"
|
||||
)
|
||||
|
||||
|
||||
def _auc_for_records(
|
||||
clf: RefugeClassification,
|
||||
records: List[RefugeClassificationRecord],
|
||||
device: str,
|
||||
) -> float:
|
||||
if not records:
|
||||
return float("nan")
|
||||
dataset = RefugeClassificationDataset(
|
||||
records,
|
||||
transform=clf.eval_transform,
|
||||
polar_transform=clf.polar_transform,
|
||||
size=clf.crop_size,
|
||||
)
|
||||
loader = DataLoader(dataset, batch_size=64, shuffle=False, num_workers=0)
|
||||
clf.backbone.to(device).eval()
|
||||
clf.classifier_head.to(device).eval()
|
||||
preds: List[float] = []
|
||||
targets: List[int] = []
|
||||
with torch.no_grad():
|
||||
for batch in tqdm(loader, desc="Eval", leave=False, unit="batch"):
|
||||
images = batch["image"].to(device)
|
||||
polars = batch["polar"].to(device)
|
||||
extra_feats = batch["features"].to(device)
|
||||
labels = batch["label"].cpu().numpy().tolist()
|
||||
feats_img = clf.backbone(images)
|
||||
feats = feats_img
|
||||
if getattr(clf, "use_polar", False):
|
||||
feats_polar = clf.backbone(polars)
|
||||
feats = torch.cat([feats, feats_polar], dim=1)
|
||||
if getattr(clf, "extra_feature_dim", 0) > 0:
|
||||
feats = torch.cat([feats, extra_feats], dim=1)
|
||||
logits = clf.classifier_head(feats)
|
||||
probs = torch.softmax(logits, dim=1)[:, 1].cpu().numpy().tolist()
|
||||
preds.extend(probs)
|
||||
targets.extend(labels)
|
||||
if len(set(targets)) < 2:
|
||||
return float("nan")
|
||||
return float(roc_auc_score(targets, preds))
|
||||
|
||||
|
||||
def evaluate(args: argparse.Namespace) -> None:
|
||||
pre = ensure_preprocessing()
|
||||
seg = _load_segmentation(pre, args)
|
||||
clf, clf_ckpt = _load_classifier(pre, seg, args)
|
||||
|
||||
def evaluate_subset(
|
||||
clf_obj: RefugeClassification,
|
||||
records: List[RefugeClassificationRecord],
|
||||
label: str,
|
||||
) -> None:
|
||||
if not records:
|
||||
print(f"[eval] No samples found for {label}; skipping.")
|
||||
return
|
||||
|
||||
base_state = {
|
||||
"backbone": clf_obj.backbone.state_dict(),
|
||||
"rotation": clf_obj.rotation_head.state_dict(),
|
||||
}
|
||||
|
||||
auc_no_ttt = _auc_for_records(clf_obj, records, device=args.device)
|
||||
|
||||
auc_ttt = float("nan")
|
||||
if args.with_ttt:
|
||||
ttt_loader = DataLoader(
|
||||
RefugeClassificationDataset(
|
||||
records,
|
||||
transform=clf_obj.ttt_transform,
|
||||
polar_transform=clf_obj.polar_transform,
|
||||
size=clf_obj.crop_size,
|
||||
),
|
||||
batch_size=16,
|
||||
shuffle=False,
|
||||
num_workers=0,
|
||||
)
|
||||
ttt_iter = tqdm(range(args.ttt_steps), desc="TTT", unit="step")
|
||||
for _ in ttt_iter:
|
||||
clf_obj.apply_ttt(ttt_loader, device=args.device, steps=1)
|
||||
auc_ttt = _auc_for_records(clf_obj, records, device=args.device)
|
||||
clf_obj.backbone.load_state_dict(base_state["backbone"])
|
||||
clf_obj.rotation_head.load_state_dict(base_state["rotation"])
|
||||
|
||||
print(
|
||||
f"{label}: AUC (no TTT) = {auc_no_ttt:.4f}"
|
||||
+ (f", AUC (TTT) = {auc_ttt:.4f}" if args.with_ttt else "")
|
||||
)
|
||||
|
||||
if args.eval_datasets:
|
||||
for dataset_name in dict.fromkeys(args.eval_datasets):
|
||||
if dataset_name.lower() == "papila":
|
||||
papila_records, papila_clf = build_papila_records(args, pre, clf_ckpt)
|
||||
if papila_clf is None:
|
||||
print("[eval] Papila evaluation aborted; no samples built.")
|
||||
else:
|
||||
evaluate_subset(papila_clf, papila_records, "PAPILA holdout")
|
||||
else:
|
||||
records = _collect_records(
|
||||
pre,
|
||||
seg,
|
||||
clf,
|
||||
dataset_name,
|
||||
"holdout",
|
||||
scale=args.crop_scale,
|
||||
)
|
||||
evaluate_subset(clf, records, f"{dataset_name.upper()} holdout")
|
||||
return
|
||||
|
||||
# Do not mix splits: report per dataset + split
|
||||
subsets = [
|
||||
("refuge1", "val"),
|
||||
("refuge2", "val"),
|
||||
("refuge2", "test"),
|
||||
]
|
||||
|
||||
for dataset_name, split in subsets:
|
||||
records = _collect_records(
|
||||
pre, seg, clf, dataset_name, split, scale=args.crop_scale
|
||||
)
|
||||
evaluate_subset(clf, records, f"{dataset_name.upper()} {split}")
|
||||
|
||||
if args.dump_masks and dataset_name == "refuge1" and split == "val":
|
||||
out_dir = Path(args.dump_masks)
|
||||
out_dir.mkdir(parents=True, exist_ok=True)
|
||||
for rec in records:
|
||||
sample = rec.sample
|
||||
if sample is None:
|
||||
continue
|
||||
pred = seg.predict_mask(sample, device=args.device).numpy()
|
||||
Image.fromarray((pred * 255).astype(np.uint8)).save(
|
||||
out_dir / f"{sample.sample_id}_pred.png"
|
||||
)
|
||||
if sample.mask_path and sample.mask_path.exists():
|
||||
Image.open(sample.mask_path).convert("L").save(
|
||||
out_dir / f"{sample.sample_id}_gt.png"
|
||||
)
|
||||
|
||||
|
||||
def evaluate_segmentation(args: argparse.Namespace) -> None:
|
||||
manifest_path = args.seg_manifest or Path("manifest.csv")
|
||||
mask_cache_dir = None if args.in_memory_cache else args.mask_cache_dir
|
||||
image_cache_dir = None if args.in_memory_cache else args.image_cache_dir
|
||||
|
||||
segmenter = UNetSegmenter(
|
||||
manifest_path=manifest_path,
|
||||
normalize=args.seg_normalize,
|
||||
device=args.device,
|
||||
mask_cache_dir=mask_cache_dir,
|
||||
image_cache_dir=image_cache_dir,
|
||||
in_memory_cache=args.in_memory_cache,
|
||||
loader_workers=args.loader_workers,
|
||||
)
|
||||
if args.in_memory_cache:
|
||||
segmenter.prebuild_in_memory_cache(
|
||||
cache_workers=max(0, int(args.cache_workers)),
|
||||
include_train=False,
|
||||
include_val=bool(args.eval_seg_splits is None or "val" in args.eval_seg_splits),
|
||||
include_holdout=bool(args.eval_seg_splits is None or "holdout" in args.eval_seg_splits),
|
||||
)
|
||||
|
||||
ckpt = resolve_unet_weights(args.seg_weights)
|
||||
if ckpt.exists():
|
||||
state = torch.load(ckpt, map_location=segmenter.device)
|
||||
state_dict = state.get("model", state)
|
||||
segmenter.model.load_state_dict(state_dict, strict=False)
|
||||
print(f"[seg-eval] Loaded weights from {ckpt}")
|
||||
else:
|
||||
raise FileNotFoundError(f"Segmentation weights not found at {ckpt}")
|
||||
|
||||
dataset_filter = args.eval_seg_datasets
|
||||
split_filter = args.eval_seg_splits
|
||||
output_dir = args.eval_seg_output or Path("analysis_data/segmenter_eval")
|
||||
metrics_path = args.eval_seg_metrics_path
|
||||
|
||||
segmenter.evaluate_dataset(
|
||||
dataset_filter=dataset_filter,
|
||||
split_filter=split_filter,
|
||||
output_dir=output_dir,
|
||||
save_overlays=not args.eval_seg_no_overlays,
|
||||
metrics_path=metrics_path,
|
||||
threshold=args.eval_seg_threshold,
|
||||
tta=args.eval_seg_tta,
|
||||
)
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description="REFUGE pipeline helper")
|
||||
parser.add_argument(
|
||||
"--train-seg", action="store_true", help="Train the segmentation model"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--train-unet-seg",
|
||||
action="store_true",
|
||||
help="Train the UNet segmenter (replacement for scripts/run_unet_segmenter.py)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--train-clf", action="store_true", help="Train the classification model"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--eval", action="store_true", help="Run evaluation on stored checkpoints"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--with-ttt",
|
||||
action="store_true",
|
||||
help="Apply test-time training during evaluation",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--ttt-steps", type=int, default=1, help="TTT epochs over evaluation loader"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--export-backbone",
|
||||
type=Path,
|
||||
default=None,
|
||||
help="Optional path to export the trained backbone weights",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--dump-masks",
|
||||
type=Path,
|
||||
default=None,
|
||||
help="Optional directory to dump predicted/GT masks during eval",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--device", default="cuda" if torch.cuda.is_available() else "cpu"
|
||||
)
|
||||
parser.add_argument("--num-workers", type=int, default=4)
|
||||
# Segmentation hyperparameters
|
||||
parser.add_argument("--seg-epochs", type=int, default=40)
|
||||
parser.add_argument("--seg-lr", type=float, default=1e-3)
|
||||
parser.add_argument("--seg-weight-decay", type=float, default=1e-5)
|
||||
parser.add_argument("--seg-image-size", type=int, default=512)
|
||||
parser.add_argument("--seg-batch-size", type=int, default=4)
|
||||
parser.add_argument(
|
||||
"--seg-manifest",
|
||||
type=Path,
|
||||
default=Path("manifest.csv"),
|
||||
help="Manifest CSV for the UNet segmenter (default: manifest.csv)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--seg-weights",
|
||||
type=Path,
|
||||
default=None,
|
||||
help="Path to UNet segmenter weights (default: models/unet_segmenter/best.pt)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--seg-normalize",
|
||||
choices=["none", "imagenet", "per_image"],
|
||||
default="none",
|
||||
help="Normalization mode used when running the UNet segmenter",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--seg-strong-aug",
|
||||
action="store_true",
|
||||
help="Enable stronger geometric augmentations when training the UNet segmenter",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--seg-train-datasets",
|
||||
nargs="+",
|
||||
default=["refuge"],
|
||||
help="Datasets to use for UNet segmenter training (default: refuge)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--seg-val-datasets",
|
||||
nargs="+",
|
||||
default=["refuge"],
|
||||
help="Datasets eligible for validation sampling (default: refuge)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--seg-holdout-datasets",
|
||||
nargs="+",
|
||||
default=["refuge"],
|
||||
help="Datasets reserved for holdout set during UNet segmenter training (default: refuge)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--seg-checkpoint-dir",
|
||||
type=Path,
|
||||
default=Path("models/v2/refuge/segmentation/per_image"),
|
||||
help="Directory to store UNet segmenter checkpoints",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--loader-workers",
|
||||
type=int,
|
||||
default=0,
|
||||
help="DataLoader workers for UNet segmenter train/eval.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--mask-cache-dir",
|
||||
type=Path,
|
||||
default=None,
|
||||
help="Optional cache dir for parsed/resized disc+cup masks.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--image-cache-dir",
|
||||
type=Path,
|
||||
default=None,
|
||||
help="Optional cache dir for resized RGB images before augmentation.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--in-memory-cache",
|
||||
action="store_true",
|
||||
help="Cache preprocessed images and masks in RAM (per DataLoader worker process).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--cache-workers",
|
||||
type=int,
|
||||
default=0,
|
||||
help="Worker threads for prebuilding in-memory cache before training/eval.",
|
||||
)
|
||||
# Classification hyperparameters
|
||||
parser.add_argument("--clf-epochs", type=int, default=30)
|
||||
parser.add_argument("--clf-lr", type=float, default=1e-4)
|
||||
parser.add_argument("--clf-weight-decay", type=float, default=1e-4)
|
||||
parser.add_argument("--clf-batch-size", type=int, default=16)
|
||||
parser.add_argument(
|
||||
"--clf-backbone",
|
||||
choices=sorted(CLASSIFIER_BACKBONES.keys()),
|
||||
default="resnet50",
|
||||
help="Backbone architecture for the REFUGE classifier",
|
||||
)
|
||||
parser.add_argument("--rotation-weight", type=float, default=0.5)
|
||||
parser.add_argument("--crop-scale", type=float, default=2.5)
|
||||
parser.add_argument("--crop-size", type=int, default=224)
|
||||
parser.add_argument(
|
||||
"--clf-cache-dir",
|
||||
type=Path,
|
||||
default=Path("analysis_data/classifier_cache"),
|
||||
help="Directory to cache classifier preprocessing artifacts",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--clf-use-all",
|
||||
action="store_true",
|
||||
help="Use all labelled samples (train+val) when building classifier dataset",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--clf-auto-val-ratio",
|
||||
type=float,
|
||||
default=0.1,
|
||||
help="Fraction for automatic validation split when no explicit val set is used",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--clf-checkpoint-path",
|
||||
type=Path,
|
||||
default=None,
|
||||
help="Optional explicit path for the classifier checkpoint (defaults to models/refuge/classifier/<backbone>/refuge_classifier_best.pt)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--eval-datasets",
|
||||
nargs="+",
|
||||
help="Datasets to evaluate during --eval (e.g. papila). Defaults to REFUGE splits.",
|
||||
)
|
||||
# Segmentation evaluation parameters
|
||||
parser.add_argument(
|
||||
"--eval-seg",
|
||||
action="store_true",
|
||||
help="Evaluate the segmentation model on specified datasets/splits",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--eval-seg-datasets",
|
||||
nargs="+",
|
||||
default=["refuge"],
|
||||
help="Segmentation datasets to evaluate (default: refuge)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--eval-seg-splits",
|
||||
nargs="+",
|
||||
choices=["train", "val", "holdout"],
|
||||
help="Segmentation splits to evaluate (default: val)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--eval-seg-output",
|
||||
type=Path,
|
||||
default=Path("analysis_data/segmenter_eval"),
|
||||
help="Directory to store segmentation metrics CSVs",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--eval-seg-threshold",
|
||||
type=float,
|
||||
default=0.5,
|
||||
help="Threshold for binarising predicted masks during segmentation eval",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--eval-seg-metrics-path",
|
||||
type=Path,
|
||||
default=None,
|
||||
help="Optional explicit CSV path for segmentation metrics output",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--eval-seg-no-overlays",
|
||||
action="store_true",
|
||||
help="Skip saving GT/pred overlay images during segmentation evaluation",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--eval-seg-tta",
|
||||
action="store_true",
|
||||
help="Enable horizontal/vertical flip TTA during segmentation evaluation",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--papila-metrics",
|
||||
type=Path,
|
||||
default=None,
|
||||
help="Optional CSV of Papila Dice metrics used to filter samples",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--papila-dice-threshold",
|
||||
type=float,
|
||||
default=0.01,
|
||||
help="Minimum Dice required (disc or cup) when filtering Papila metrics",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--papila-positive-labels",
|
||||
nargs="+",
|
||||
default=["glaucoma", "glaucoma suspect", "suspect"],
|
||||
help="Papila label values treated as positive when labels are non-numeric",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--papila-image-dir",
|
||||
type=Path,
|
||||
default=Path("Papila/FundusImages"),
|
||||
help="Path to Papila fundus images",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--papila-clinical-dir",
|
||||
type=Path,
|
||||
default=Path("Papila/ClinicalData"),
|
||||
help="Path to Papila clinical CSVs",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--papila-label-col",
|
||||
type=str,
|
||||
default="Diagnosis",
|
||||
help="Column name containing Papila labels",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--papila-use-gt",
|
||||
action="store_true",
|
||||
help="Use Papila ground-truth contours when evaluating classifiers",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--papila-contours-dir",
|
||||
type=Path,
|
||||
default=Path("Papila/ExpertsSegmentations/Contours"),
|
||||
help="Directory containing Papila contour text files",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
|
||||
if not any(
|
||||
[
|
||||
args.train_seg,
|
||||
args.train_unet_seg,
|
||||
args.train_clf,
|
||||
args.eval,
|
||||
args.eval_seg,
|
||||
args.export_backbone,
|
||||
]
|
||||
):
|
||||
raise SystemExit(
|
||||
"Specify at least one action: --train-seg, --train-unet-seg, --train-clf, --eval, --eval-seg, or --export-backbone"
|
||||
)
|
||||
|
||||
if args.train_seg:
|
||||
train_segmentation(args)
|
||||
|
||||
if args.train_unet_seg:
|
||||
train_unet_segmenter(args)
|
||||
|
||||
if args.train_clf:
|
||||
train_classifier(args)
|
||||
|
||||
if args.eval:
|
||||
evaluate(args)
|
||||
|
||||
if args.eval_seg:
|
||||
evaluate_segmentation(args)
|
||||
|
||||
if args.export_backbone:
|
||||
pre = ensure_preprocessing()
|
||||
seg = _load_segmentation(pre, args)
|
||||
clf = _load_classifier(pre, seg, args)
|
||||
out_path = args.export_backbone
|
||||
out_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
torch.save(clf.extract_backbone().state_dict(), out_path)
|
||||
print(f"Backbone weights exported to {out_path}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Executable
+156
@@ -0,0 +1,156 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Train and evaluate the U-Net optic disc/cup segmenter."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
from pathlib import Path
|
||||
import torch
|
||||
|
||||
import sys
|
||||
|
||||
REPO_ROOT = Path(__file__).resolve().parents[3]
|
||||
if str(REPO_ROOT) not in sys.path:
|
||||
sys.path.insert(0, str(REPO_ROOT))
|
||||
|
||||
from classes.unet_segmenter import UNetSegmenter
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description="UNet segmenter runner")
|
||||
parser.add_argument("--manifest", type=Path, required=True, help="Path to manifest CSV")
|
||||
parser.add_argument("--train", action="store_true", help="Train the segmenter")
|
||||
parser.add_argument("--evaluate", action="store_true", help="Evaluate on holdout set")
|
||||
parser.add_argument("--epochs", type=int, default=40)
|
||||
parser.add_argument("--batch-size", type=int, default=4)
|
||||
parser.add_argument("--lr", type=float, default=1e-3)
|
||||
parser.add_argument("--weight-decay", type=float, default=1e-5)
|
||||
parser.add_argument("--disc-weight", type=float, default=1.0)
|
||||
parser.add_argument("--cup-weight", type=float, default=1.0)
|
||||
parser.add_argument("--checkpoint-dir", type=Path, default=Path("models/unet_segmenter"))
|
||||
parser.add_argument("--eval-output", type=Path, default=Path("analysis_data/segmenter_eval"))
|
||||
parser.add_argument(
|
||||
"--normalize",
|
||||
choices=["none", "imagenet", "per_image"],
|
||||
default="none",
|
||||
help="Image normalization mode for train/eval",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--strong-aug",
|
||||
action="store_true",
|
||||
help="Enable stronger train-time augmentations (flips/rotations)",
|
||||
)
|
||||
parser.add_argument("--train-datasets", nargs="+", help="Datasets to use for training/validation (default: all)")
|
||||
parser.add_argument("--val-datasets", nargs="+", help="Datasets eligible for validation sampling (default: match training)")
|
||||
parser.add_argument("--holdout-datasets", nargs="+", help="Restrict holdout entries to these datasets (default: all)")
|
||||
parser.add_argument(
|
||||
"--val-ratio",
|
||||
type=float,
|
||||
default=0.1,
|
||||
help="Fraction of training data reserved for validation (default: 0.1)",
|
||||
)
|
||||
parser.add_argument("--eval-datasets", nargs="+", help="Datasets to evaluate (default: holdout split only)")
|
||||
parser.add_argument("--eval-splits", nargs="+", help="Splits to evaluate (default: holdout or all when --eval-datasets is set)")
|
||||
parser.add_argument("--eval-metrics-path", type=Path, help="Optional CSV path for evaluation metrics output")
|
||||
parser.add_argument("--no-eval-overlays", action="store_true", help="Skip writing overlay images during evaluation")
|
||||
parser.add_argument("--threshold", type=float, default=0.5, help="Probability threshold for binarizing predictions")
|
||||
parser.add_argument("--tta", action="store_true", help="Enable simple test-time augmentation (H/V flips) during evaluation")
|
||||
parser.add_argument(
|
||||
"--weights",
|
||||
type=Path,
|
||||
help="Optional model weights (.pt) for eval-only runs; defaults to <checkpoint-dir>/best.pt",
|
||||
)
|
||||
parser.add_argument("--device", choices=["auto", "cuda", "cpu"], default="auto", help="Execution device for UNet (default: auto).")
|
||||
parser.add_argument("--loader-workers", type=int, default=0, help="DataLoader workers for train/eval.")
|
||||
parser.add_argument("--mask-cache-dir", type=Path, default=None, help="Optional cache dir for parsed/resized disc+cup masks.")
|
||||
parser.add_argument("--image-cache-dir", type=Path, default=None, help="Optional cache dir for resized RGB images before augmentation.")
|
||||
parser.add_argument("--in-memory-cache", action="store_true", help="Cache preprocessed images and masks in RAM (per DataLoader worker process).")
|
||||
parser.add_argument("--cache-workers", type=int, default=0, help="Worker threads for prebuilding in-memory cache before training/eval.")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
if args.device == "auto":
|
||||
selected_device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
else:
|
||||
selected_device = args.device
|
||||
if selected_device == "cuda" and not torch.cuda.is_available():
|
||||
raise RuntimeError("Requested --device cuda but CUDA is not available.")
|
||||
|
||||
print(
|
||||
f"[UNet] device={selected_device} "
|
||||
f"(cuda_available={torch.cuda.is_available()}, workers={args.loader_workers})"
|
||||
)
|
||||
if selected_device == "cuda":
|
||||
idx = torch.cuda.current_device()
|
||||
print(f"[UNet] gpu={torch.cuda.get_device_name(idx)}")
|
||||
|
||||
mask_cache_dir = None if args.in_memory_cache else args.mask_cache_dir
|
||||
image_cache_dir = None if args.in_memory_cache else args.image_cache_dir
|
||||
if args.in_memory_cache and (args.mask_cache_dir or args.image_cache_dir):
|
||||
print("[UNet] in_memory_cache enabled: disk caches disabled for this run.")
|
||||
|
||||
segmenter = UNetSegmenter(
|
||||
manifest_path=args.manifest,
|
||||
device=selected_device,
|
||||
cup_weight=args.cup_weight,
|
||||
disc_weight=args.disc_weight,
|
||||
val_ratio=args.val_ratio,
|
||||
train_datasets=args.train_datasets,
|
||||
val_datasets=args.val_datasets,
|
||||
holdout_datasets=args.holdout_datasets,
|
||||
normalize=args.normalize,
|
||||
use_stronger_aug=args.strong_aug,
|
||||
mask_cache_dir=mask_cache_dir,
|
||||
image_cache_dir=image_cache_dir,
|
||||
in_memory_cache=args.in_memory_cache,
|
||||
loader_workers=args.loader_workers,
|
||||
)
|
||||
if mask_cache_dir:
|
||||
print(f"[UNet] mask_cache_dir={mask_cache_dir}")
|
||||
if image_cache_dir:
|
||||
print(f"[UNet] image_cache_dir={image_cache_dir}")
|
||||
if args.in_memory_cache:
|
||||
print("[UNet] in_memory_cache=enabled (note: memory use scales with loader workers)")
|
||||
segmenter.prebuild_in_memory_cache(
|
||||
cache_workers=max(0, int(args.cache_workers)),
|
||||
include_train=bool(args.train),
|
||||
include_val=bool(args.train),
|
||||
include_holdout=bool(args.evaluate),
|
||||
)
|
||||
|
||||
if args.train:
|
||||
segmenter.train(
|
||||
epochs=args.epochs,
|
||||
batch_size=args.batch_size,
|
||||
lr=args.lr,
|
||||
weight_decay=args.weight_decay,
|
||||
checkpoint_dir=args.checkpoint_dir,
|
||||
)
|
||||
|
||||
if args.evaluate:
|
||||
if not args.train:
|
||||
ckpt = args.weights or (args.checkpoint_dir / "best.pt")
|
||||
if ckpt and ckpt.exists():
|
||||
state = torch.load(ckpt, map_location=segmenter.device)
|
||||
state_dict = state.get("model", state)
|
||||
segmenter.model.load_state_dict(state_dict, strict=False)
|
||||
print(f"Loaded weights from {ckpt}")
|
||||
else:
|
||||
print(f"[warn] No checkpoint found at {ckpt}. Evaluating untrained weights.")
|
||||
|
||||
split_filter = {"holdout"} if args.eval_splits is None else args.eval_splits
|
||||
segmenter.evaluate_dataset(
|
||||
dataset_filter=args.eval_datasets,
|
||||
split_filter=split_filter,
|
||||
output_dir=args.eval_output,
|
||||
save_overlays=not args.no_eval_overlays,
|
||||
metrics_path=args.eval_metrics_path,
|
||||
threshold=args.threshold,
|
||||
tta=args.tta,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user