Compare commits
7 Commits
a721e52909
..
main
| Author | SHA1 | Date | |
|---|---|---|---|
 rpotter6298
|
3d7777f010 | ||
|
rpotter6298
|
708fbc70ce | ||
|
rpotter6298
|
3d954a4606 | ||
|
rpotter6298
|
280060db82 | ||
|
rpotter6298
|
32a801a572 | ||
|
rpotter6298
|
4dd2dbc734 | ||
|
rpotter6298
|
86ed453736 |
+3
-1
@@ -17,4 +17,6 @@ scripts/utility/backup_mirror_with_archive.sh
|
||||
v3/distributed/logs/*
|
||||
v4/configs/**/
|
||||
v4/distributed/logs/*
|
||||
v4/results/*
|
||||
v4/results/*
|
||||
manuscript/*
|
||||
binocular_analog/*
|
||||
@@ -1,849 +0,0 @@
|
||||
"""REFUGE glaucoma classification with rotation-based TTT."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Callable, Dict, Iterable, List, Optional, Sequence, Tuple
|
||||
import random
|
||||
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
import torch
|
||||
from torch import nn
|
||||
from torch.utils.data import DataLoader, Dataset
|
||||
from torchvision import models, transforms
|
||||
from torchvision.transforms import functional as TF
|
||||
import torch.nn.functional as F
|
||||
from sklearn.metrics import roc_auc_score
|
||||
from skimage.transform import warp_polar
|
||||
from tqdm import tqdm
|
||||
|
||||
from classes.geometry_features import (
|
||||
FEATURE_DIM,
|
||||
EPS,
|
||||
compute_geometry_features,
|
||||
disc_cup_from_mask_image,
|
||||
)
|
||||
from classes.refuge_preprocessing import RefugePreprocessing, RefugeSample
|
||||
from classes.refuge_segmentation import RefugeSegmentation
|
||||
from classes.unet_segmenter import UNetSegmenter
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Dataset utilities
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _default_image_transform(size: int = 256) -> transforms.Compose:
|
||||
return transforms.Compose(
|
||||
[
|
||||
transforms.Resize((size, size)),
|
||||
transforms.ToTensor(),
|
||||
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
def _augment_image_transform(size: int = 256) -> transforms.Compose:
|
||||
return transforms.Compose(
|
||||
[
|
||||
transforms.Resize((size, size)),
|
||||
transforms.RandomHorizontalFlip(),
|
||||
transforms.RandomRotation(10),
|
||||
transforms.ColorJitter(0.1, 0.1, 0.1, 0.05),
|
||||
transforms.ToTensor(),
|
||||
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
def _crop_from_geometry(image: Image.Image, geometry: Dict[str, float], size: int = 256) -> Image.Image:
|
||||
cx, cy = geometry["centre_x"], geometry["centre_y"]
|
||||
r = geometry["crop_radius"]
|
||||
left = max(0.0, cx - r)
|
||||
upper = max(0.0, cy - r)
|
||||
right = min(image.width, cx + r)
|
||||
lower = min(image.height, cy + r)
|
||||
crop = image.crop((left, upper, right, lower))
|
||||
return crop.resize((size, size), Image.BILINEAR)
|
||||
|
||||
|
||||
def _geometry_from_mask(mask: np.ndarray, scale: float) -> Dict[str, float]:
|
||||
mask = np.asarray(mask) > 0
|
||||
coords = np.argwhere(mask)
|
||||
if coords.size == 0:
|
||||
raise RuntimeError("Empty mask; cannot derive geometry")
|
||||
ys, xs = coords[:, 0], coords[:, 1]
|
||||
centre_x = float(xs.mean())
|
||||
centre_y = float(ys.mean())
|
||||
width = float(xs.max() - xs.min())
|
||||
height = float(ys.max() - ys.min())
|
||||
diameter = max(width, height)
|
||||
radius = diameter / 2.0
|
||||
crop_radius = radius * scale
|
||||
return {
|
||||
"centre_x": centre_x,
|
||||
"centre_y": centre_y,
|
||||
"radius": radius,
|
||||
"crop_radius": crop_radius,
|
||||
"crop_size": crop_radius * 2.0,
|
||||
}
|
||||
|
||||
|
||||
def _compute_feature_vector(disc_mask: np.ndarray, cup_mask: np.ndarray) -> np.ndarray:
|
||||
return compute_geometry_features(disc_mask, cup_mask)
|
||||
|
||||
|
||||
def _compute_polar_image(crop: Image.Image, size: int) -> Image.Image:
|
||||
arr = np.asarray(crop).astype(np.float32) / 255.0
|
||||
radius = min(arr.shape[0], arr.shape[1]) / 2.0
|
||||
polar = warp_polar(
|
||||
arr,
|
||||
radius=radius,
|
||||
scaling="linear",
|
||||
channel_axis=-1,
|
||||
)
|
||||
polar = np.clip(polar, 0.0, 1.0)
|
||||
polar_img = Image.fromarray((polar * 255).astype(np.uint8))
|
||||
return polar_img.resize((size, size), Image.BILINEAR)
|
||||
|
||||
|
||||
def _crop_mask_from_geometry(mask: np.ndarray, geometry: Dict[str, float], size: int) -> np.ndarray:
|
||||
mask_img = Image.fromarray((mask > 0).astype(np.uint8) * 255)
|
||||
cx, cy = geometry["centre_x"], geometry["centre_y"]
|
||||
r = geometry["crop_radius"]
|
||||
left = max(0.0, cx - r)
|
||||
upper = max(0.0, cy - r)
|
||||
right = min(mask_img.width, cx + r)
|
||||
lower = min(mask_img.height, cy + r)
|
||||
crop = mask_img.crop((left, upper, right, lower)).resize((size, size), Image.NEAREST)
|
||||
return (np.asarray(crop) > 0).astype(np.uint8)
|
||||
|
||||
|
||||
@dataclass
|
||||
class RefugeClassificationRecord:
|
||||
sample: RefugeSample
|
||||
geometry: Dict[str, float]
|
||||
disc_mask: Optional[np.ndarray] = None
|
||||
cup_mask: Optional[np.ndarray] = None
|
||||
|
||||
|
||||
class RefugeClassificationDataset(Dataset):
|
||||
def __init__(
|
||||
self,
|
||||
records: Sequence[RefugeClassificationRecord],
|
||||
transform: transforms.Compose,
|
||||
polar_transform: transforms.Compose,
|
||||
size: int = 256,
|
||||
) -> None:
|
||||
self.records = list(records)
|
||||
self.transform = transform
|
||||
self.polar_transform = polar_transform
|
||||
self.size = size
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self.records)
|
||||
|
||||
def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
|
||||
rec = self.records[idx]
|
||||
image = Image.open(rec.sample.image_path).convert("RGB")
|
||||
crop = _crop_from_geometry(image, rec.geometry, size=self.size)
|
||||
polar_image = _compute_polar_image(crop, size=self.size)
|
||||
tensor = self.transform(crop)
|
||||
polar_tensor = self.polar_transform(polar_image)
|
||||
|
||||
features = np.zeros((FEATURE_DIM,), dtype=np.float32)
|
||||
if rec.disc_mask is not None and rec.cup_mask is not None:
|
||||
disc_crop = _crop_mask_from_geometry(rec.disc_mask, rec.geometry, self.size)
|
||||
cup_crop = _crop_mask_from_geometry(rec.cup_mask, rec.geometry, self.size)
|
||||
features = _compute_feature_vector(disc_crop, cup_crop)
|
||||
|
||||
feature_tensor = torch.from_numpy(features).float()
|
||||
label = rec.sample.label
|
||||
if label is None:
|
||||
raise ValueError(f"Sample {rec.sample.sample_id} is missing glaucoma label")
|
||||
return {
|
||||
"image": tensor,
|
||||
"polar": polar_tensor,
|
||||
"features": feature_tensor,
|
||||
"label": torch.tensor(label, dtype=torch.long),
|
||||
"sample_id": rec.sample.sample_id,
|
||||
}
|
||||
|
||||
|
||||
class RefugeTTTDataset(Dataset):
|
||||
"""Dataset providing unlabeled crops for test-time training."""
|
||||
|
||||
def __init__(self, records: Sequence[RefugeClassificationRecord], transform: transforms.Compose, size: int = 256) -> None:
|
||||
self.records = list(records)
|
||||
self.transform = transform
|
||||
self.size = size
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self.records)
|
||||
|
||||
def __getitem__(self, idx: int) -> torch.Tensor:
|
||||
rec = self.records[idx]
|
||||
image = Image.open(rec.sample.image_path).convert("RGB")
|
||||
crop = _crop_from_geometry(image, rec.geometry, size=self.size)
|
||||
return self.transform(crop)
|
||||
|
||||
|
||||
class UNetGeometryProvider:
|
||||
"""Callable wrapper that derives disc geometry using a trained UNetSegmenter."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
segmenter: UNetSegmenter,
|
||||
threshold: float = 0.5,
|
||||
tta: bool = False,
|
||||
) -> None:
|
||||
self.segmenter = segmenter
|
||||
self.threshold = threshold
|
||||
self.tta = tta
|
||||
self.segmenter.model.eval()
|
||||
|
||||
def __call__(self, sample: RefugeSample, scale: float) -> Tuple[Dict[str, float], np.ndarray, np.ndarray]:
|
||||
image = Image.open(sample.image_path).convert("RGB")
|
||||
resized = self.segmenter.preprocess_image(image)
|
||||
tensor = transforms.ToTensor()(resized)
|
||||
tensor = self.segmenter._normalize_tensor(tensor)
|
||||
tensor = tensor.unsqueeze(0).to(self.segmenter.device)
|
||||
with torch.no_grad():
|
||||
logits = self.segmenter.model(tensor)
|
||||
if self.tta:
|
||||
t_h = torch.flip(tensor, dims=[3])
|
||||
log_h = self.segmenter.model(t_h)
|
||||
log_h = torch.flip(log_h, dims=[3])
|
||||
t_v = torch.flip(tensor, dims=[2])
|
||||
log_v = self.segmenter.model(t_v)
|
||||
log_v = torch.flip(log_v, dims=[2])
|
||||
logits = (logits + log_h + log_v) / 3.0
|
||||
probs = torch.sigmoid(logits)[0].cpu().numpy()
|
||||
|
||||
disc_pred = (probs[0] > self.threshold).astype(np.uint8) * 255
|
||||
cup_pred = (probs[1] > self.threshold).astype(np.uint8) * 255
|
||||
disc_img = Image.fromarray(disc_pred, mode="L").resize(image.size, Image.NEAREST)
|
||||
cup_img = Image.fromarray(cup_pred, mode="L").resize(image.size, Image.NEAREST)
|
||||
disc_mask = (np.array(disc_img, dtype=np.uint8) > 0).astype(np.uint8)
|
||||
cup_mask = (np.array(cup_img, dtype=np.uint8) > 0).astype(np.uint8)
|
||||
cup_mask = (cup_mask > 0) & (disc_mask > 0)
|
||||
cup_mask = cup_mask.astype(np.uint8)
|
||||
geom = _geometry_from_mask(disc_mask, scale)
|
||||
return geom, disc_mask, cup_mask
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Classification module
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class ArcMarginProduct(nn.Module):
|
||||
"""Additive angular margin (ArcFace) head."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
in_features: int,
|
||||
out_features: int,
|
||||
s: float = 30.0,
|
||||
m: float = 0.5,
|
||||
easy_margin: bool = False,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.in_features = in_features
|
||||
self.out_features = out_features
|
||||
self.s = float(s)
|
||||
self.m = float(m)
|
||||
self.easy_margin = easy_margin
|
||||
self.weight = nn.Parameter(torch.empty(out_features, in_features))
|
||||
nn.init.xavier_uniform_(self.weight)
|
||||
|
||||
self.cos_m = math.cos(m)
|
||||
self.sin_m = math.sin(m)
|
||||
self.th = math.cos(math.pi - m)
|
||||
self.mm = math.sin(math.pi - m) * m
|
||||
|
||||
def forward(self, input: torch.Tensor, label: Optional[torch.Tensor] = None) -> torch.Tensor:
|
||||
cosine = F.linear(F.normalize(input), F.normalize(self.weight))
|
||||
if label is None:
|
||||
return cosine * self.s
|
||||
|
||||
sine = torch.sqrt(torch.clamp(1.0 - cosine.pow(2), min=0.0))
|
||||
phi = cosine * self.cos_m - sine * self.sin_m
|
||||
if self.easy_margin:
|
||||
phi = torch.where(cosine > 0, phi, cosine)
|
||||
else:
|
||||
phi = torch.where(cosine > self.th, phi, cosine - self.mm)
|
||||
|
||||
one_hot = torch.zeros_like(cosine)
|
||||
one_hot.scatter_(1, label.view(-1, 1), 1.0)
|
||||
logits = (one_hot * phi) + ((1.0 - one_hot) * cosine)
|
||||
logits *= self.s
|
||||
return logits
|
||||
|
||||
|
||||
class RefugeClassification:
|
||||
"""Train and evaluate REFUGE glaucoma classifiers with TTT support."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
preprocessing: RefugePreprocessing,
|
||||
segmentation: RefugeSegmentation,
|
||||
backbone: Optional[nn.Module] = None,
|
||||
geometry_fn: Optional[
|
||||
Callable[
|
||||
[RefugeSample, float],
|
||||
Tuple[Dict[str, float], Optional[np.ndarray], Optional[np.ndarray]],
|
||||
]
|
||||
] = None,
|
||||
cache_dir: Optional[Path] = None,
|
||||
use_all_labeled: bool = False,
|
||||
auto_val_ratio: float = 0.1,
|
||||
use_margin: bool = False,
|
||||
margin_s: float = 30.0,
|
||||
margin_m: float = 0.5,
|
||||
) -> None:
|
||||
self.preprocessing = preprocessing
|
||||
self.segmentation = segmentation
|
||||
if backbone is not None:
|
||||
self.backbone = backbone
|
||||
in_features = getattr(self.backbone, "_feature_dim", None)
|
||||
if in_features is None:
|
||||
if hasattr(self.backbone, "fc") and hasattr(self.backbone.fc, "in_features"):
|
||||
in_features = self.backbone.fc.in_features # type: ignore[attr-defined]
|
||||
self.backbone.fc = nn.Identity() # type: ignore[attr-defined]
|
||||
else:
|
||||
raise ValueError(
|
||||
"Provided backbone must have '_feature_dim' or expose fc.in_features"
|
||||
)
|
||||
else:
|
||||
self.backbone = self._default_backbone()
|
||||
in_features = getattr(self.backbone, "_feature_dim", None)
|
||||
if in_features is None:
|
||||
in_features = self.backbone.fc.in_features # type: ignore[attr-defined]
|
||||
self.backbone.fc = nn.Identity() # type: ignore[attr-defined]
|
||||
self.feature_dim = in_features
|
||||
self.use_polar = True
|
||||
self.extra_feature_dim = FEATURE_DIM
|
||||
combined_dim = self.feature_dim * (1 + int(self.use_polar)) + self.extra_feature_dim
|
||||
self.margin_s = float(margin_s)
|
||||
self.margin_m = float(margin_m)
|
||||
self.use_margin = bool(use_margin)
|
||||
if self.use_margin:
|
||||
self.classifier_head = ArcMarginProduct(
|
||||
combined_dim, 2, s=self.margin_s, m=self.margin_m
|
||||
)
|
||||
else:
|
||||
self.classifier_head = nn.Linear(combined_dim, 2)
|
||||
self.rotation_head = nn.Linear(self.feature_dim, 4)
|
||||
|
||||
self.train_dataset: Optional[RefugeClassificationDataset] = None
|
||||
self.val_dataset: Optional[RefugeClassificationDataset] = None
|
||||
self.train_loader: Optional[DataLoader] = None
|
||||
self.val_loader: Optional[DataLoader] = None
|
||||
self.ttt_transform = _default_image_transform()
|
||||
self.train_transform = _augment_image_transform()
|
||||
self.eval_transform = _default_image_transform()
|
||||
self.polar_transform = _default_image_transform()
|
||||
self.crop_scale = 2.5
|
||||
self.crop_size = 256
|
||||
self.geometry_cache: Dict[
|
||||
str, Tuple[Dict[str, float], Optional[np.ndarray], Optional[np.ndarray]]
|
||||
] = {}
|
||||
self.train_records: List[RefugeClassificationRecord] = []
|
||||
self.val_records: List[RefugeClassificationRecord] = []
|
||||
self._geometry_fn = geometry_fn
|
||||
self.cache_dir = cache_dir
|
||||
if self.cache_dir is not None:
|
||||
self.cache_dir.mkdir(parents=True, exist_ok=True)
|
||||
self.use_all_labeled = use_all_labeled
|
||||
self.auto_val_ratio = auto_val_ratio
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
@staticmethod
|
||||
def _default_backbone() -> nn.Module:
|
||||
weights = models.ResNet50_Weights.IMAGENET1K_V2
|
||||
model = models.resnet50(weights=weights)
|
||||
in_features = model.fc.in_features
|
||||
model.fc = nn.Identity()
|
||||
setattr(model, "_feature_dim", in_features)
|
||||
return model
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
def build_datasets(
|
||||
self,
|
||||
crop_scale: float = 2.5,
|
||||
crop_size: int = 256,
|
||||
batch_size: int = 16,
|
||||
num_workers: int = 4,
|
||||
) -> None:
|
||||
self.crop_scale = crop_scale
|
||||
self.crop_size = crop_size
|
||||
self.train_transform = _augment_image_transform(crop_size)
|
||||
self.eval_transform = _default_image_transform(crop_size)
|
||||
self.ttt_transform = _default_image_transform(crop_size)
|
||||
self.polar_transform = _default_image_transform(crop_size)
|
||||
|
||||
manifest = list(self.preprocessing.build_manifest())
|
||||
train_records: List[RefugeClassificationRecord] = []
|
||||
val_records: List[RefugeClassificationRecord] = []
|
||||
|
||||
allowed_splits = {"train", "val"}
|
||||
candidates = [
|
||||
sample
|
||||
for sample in manifest
|
||||
if sample.label is not None and sample.split in allowed_splits
|
||||
]
|
||||
|
||||
print(
|
||||
f"[classifier] Building datasets from {len(candidates)} labelled samples (train/val)"
|
||||
)
|
||||
|
||||
skipped: List[str] = []
|
||||
for sample in tqdm(
|
||||
candidates,
|
||||
desc="Preparing records",
|
||||
unit="sample",
|
||||
leave=False,
|
||||
):
|
||||
try:
|
||||
geom, disc_mask, cup_mask = self._resolve_geometry(sample, crop_scale)
|
||||
except RuntimeError:
|
||||
skipped.append(sample.sample_id)
|
||||
continue
|
||||
record = RefugeClassificationRecord(
|
||||
sample=sample,
|
||||
geometry=geom,
|
||||
disc_mask=disc_mask,
|
||||
cup_mask=cup_mask,
|
||||
)
|
||||
if sample.split == "train" or (
|
||||
self.use_all_labeled and sample.split == "val"
|
||||
):
|
||||
train_records.append(record)
|
||||
else:
|
||||
val_records.append(record)
|
||||
|
||||
if skipped:
|
||||
print(
|
||||
f"[classifier] WARNING: {len(skipped)}/{len(candidates)} samples skipped "
|
||||
f"due to empty segmentation mask: {skipped}"
|
||||
)
|
||||
|
||||
if (not val_records or self.use_all_labeled) and train_records and self.auto_val_ratio > 0.0:
|
||||
rng = random.Random(42)
|
||||
label_groups: Dict[int, List[RefugeClassificationRecord]] = {}
|
||||
for rec in train_records:
|
||||
label = int(rec.sample.label or 0)
|
||||
label_groups.setdefault(label, []).append(rec)
|
||||
|
||||
new_train: List[RefugeClassificationRecord] = []
|
||||
new_val: List[RefugeClassificationRecord] = []
|
||||
for recs in label_groups.values():
|
||||
rng.shuffle(recs)
|
||||
if len(recs) <= 1:
|
||||
new_train.extend(recs)
|
||||
continue
|
||||
val_count = max(1, int(round(len(recs) * self.auto_val_ratio)))
|
||||
if val_count >= len(recs):
|
||||
val_count = len(recs) - 1
|
||||
new_val.extend(recs[:val_count])
|
||||
new_train.extend(recs[val_count:])
|
||||
|
||||
if not new_val:
|
||||
# Fallback: ensure at least one validation sample if possible
|
||||
if len(new_train) > 1:
|
||||
new_val.append(new_train.pop())
|
||||
|
||||
if new_val:
|
||||
val_records = new_val
|
||||
train_records = new_train
|
||||
|
||||
self.train_records = train_records
|
||||
self.val_records = val_records
|
||||
|
||||
print(
|
||||
f"[classifier] Records ready → train: {len(train_records)}, val: {len(val_records)}"
|
||||
)
|
||||
|
||||
self.train_dataset = RefugeClassificationDataset(
|
||||
train_records,
|
||||
transform=self.train_transform,
|
||||
polar_transform=self.polar_transform,
|
||||
size=crop_size,
|
||||
)
|
||||
self.val_dataset = RefugeClassificationDataset(
|
||||
val_records,
|
||||
transform=self.eval_transform,
|
||||
polar_transform=self.polar_transform,
|
||||
size=crop_size,
|
||||
)
|
||||
|
||||
self.train_loader = DataLoader(
|
||||
self.train_dataset,
|
||||
batch_size=batch_size,
|
||||
shuffle=True,
|
||||
num_workers=num_workers,
|
||||
pin_memory=True,
|
||||
)
|
||||
self.val_loader = DataLoader(
|
||||
self.val_dataset,
|
||||
batch_size=batch_size,
|
||||
shuffle=False,
|
||||
num_workers=num_workers,
|
||||
pin_memory=True,
|
||||
)
|
||||
|
||||
print(
|
||||
"[classifier] DataLoaders prepared — training batches will start shortly"
|
||||
)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
def _resolve_geometry(
|
||||
self, sample: RefugeSample, scale: float
|
||||
) -> Tuple[Dict[str, float], Optional[np.ndarray], Optional[np.ndarray]]:
|
||||
key = self._cache_key(sample.sample_id, scale)
|
||||
cached = self.geometry_cache.get(key)
|
||||
if cached is not None:
|
||||
return cached
|
||||
|
||||
cache_path = self._cache_path(sample.sample_id, scale)
|
||||
if cache_path is not None and cache_path.exists():
|
||||
data = np.load(cache_path, allow_pickle=False)
|
||||
geom = {
|
||||
"centre_x": float(data["centre_x"]),
|
||||
"centre_y": float(data["centre_y"]),
|
||||
"radius": float(data["radius"]),
|
||||
"crop_radius": float(data["crop_radius"]),
|
||||
"crop_size": float(data["crop_size"]),
|
||||
}
|
||||
disc_mask = None
|
||||
cup_mask = None
|
||||
if int(data["has_disc"]):
|
||||
disc_mask = data["disc_mask"].astype(np.uint8)
|
||||
if int(data["has_cup"]):
|
||||
cup_mask = data["cup_mask"].astype(np.uint8)
|
||||
self.geometry_cache[key] = (geom, disc_mask, cup_mask)
|
||||
return geom, disc_mask, cup_mask
|
||||
|
||||
disc_mask: Optional[np.ndarray] = None
|
||||
cup_mask: Optional[np.ndarray] = None
|
||||
|
||||
if sample.mask_path and sample.mask_path.exists():
|
||||
mask_img = Image.open(sample.mask_path).convert("RGB")
|
||||
disc_mask, cup_mask = disc_cup_from_mask_image(mask_img)
|
||||
geom = _geometry_from_mask(disc_mask, scale)
|
||||
elif self._geometry_fn is not None:
|
||||
geom, disc_mask, cup_mask = self._geometry_fn(sample, scale)
|
||||
else:
|
||||
geom = self.segmentation.infer_disc_geometry(sample, scale=scale)
|
||||
try:
|
||||
pred_mask = self.segmentation.predict_mask(sample).numpy()
|
||||
disc_mask = pred_mask.astype(np.uint8)
|
||||
except Exception:
|
||||
disc_mask = None
|
||||
cup_mask = None
|
||||
|
||||
if cache_path is not None:
|
||||
try:
|
||||
np.savez_compressed(
|
||||
cache_path,
|
||||
centre_x=geom["centre_x"],
|
||||
centre_y=geom["centre_y"],
|
||||
radius=geom["radius"],
|
||||
crop_radius=geom["crop_radius"],
|
||||
crop_size=geom.get("crop_size", geom["crop_radius"] * 2.0),
|
||||
disc_mask=disc_mask if disc_mask is not None else np.array([], dtype=np.uint8),
|
||||
cup_mask=cup_mask if cup_mask is not None else np.array([], dtype=np.uint8),
|
||||
has_disc=int(disc_mask is not None),
|
||||
has_cup=int(cup_mask is not None),
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
self.geometry_cache[key] = (geom, disc_mask, cup_mask)
|
||||
return geom, disc_mask, cup_mask
|
||||
|
||||
def set_geometry_fn(
|
||||
self,
|
||||
geometry_fn: Optional[
|
||||
Callable[
|
||||
[RefugeSample, float],
|
||||
Tuple[Dict[str, float], Optional[np.ndarray], Optional[np.ndarray]],
|
||||
]
|
||||
],
|
||||
) -> None:
|
||||
self._geometry_fn = geometry_fn
|
||||
self.geometry_cache.clear()
|
||||
|
||||
def build_records_for_samples(
|
||||
self,
|
||||
samples: Sequence[RefugeSample],
|
||||
crop_scale: Optional[float] = None,
|
||||
progress_prefix: Optional[str] = None,
|
||||
) -> List[RefugeClassificationRecord]:
|
||||
scale = crop_scale if crop_scale is not None else self.crop_scale
|
||||
records: List[RefugeClassificationRecord] = []
|
||||
skipped: List[str] = []
|
||||
iterator: Iterable[RefugeSample]
|
||||
if progress_prefix is not None:
|
||||
iterator = tqdm(samples, desc=progress_prefix, unit="sample", leave=False)
|
||||
else:
|
||||
iterator = samples
|
||||
labeled = [s for s in samples if s.label is not None]
|
||||
for sample in iterator:
|
||||
if sample.label is None:
|
||||
continue
|
||||
try:
|
||||
geom, disc_mask, cup_mask = self._resolve_geometry(sample, scale)
|
||||
except RuntimeError:
|
||||
skipped.append(sample.sample_id)
|
||||
continue
|
||||
records.append(
|
||||
RefugeClassificationRecord(
|
||||
sample=sample,
|
||||
geometry=geom,
|
||||
disc_mask=disc_mask,
|
||||
cup_mask=cup_mask,
|
||||
)
|
||||
)
|
||||
prefix = f"[{progress_prefix}]" if progress_prefix else "[classifier]"
|
||||
if skipped:
|
||||
print(
|
||||
f"{prefix} WARNING: {len(skipped)}/{len(labeled)} samples skipped "
|
||||
f"due to empty segmentation mask: {skipped}"
|
||||
)
|
||||
else:
|
||||
print(f"{prefix} All {len(labeled)} samples processed successfully.")
|
||||
return records
|
||||
|
||||
def clear_disk_cache(self) -> None:
|
||||
"""Delete all cached geometry/mask .npz files in cache_dir."""
|
||||
if self.cache_dir is None or not self.cache_dir.exists():
|
||||
return
|
||||
removed = 0
|
||||
for f in self.cache_dir.glob("*.npz"):
|
||||
f.unlink()
|
||||
removed += 1
|
||||
self.geometry_cache.clear()
|
||||
print(f"[classifier] Cleared {removed} cached geometry files from {self.cache_dir}")
|
||||
|
||||
def _cache_key(self, sample_id: str, scale: float) -> str:
|
||||
scale_tag = int(round(scale * 100))
|
||||
return f"{sample_id}_s{scale_tag}"
|
||||
|
||||
def _cache_path(self, sample_id: str, scale: float) -> Optional[Path]:
|
||||
if self.cache_dir is None:
|
||||
return None
|
||||
return self.cache_dir / f"{self._cache_key(sample_id, scale)}.npz"
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
def train(
|
||||
self,
|
||||
epochs: int = 30,
|
||||
lr: float = 1e-4,
|
||||
weight_decay: float = 1e-4,
|
||||
device: Optional[str] = None,
|
||||
rotation_weight: float = 0.5,
|
||||
checkpoint_dir: Optional[Path] = None,
|
||||
) -> Dict[str, float]:
|
||||
if self.train_loader is None or self.val_loader is None:
|
||||
self.build_datasets()
|
||||
|
||||
device = device or ("cuda" if torch.cuda.is_available() else "cpu")
|
||||
self.backbone.to(device)
|
||||
self.classifier_head.to(device)
|
||||
self.rotation_head.to(device)
|
||||
|
||||
params = list(self.backbone.parameters()) + list(self.classifier_head.parameters()) + list(self.rotation_head.parameters())
|
||||
optimizer = torch.optim.Adam(params, lr=lr, weight_decay=weight_decay)
|
||||
clf_loss = nn.CrossEntropyLoss()
|
||||
rot_loss = nn.CrossEntropyLoss()
|
||||
|
||||
best_auc = 0.0
|
||||
history: Dict[str, float] = {}
|
||||
|
||||
epoch_iter = tqdm(range(1, epochs + 1), desc="Epochs", unit="epoch")
|
||||
|
||||
print(
|
||||
f"[classifier] Starting training for {epochs} epochs with batch size {self.train_loader.batch_size}"
|
||||
)
|
||||
|
||||
for epoch in epoch_iter:
|
||||
self.backbone.train()
|
||||
self.classifier_head.train()
|
||||
self.rotation_head.train()
|
||||
running_loss = 0.0
|
||||
|
||||
batch_iter = tqdm(
|
||||
self.train_loader, # type: ignore[arg-type]
|
||||
desc=f"Train {epoch}/{epochs}",
|
||||
leave=False,
|
||||
unit="batch",
|
||||
)
|
||||
|
||||
for batch in batch_iter:
|
||||
images = batch["image"].to(device)
|
||||
polars = batch["polar"].to(device)
|
||||
extra_feats = batch["features"].to(device)
|
||||
labels = batch["label"].to(device)
|
||||
optimizer.zero_grad()
|
||||
|
||||
feats_img = self.backbone(images)
|
||||
feats = feats_img
|
||||
if self.use_polar:
|
||||
feats_polar = self.backbone(polars)
|
||||
feats = torch.cat([feats, feats_polar], dim=1)
|
||||
if self.extra_feature_dim > 0:
|
||||
feats = torch.cat([feats, extra_feats], dim=1)
|
||||
if self.use_margin:
|
||||
logits = self.classifier_head(feats, labels)
|
||||
else:
|
||||
logits = self.classifier_head(feats)
|
||||
loss_cls = clf_loss(logits, labels)
|
||||
|
||||
rot_imgs, rot_labels = self._build_rotation_batch(images)
|
||||
feats_rot = self.backbone(rot_imgs)
|
||||
logits_rot = self.rotation_head(feats_rot)
|
||||
loss_rot = rot_loss(logits_rot, rot_labels)
|
||||
|
||||
loss = loss_cls + rotation_weight * loss_rot
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
running_loss += loss.item() * images.size(0)
|
||||
|
||||
train_loss = running_loss / len(self.train_loader.dataset) # type: ignore[arg-type]
|
||||
metrics = self.evaluate(device=device)
|
||||
history[f"epoch_{epoch}_loss"] = train_loss
|
||||
history[f"epoch_{epoch}_auc"] = metrics.get("auc", float("nan"))
|
||||
|
||||
auc_val = metrics.get("auc", 0.0)
|
||||
epoch_iter.set_postfix(loss=f"{train_loss:.4f}", auc=f"{auc_val:.4f}")
|
||||
|
||||
if auc_val > best_auc:
|
||||
best_auc = metrics["auc"]
|
||||
if checkpoint_dir is not None:
|
||||
checkpoint_dir.mkdir(parents=True, exist_ok=True)
|
||||
torch.save({
|
||||
"backbone": self.backbone.state_dict(),
|
||||
"classifier": self.classifier_head.state_dict(),
|
||||
"rotation": self.rotation_head.state_dict(),
|
||||
}, checkpoint_dir / "refuge_classifier_best.pt")
|
||||
|
||||
return {"best_auc": best_auc, **history}
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
def evaluate(
|
||||
self,
|
||||
split: str = "val",
|
||||
apply_ttt: bool = False,
|
||||
device: Optional[str] = None,
|
||||
) -> Dict[str, float]:
|
||||
if split != "val":
|
||||
raise ValueError("Only validation split supported currently")
|
||||
if self.val_loader is None:
|
||||
self.build_datasets()
|
||||
|
||||
device = device or ("cuda" if torch.cuda.is_available() else "cpu")
|
||||
self.backbone.to(device)
|
||||
self.classifier_head.to(device)
|
||||
self.rotation_head.to(device)
|
||||
|
||||
if apply_ttt:
|
||||
ttt_ds = RefugeTTTDataset(self.val_records, transform=self.ttt_transform, size=self.crop_size)
|
||||
ttt_loader = DataLoader(ttt_ds, batch_size=32, shuffle=False)
|
||||
self.apply_ttt(ttt_loader, device=device)
|
||||
|
||||
self.backbone.eval()
|
||||
self.classifier_head.eval()
|
||||
preds: List[float] = []
|
||||
targets: List[int] = []
|
||||
|
||||
with torch.no_grad():
|
||||
val_iter = tqdm(self.val_loader, desc="Validate", leave=False, unit="batch")
|
||||
for batch in val_iter: # type: ignore[arg-type]
|
||||
images = batch["image"].to(device)
|
||||
labels = batch["label"].to(device)
|
||||
polars = batch["polar"].to(device)
|
||||
extra_feats = batch["features"].to(device)
|
||||
feats_img = self.backbone(images)
|
||||
feats = feats_img
|
||||
if self.use_polar:
|
||||
feats_polar = self.backbone(polars)
|
||||
feats = torch.cat([feats, feats_polar], dim=1)
|
||||
if self.extra_feature_dim > 0:
|
||||
feats = torch.cat([feats, extra_feats], dim=1)
|
||||
if self.use_margin:
|
||||
logits = self.classifier_head(feats)
|
||||
else:
|
||||
logits = self.classifier_head(feats)
|
||||
probs = torch.softmax(logits, dim=1)[:, 1]
|
||||
preds.extend(probs.cpu().numpy().tolist())
|
||||
targets.extend(labels.cpu().numpy().tolist())
|
||||
|
||||
auc = 0.0
|
||||
try:
|
||||
if len(set(targets)) > 1:
|
||||
auc = float(roc_auc_score(targets, preds))
|
||||
except ValueError:
|
||||
auc = 0.0
|
||||
|
||||
return {"auc": auc}
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
def apply_ttt(self, loader: DataLoader, device: Optional[str] = None, steps: int = 1, lr: float = 1e-5) -> None:
|
||||
device = device or ("cuda" if torch.cuda.is_available() else "cpu")
|
||||
self.backbone.to(device)
|
||||
self.rotation_head.to(device)
|
||||
self.backbone.train()
|
||||
self.rotation_head.train()
|
||||
|
||||
optimizer = torch.optim.Adam(list(self.backbone.parameters()) + list(self.rotation_head.parameters()), lr=lr)
|
||||
criterion = nn.CrossEntropyLoss()
|
||||
|
||||
for _ in range(steps):
|
||||
for batch in tqdm(loader, desc="TTT adapt", leave=False, unit="batch"):
|
||||
if isinstance(batch, dict):
|
||||
images = batch["image"].to(device)
|
||||
else:
|
||||
images = batch.to(device)
|
||||
optimizer.zero_grad()
|
||||
rot_imgs, rot_labels = self._build_rotation_batch(images)
|
||||
feats = self.backbone(rot_imgs)
|
||||
logits = self.rotation_head(feats)
|
||||
loss = criterion(logits, rot_labels)
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
def extract_backbone(self) -> nn.Module:
|
||||
return self.backbone
|
||||
|
||||
def save_checkpoint(self, output_dir: Path) -> None:
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
torch.save({
|
||||
"backbone": self.backbone.state_dict(),
|
||||
"classifier": self.classifier_head.state_dict(),
|
||||
"rotation": self.rotation_head.state_dict(),
|
||||
}, output_dir / "refuge_classifier.pt")
|
||||
|
||||
def load_checkpoint(self, checkpoint_path: Path) -> None:
|
||||
payload = torch.load(checkpoint_path, map_location="cpu")
|
||||
self.backbone.load_state_dict(payload["backbone"])
|
||||
self.classifier_head.load_state_dict(payload["classifier"])
|
||||
self.rotation_head.load_state_dict(payload["rotation"])
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
def _build_rotation_batch(self, images: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
rotations = [0, 90, 180, 270]
|
||||
rotated = []
|
||||
labels = []
|
||||
for idx, angle in enumerate(rotations):
|
||||
rot = TF.rotate(images, angle)
|
||||
rotated.append(rot)
|
||||
labels.append(torch.full((images.size(0),), idx, dtype=torch.long, device=images.device))
|
||||
batch = torch.cat(rotated, dim=0)
|
||||
batch_labels = torch.cat(labels, dim=0)
|
||||
return batch, batch_labels
|
||||
@@ -1,306 +0,0 @@
|
||||
"""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)
|
||||
@@ -1,383 +0,0 @@
|
||||
"""REFUGE optic disc / cup segmentation utilities."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Dict, Iterable, List, Optional, Sequence, Tuple
|
||||
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
import torch
|
||||
from torch import nn
|
||||
from torch.utils.data import DataLoader, Dataset
|
||||
from torchvision import transforms
|
||||
|
||||
from classes.refuge_preprocessing import RefugePreprocessing, RefugeSample
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Dataset helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _load_rgb(path: Path) -> Image.Image:
|
||||
img = Image.open(path)
|
||||
if img.mode != "RGB":
|
||||
img = img.convert("RGB")
|
||||
return img
|
||||
|
||||
|
||||
def _load_mask_array(path: Path) -> np.ndarray:
|
||||
mask_img = Image.open(path).convert("L")
|
||||
mask = np.array(mask_img, dtype=np.float32)
|
||||
# REFUGE masks encode disc/cup with different intensities; treat any
|
||||
# positive value as disc for coarse localisation.
|
||||
mask = np.where(mask > 0, 1.0, 0.0)
|
||||
return mask
|
||||
|
||||
|
||||
@dataclass
|
||||
class RefugeSegmentationSample:
|
||||
sample: RefugeSample
|
||||
image_path: Path
|
||||
mask_path: Path
|
||||
|
||||
|
||||
class RefugeSegmentationDataset(Dataset):
|
||||
"""Simple segmentation dataset returning tensors."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
samples: Sequence[RefugeSegmentationSample],
|
||||
image_size: int = 512,
|
||||
) -> None:
|
||||
self.samples = list(samples)
|
||||
self.image_size = image_size
|
||||
self.to_tensor = transforms.ToTensor()
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self.samples)
|
||||
|
||||
def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
|
||||
rec = self.samples[idx]
|
||||
image = _load_rgb(rec.image_path)
|
||||
mask_arr = _load_mask_array(rec.mask_path)
|
||||
|
||||
if self.image_size is not None:
|
||||
image = image.resize((self.image_size, self.image_size), Image.BILINEAR)
|
||||
mask_img = Image.fromarray(mask_arr).resize(
|
||||
(self.image_size, self.image_size), Image.NEAREST
|
||||
)
|
||||
mask_arr = np.array(mask_img, dtype=np.float32)
|
||||
|
||||
image_tensor = self.to_tensor(image)
|
||||
mask_tensor = torch.from_numpy(mask_arr).unsqueeze(0) # [1,H,W]
|
||||
return {
|
||||
"image": image_tensor,
|
||||
"mask": mask_tensor,
|
||||
"sample_id": rec.sample.sample_id,
|
||||
}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Model definition (lightweight U-Net)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class DoubleConv(nn.Module):
|
||||
def __init__(self, in_channels: int, out_channels: int):
|
||||
super().__init__()
|
||||
self.net = nn.Sequential(
|
||||
nn.Conv2d(in_channels, out_channels, 3, padding=1, bias=False),
|
||||
nn.BatchNorm2d(out_channels),
|
||||
nn.ReLU(inplace=True),
|
||||
nn.Conv2d(out_channels, out_channels, 3, padding=1, bias=False),
|
||||
nn.BatchNorm2d(out_channels),
|
||||
nn.ReLU(inplace=True),
|
||||
)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
return self.net(x)
|
||||
|
||||
|
||||
class UNet(nn.Module):
|
||||
def __init__(self, in_channels: int = 3, base_channels: int = 64):
|
||||
super().__init__()
|
||||
self.enc1 = DoubleConv(in_channels, base_channels)
|
||||
self.enc2 = DoubleConv(base_channels, base_channels * 2)
|
||||
self.enc3 = DoubleConv(base_channels * 2, base_channels * 4)
|
||||
self.enc4 = DoubleConv(base_channels * 4, base_channels * 8)
|
||||
|
||||
self.pool = nn.MaxPool2d(2)
|
||||
self.bottleneck = DoubleConv(base_channels * 8, base_channels * 16)
|
||||
|
||||
self.up4 = nn.ConvTranspose2d(base_channels * 16, base_channels * 8, 2, stride=2)
|
||||
self.dec4 = DoubleConv(base_channels * 16, base_channels * 8)
|
||||
self.up3 = nn.ConvTranspose2d(base_channels * 8, base_channels * 4, 2, stride=2)
|
||||
self.dec3 = DoubleConv(base_channels * 8, base_channels * 4)
|
||||
self.up2 = nn.ConvTranspose2d(base_channels * 4, base_channels * 2, 2, stride=2)
|
||||
self.dec2 = DoubleConv(base_channels * 4, base_channels * 2)
|
||||
self.up1 = nn.ConvTranspose2d(base_channels * 2, base_channels, 2, stride=2)
|
||||
self.dec1 = DoubleConv(base_channels * 2, base_channels)
|
||||
|
||||
self.out = nn.Conv2d(base_channels, 1, 1)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
e1 = self.enc1(x)
|
||||
e2 = self.enc2(self.pool(e1))
|
||||
e3 = self.enc3(self.pool(e2))
|
||||
e4 = self.enc4(self.pool(e3))
|
||||
b = self.bottleneck(self.pool(e4))
|
||||
|
||||
d4 = self.up4(b)
|
||||
d4 = torch.cat([d4, e4], dim=1)
|
||||
d4 = self.dec4(d4)
|
||||
d3 = self.up3(d4)
|
||||
d3 = torch.cat([d3, e3], dim=1)
|
||||
d3 = self.dec3(d3)
|
||||
d2 = self.up2(d3)
|
||||
d2 = torch.cat([d2, e2], dim=1)
|
||||
d2 = self.dec2(d2)
|
||||
d1 = self.up1(d2)
|
||||
d1 = torch.cat([d1, e1], dim=1)
|
||||
d1 = self.dec1(d1)
|
||||
return self.out(d1)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Segmentation manager
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class RefugeSegmentation:
|
||||
"""Train and run coarse-to-fine OD/OC segmentation for REFUGE."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
preprocessing: RefugePreprocessing,
|
||||
model: Optional[nn.Module] = None,
|
||||
) -> None:
|
||||
self.preprocessing = preprocessing
|
||||
self.model = model or UNet()
|
||||
self.train_dataset: Optional[RefugeSegmentationDataset] = None
|
||||
self.val_dataset: Optional[RefugeSegmentationDataset] = None
|
||||
self.train_loader: Optional[DataLoader] = None
|
||||
self.val_loader: Optional[DataLoader] = None
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
def build_datasets(
|
||||
self,
|
||||
image_size: int = 512,
|
||||
batch_size: int = 8,
|
||||
num_workers: int = 4,
|
||||
) -> None:
|
||||
manifest = self.preprocessing.build_manifest()
|
||||
|
||||
train_samples: List[RefugeSegmentationSample] = []
|
||||
val_samples: List[RefugeSegmentationSample] = []
|
||||
|
||||
for sample in manifest:
|
||||
if not sample.mask_path or not sample.mask_path.exists():
|
||||
continue
|
||||
rec = RefugeSegmentationSample(sample=sample, image_path=sample.image_path, mask_path=sample.mask_path)
|
||||
if sample.split == "train":
|
||||
train_samples.append(rec)
|
||||
elif sample.split in {"val", "validation"}:
|
||||
val_samples.append(rec)
|
||||
|
||||
if not val_samples:
|
||||
# Fall back to using a subset of training data for validation
|
||||
split = max(1, int(0.1 * len(train_samples)))
|
||||
val_samples = train_samples[:split]
|
||||
train_samples = train_samples[split:]
|
||||
|
||||
self.train_dataset = RefugeSegmentationDataset(train_samples, image_size=image_size)
|
||||
self.val_dataset = RefugeSegmentationDataset(val_samples, image_size=image_size)
|
||||
self.train_loader = DataLoader(
|
||||
self.train_dataset,
|
||||
batch_size=batch_size,
|
||||
shuffle=True,
|
||||
num_workers=num_workers,
|
||||
pin_memory=True,
|
||||
)
|
||||
self.val_loader = DataLoader(
|
||||
self.val_dataset,
|
||||
batch_size=batch_size,
|
||||
shuffle=False,
|
||||
num_workers=num_workers,
|
||||
pin_memory=True,
|
||||
)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
def train(
|
||||
self,
|
||||
epochs: int = 40,
|
||||
lr: float = 1e-3,
|
||||
weight_decay: float = 1e-5,
|
||||
device: Optional[str] = None,
|
||||
checkpoint_dir: Optional[Path] = None,
|
||||
) -> Dict[str, float]:
|
||||
if self.train_loader is None or self.val_loader is None:
|
||||
self.build_datasets()
|
||||
|
||||
device = device or ("cuda" if torch.cuda.is_available() else "cpu")
|
||||
self.model.to(device)
|
||||
criterion = nn.BCEWithLogitsLoss()
|
||||
optimizer = torch.optim.Adam(self.model.parameters(), lr=lr, weight_decay=weight_decay)
|
||||
|
||||
best_dice = 0.0
|
||||
history: Dict[str, float] = {}
|
||||
|
||||
for epoch in range(1, epochs + 1):
|
||||
print(f"[Seg] Processing epoch {epoch}/{epochs}")
|
||||
self.model.train()
|
||||
running_loss = 0.0
|
||||
for batch in self.train_loader: # type: ignore[arg-type]
|
||||
images = batch["image"].to(device)
|
||||
masks = batch["mask"].to(device)
|
||||
optimizer.zero_grad()
|
||||
logits = self.model(images)
|
||||
loss = criterion(logits, masks)
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
running_loss += loss.item() * images.size(0)
|
||||
|
||||
train_loss = running_loss / len(self.train_loader.dataset) # type: ignore[arg-type]
|
||||
val_metrics = self.evaluate(device=device)
|
||||
history[f"epoch_{epoch}_loss"] = train_loss
|
||||
history[f"epoch_{epoch}_dice"] = val_metrics.get("dice", float("nan"))
|
||||
|
||||
if val_metrics.get("dice", 0.0) > best_dice:
|
||||
best_dice = val_metrics["dice"]
|
||||
if checkpoint_dir is not None:
|
||||
checkpoint_dir.mkdir(parents=True, exist_ok=True)
|
||||
torch.save(self.model.state_dict(), checkpoint_dir / "refuge_segmentation_best.pt")
|
||||
|
||||
return {"best_dice": best_dice, **history}
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
def evaluate(self, split: str = "val", device: Optional[str] = None) -> Dict[str, float]:
|
||||
if split != "val":
|
||||
raise ValueError("Only validation split supported currently")
|
||||
if self.val_loader is None:
|
||||
self.build_datasets()
|
||||
|
||||
device = device or ("cuda" if torch.cuda.is_available() else "cpu")
|
||||
self.model.to(device)
|
||||
self.model.eval()
|
||||
|
||||
dices: List[float] = []
|
||||
criterion = nn.BCEWithLogitsLoss()
|
||||
losses: List[float] = []
|
||||
|
||||
with torch.no_grad():
|
||||
for batch in self.val_loader: # type: ignore[arg-type]
|
||||
images = batch["image"].to(device)
|
||||
masks = batch["mask"].to(device)
|
||||
logits = self.model(images)
|
||||
loss = criterion(logits, masks)
|
||||
losses.append(loss.item() * images.size(0))
|
||||
probs = torch.sigmoid(logits)
|
||||
preds = (probs > 0.5).float()
|
||||
dice = self._dice_coefficient(preds, masks)
|
||||
dices.extend(dice)
|
||||
|
||||
mean_dice = float(np.mean(dices)) if dices else 0.0
|
||||
mean_loss = float(np.sum(losses) / len(self.val_loader.dataset)) # type: ignore[arg-type]
|
||||
return {"dice": mean_dice, "loss": mean_loss}
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
def predict_mask(self, sample: RefugeSample, device: Optional[str] = None) -> torch.Tensor:
|
||||
if self.train_dataset is None:
|
||||
self.build_datasets()
|
||||
device = device or ("cuda" if torch.cuda.is_available() else "cpu")
|
||||
self.model.to(device)
|
||||
self.model.eval()
|
||||
|
||||
image = _load_rgb(sample.image_path)
|
||||
original_size = image.size # (width, height)
|
||||
image_resized = image.resize((self.train_dataset.image_size, self.train_dataset.image_size), Image.BILINEAR) # type: ignore[union-attr]
|
||||
tensor = transforms.ToTensor()(image_resized).unsqueeze(0).to(device)
|
||||
|
||||
with torch.no_grad():
|
||||
logits = self.model(tensor)
|
||||
mask_resized = torch.sigmoid(logits)[0, 0]
|
||||
|
||||
mask_np = mask_resized.cpu().numpy()
|
||||
mask_np = (mask_np > 0.5).astype(np.float32)
|
||||
mask_img = Image.fromarray(mask_np)
|
||||
mask_img = mask_img.resize(original_size, Image.NEAREST)
|
||||
return torch.from_numpy(np.array(mask_img, dtype=np.float32))
|
||||
|
||||
def infer_disc_geometry(
|
||||
self,
|
||||
sample: RefugeSample,
|
||||
scale: float = 2.5,
|
||||
) -> Dict[str, float]:
|
||||
if sample.mask_path and sample.mask_path.exists():
|
||||
mask = _load_mask_array(sample.mask_path)
|
||||
else:
|
||||
mask = self.predict_mask(sample).numpy()
|
||||
|
||||
coords = np.argwhere(mask > 0.5)
|
||||
if coords.size == 0:
|
||||
raise RuntimeError(f"Unable to locate disc for sample {sample.sample_id}")
|
||||
|
||||
ys, xs = coords[:, 0], coords[:, 1]
|
||||
centre_x = float(xs.mean())
|
||||
centre_y = float(ys.mean())
|
||||
width = float(xs.max() - xs.min())
|
||||
height = float(ys.max() - ys.min())
|
||||
diameter = max(width, height)
|
||||
radius = diameter / 2.0
|
||||
crop_radius = radius * scale
|
||||
return {
|
||||
"centre_x": centre_x,
|
||||
"centre_y": centre_y,
|
||||
"radius": radius,
|
||||
"crop_radius": crop_radius,
|
||||
"crop_size": crop_radius * 2.0,
|
||||
}
|
||||
|
||||
def batch_crops(
|
||||
self,
|
||||
samples: Iterable[RefugeSample],
|
||||
scale: float = 2.5,
|
||||
output_dir: Optional[Path] = None,
|
||||
size: int = 256,
|
||||
) -> Dict[str, Path]:
|
||||
output_paths: Dict[str, Path] = {}
|
||||
if output_dir is not None:
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
for sample in samples:
|
||||
geom = self.infer_disc_geometry(sample, scale=scale)
|
||||
image = _load_rgb(sample.image_path)
|
||||
cx, cy = geom["centre_x"], geom["centre_y"]
|
||||
r = geom["crop_radius"]
|
||||
left = max(0.0, cx - r)
|
||||
upper = max(0.0, cy - r)
|
||||
right = min(image.width, cx + r)
|
||||
lower = min(image.height, cy + r)
|
||||
crop = image.crop((left, upper, right, lower)).resize((size, size), Image.BILINEAR)
|
||||
if output_dir is not None:
|
||||
out_path = output_dir / f"{sample.sample_id}_crop.png"
|
||||
crop.save(out_path)
|
||||
output_paths[sample.sample_id] = out_path
|
||||
return output_paths
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
@staticmethod
|
||||
def _dice_coefficient(preds: torch.Tensor, targets: torch.Tensor) -> List[float]:
|
||||
eps = 1e-6
|
||||
dices = []
|
||||
preds = preds.view(preds.size(0), -1)
|
||||
targets = targets.view(targets.size(0), -1)
|
||||
for p, t in zip(preds, targets):
|
||||
intersection = float((p * t).sum().item())
|
||||
union = float(p.sum().item() + t.sum().item())
|
||||
dice = (2.0 * intersection + eps) / (union + eps)
|
||||
dices.append(dice)
|
||||
return dices
|
||||
@@ -1,123 +0,0 @@
|
||||
# se_block.py
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
class SEGateLogger:
|
||||
"""
|
||||
Lightweight stats over SE gates.
|
||||
Use: logger.accumulate(gates) each batch; logger.get() at epoch end.
|
||||
"""
|
||||
def __init__(self, enabled: bool = True, track_channels: bool = False, dim: int | None = None):
|
||||
self.enabled = enabled
|
||||
self.track_channels = track_channels
|
||||
self.dim = dim
|
||||
self.reset()
|
||||
|
||||
def reset(self):
|
||||
self._n = 0
|
||||
self._sum = 0.0
|
||||
self._sum2 = 0.0
|
||||
self._lt02 = 0
|
||||
self._gt08 = 0
|
||||
# optional per-channel
|
||||
self._ch_sum = None
|
||||
self._ch_count = 0
|
||||
if self.track_channels and self.dim is not None:
|
||||
self._ch_sum = torch.zeros(self.dim, dtype=torch.float32)
|
||||
|
||||
@torch.no_grad()
|
||||
def accumulate(self, gates: torch.Tensor):
|
||||
if not self.enabled:
|
||||
return
|
||||
# gates expected shape [N, C]; if a map/sequence gate is passed, reduce to [N, C]
|
||||
if gates.dim() == 4: # [N,C,H,W] gates (uncommon)
|
||||
g = gates.mean(dim=(2,3))
|
||||
elif gates.dim() == 3: # [N,T,C] gates (sequence)
|
||||
g = gates.mean(dim=1)
|
||||
elif gates.dim() == 2: # [N,C]
|
||||
g = gates
|
||||
else:
|
||||
g = gates.view(gates.size(0), -1)
|
||||
|
||||
g = g.detach()
|
||||
self._n += g.numel()
|
||||
self._sum += g.sum().item()
|
||||
self._sum2 += (g*g).sum().item()
|
||||
self._lt02 += (g < 0.2).sum().item()
|
||||
self._gt08 += (g > 0.8).sum().item()
|
||||
|
||||
if self._ch_sum is not None:
|
||||
self._ch_sum += g.sum(dim=0).cpu()
|
||||
self._ch_count += g.size(0)
|
||||
|
||||
def get(self, reset: bool = True):
|
||||
if self._n == 0:
|
||||
return None
|
||||
mean = self._sum / self._n
|
||||
var = max(0.0, self._sum2 / self._n - mean * mean)
|
||||
out = {
|
||||
"mean": mean,
|
||||
"std": var ** 0.5,
|
||||
"pct_lt_0.2": self._lt02 / self._n,
|
||||
"pct_gt_0.8": self._gt08 / self._n,
|
||||
}
|
||||
if self._ch_sum is not None and self._ch_count > 0:
|
||||
out["channel_mean"] = (self._ch_sum / float(self._ch_count)).tolist()
|
||||
if reset:
|
||||
self.reset()
|
||||
return out
|
||||
|
||||
class SEBlock(nn.Module):
|
||||
"""
|
||||
SE-style channel gating that works for vectors and maps.
|
||||
|
||||
Input:
|
||||
- [N, C] (vector) -> squeeze = identity
|
||||
- [N, C, H, W] (image map) -> squeeze over H,W
|
||||
- [N, T, C] (sequence) -> squeeze over T
|
||||
|
||||
Gate modes:
|
||||
- residual (default): gate = 1 + tanh(MLP(s)) in (0, 2) [identity at init]
|
||||
- plain: gate = sigmoid(MLP(s)) in (0, 1)
|
||||
"""
|
||||
def __init__(self, dim: int, reduction: int = 16, residual: bool = True, identity_init: bool = True):
|
||||
super().__init__()
|
||||
hid = max(1, dim // max(1, reduction))
|
||||
self.fc1 = nn.Linear(dim, hid, bias=True)
|
||||
self.act = nn.ReLU(inplace=True)
|
||||
self.fc2 = nn.Linear(hid, dim, bias=True)
|
||||
self.residual = residual
|
||||
|
||||
if residual and identity_init:
|
||||
# make MLP output ~0 at start → gate ≈ 1.0
|
||||
nn.init.zeros_(self.fc2.weight)
|
||||
nn.init.zeros_(self.fc2.bias)
|
||||
|
||||
def _squeeze(self, x: torch.Tensor) -> torch.Tensor:
|
||||
if x.dim() == 2: # [N,C]
|
||||
return x
|
||||
if x.dim() == 4: # [N,C,H,W]
|
||||
return x.mean(dim=(2,3))
|
||||
if x.dim() == 3: # [N,T,C]
|
||||
return x.mean(dim=1)
|
||||
# fallback: flatten non-batch dims into channels
|
||||
return x.view(x.size(0), -1)
|
||||
|
||||
def _broadcast(self, gate: torch.Tensor, like: torch.Tensor) -> torch.Tensor:
|
||||
if like.dim() == 2:
|
||||
return gate
|
||||
if like.dim() == 3:
|
||||
return gate.unsqueeze(1) # [N,1,C]
|
||||
if like.dim() == 4:
|
||||
return gate.unsqueeze(-1).unsqueeze(-1) # [N,C,1,1]
|
||||
return gate.view_as(like)
|
||||
|
||||
def forward(self, x: torch.Tensor):
|
||||
s = self._squeeze(x) # [N,C]
|
||||
u = self.fc2(self.act(self.fc1(s))) # [N,C]
|
||||
if self.residual:
|
||||
gate = 1.0 + torch.tanh(u) # (0, 2) with identity at 1.0
|
||||
else:
|
||||
gate = torch.sigmoid(u) # (0, 1)
|
||||
y = x * self._broadcast(gate, x)
|
||||
return y, gate # return both the reweighted tensor and the gate for logging
|
||||
@@ -1,107 +0,0 @@
|
||||
from .network_manager import (
|
||||
FoldResult,
|
||||
LoaderBundle,
|
||||
NetworkManager,
|
||||
PatientSplit,
|
||||
)
|
||||
from .split_manager import (
|
||||
PatientFirstSplitManager,
|
||||
SplitPlan,
|
||||
build_patient_split_plans,
|
||||
)
|
||||
from .profiles import (
|
||||
DatasetProfile,
|
||||
SimpleDatasetProfile,
|
||||
SlotDescriptor,
|
||||
PapilaProfile,
|
||||
build_papila_profile,
|
||||
)
|
||||
from .loader_factory import SlotLoaderFactory
|
||||
from .slot_dataset import SlotDataset, slot_collate
|
||||
from .papila_data import PapilaData
|
||||
from .papila_builders import build_papila_data
|
||||
from .data_bundle import DataBundle
|
||||
from .dataset import ClinicalDataset
|
||||
from .config_builder import (
|
||||
ConfigAssembly,
|
||||
assemble_config,
|
||||
load_config,
|
||||
resolve_imports,
|
||||
)
|
||||
from .filters import RegexFilter, ColumnFilter, apply_regex_filters, apply_column_filters
|
||||
from .transforms import (
|
||||
ImageTransformConfig,
|
||||
backbone_transform_config,
|
||||
build_backbone_transform,
|
||||
build_eval_transform,
|
||||
build_imagenet_transform,
|
||||
ResizeTransform,
|
||||
CenterCropTransform,
|
||||
ROICropTransform,
|
||||
JitterBundleTransform,
|
||||
UnetMaskProvider,
|
||||
TRANSFORM_REGISTRY,
|
||||
build_transform_chain,
|
||||
)
|
||||
from .model_builder import V2ModelBundle, build_model_bundle
|
||||
from .towers import ImageTower, MDTower, SiameseImageTower, build_backbone
|
||||
from .bridges import Bridge, VoteBridge
|
||||
from .models import SingleEyeHT, BilateralHT
|
||||
from .v2_hypertower import V2HyperTower, V2ModeComparisonOps, V2ModeComparator
|
||||
from .hypertower_logger import HypertowerLogger
|
||||
|
||||
__all__ = [
|
||||
"NetworkManager",
|
||||
"PatientSplit",
|
||||
"LoaderBundle",
|
||||
"FoldResult",
|
||||
"PatientFirstSplitManager",
|
||||
"SplitPlan",
|
||||
"build_patient_split_plans",
|
||||
"DatasetProfile",
|
||||
"SimpleDatasetProfile",
|
||||
"SlotDescriptor",
|
||||
"PapilaProfile",
|
||||
"build_papila_profile",
|
||||
"PapilaData",
|
||||
"build_papila_data",
|
||||
"DataBundle",
|
||||
"ClinicalDataset",
|
||||
"SlotLoaderFactory",
|
||||
"SlotDataset",
|
||||
"slot_collate",
|
||||
"ConfigAssembly",
|
||||
"assemble_config",
|
||||
"load_config",
|
||||
"resolve_imports",
|
||||
"RegexFilter",
|
||||
"ColumnFilter",
|
||||
"apply_regex_filters",
|
||||
"apply_column_filters",
|
||||
"ImageTransformConfig",
|
||||
"backbone_transform_config",
|
||||
"build_backbone_transform",
|
||||
"build_eval_transform",
|
||||
"build_imagenet_transform",
|
||||
"ResizeTransform",
|
||||
"CenterCropTransform",
|
||||
"ROICropTransform",
|
||||
"JitterBundleTransform",
|
||||
"UnetMaskProvider",
|
||||
"TRANSFORM_REGISTRY",
|
||||
"build_transform_chain",
|
||||
"V2ModelBundle",
|
||||
"build_model_bundle",
|
||||
"ImageTower",
|
||||
"MDTower",
|
||||
"SiameseImageTower",
|
||||
"build_backbone",
|
||||
"Bridge",
|
||||
"VoteBridge",
|
||||
"SingleEyeHT",
|
||||
"BilateralHT",
|
||||
"V2HyperTower",
|
||||
"V2ModeComparisonOps",
|
||||
"V2ModeComparator",
|
||||
"HypertowerLogger",
|
||||
]
|
||||
@@ -1,178 +0,0 @@
|
||||
# classes/backbones.py
|
||||
from __future__ import annotations
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Callable, Dict, List
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
from torchvision import models
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class BackboneSpec:
|
||||
ctor: Callable # torchvision constructor
|
||||
weights_default: object # torchvision Weights enum DEFAULT member
|
||||
strip: Callable[[nn.Module], tuple] # fn(model)->(out_dim, model_no_head)
|
||||
blocks: Callable[[nn.Module], List[nn.Module]] # fn(model)->ordered blocks for freezing
|
||||
|
||||
REFUGELIKE_BACKBONE_PATH = Path("models/v2/refuge/refugelike_backbone.pt")
|
||||
REFUGE_DENSENET_PATH = Path("models/refuge/classifier/refuge_densenet_backbone.pt")
|
||||
REFUGE_EFFICIENT_B0_PATH = Path("models/refuge/classifier/refuge_efficient_b0_backbone.pt")
|
||||
REFUGE_EFFICIENT_B7_PATH = Path("models/refuge/classifier/refuge_efficient_b7_backbone.pt")
|
||||
|
||||
# --- strip fns ---
|
||||
def _strip_efficientnet_b0(m: models.EfficientNet):
|
||||
from torch import nn as _nn
|
||||
out_dim = m.classifier[1].in_features
|
||||
m.classifier = _nn.Identity()
|
||||
return out_dim, m
|
||||
|
||||
def _strip_resnet(m: models.ResNet):
|
||||
out_dim = m.fc.in_features
|
||||
m.fc = nn.Identity()
|
||||
return out_dim, m
|
||||
|
||||
def _strip_densenet(m: models.DenseNet):
|
||||
out_dim = m.classifier.in_features
|
||||
m.classifier = nn.Identity()
|
||||
return out_dim, m
|
||||
|
||||
def _strip_vgg(m: models.VGG):
|
||||
out_dim = m.classifier[0].in_features # 25088 for VGG16 at 224×224
|
||||
m.classifier = nn.Identity()
|
||||
return out_dim, m
|
||||
|
||||
def _strip_mobilenet_v2(m: models.MobileNetV2):
|
||||
out_dim = m.classifier[1].in_features
|
||||
m.classifier = nn.Identity()
|
||||
return out_dim, m
|
||||
|
||||
def _strip_inception_v3(m: models.Inception3):
|
||||
out_dim = m.fc.in_features
|
||||
m.fc = nn.Identity()
|
||||
if hasattr(m, "AuxLogits"):
|
||||
m.aux_logits = False
|
||||
return out_dim, m
|
||||
|
||||
# --- block splitters for ratio-based freezing ---
|
||||
def _blocks_efficientnet_b0(m: models.EfficientNet):
|
||||
return list(m.features)
|
||||
|
||||
def _blocks_resnet(m: models.ResNet):
|
||||
stem = nn.Sequential(m.conv1, m.bn1, m.relu, m.maxpool)
|
||||
return [stem, m.layer1, m.layer2, m.layer3, m.layer4]
|
||||
|
||||
def _blocks_densenet(m: models.DenseNet):
|
||||
f = m.features
|
||||
stem = nn.Sequential(f.conv0, f.norm0, f.relu0, f.pool0)
|
||||
return [stem, f.denseblock1, f.transition1, f.denseblock2, f.transition2,
|
||||
f.denseblock3, f.transition3, f.denseblock4, f.norm5]
|
||||
|
||||
def _blocks_vgg(m: models.VGG):
|
||||
stages, cur = [], []
|
||||
for mod in m.features:
|
||||
cur.append(mod)
|
||||
if isinstance(mod, nn.MaxPool2d):
|
||||
stages.append(nn.Sequential(*cur)); cur = []
|
||||
if cur: stages.append(nn.Sequential(*cur))
|
||||
return stages
|
||||
|
||||
def _blocks_mobilenet_v2(m: models.MobileNetV2):
|
||||
return list(m.features)
|
||||
|
||||
def _blocks_inception_v3(m: models.Inception3):
|
||||
blocks = []
|
||||
for name, child in m.named_children():
|
||||
if name in ("fc", "AuxLogits"):
|
||||
continue
|
||||
blocks.append(child)
|
||||
return blocks
|
||||
|
||||
# --- registry (covers paper models available in torchvision) ---
|
||||
BACKBONES: Dict[str, BackboneSpec] = {
|
||||
"efficientnet_b0": BackboneSpec(
|
||||
ctor=models.efficientnet_b0,
|
||||
weights_default=models.EfficientNet_B0_Weights.DEFAULT,
|
||||
strip=_strip_efficientnet_b0,
|
||||
blocks=_blocks_efficientnet_b0,
|
||||
),
|
||||
"resnet50": BackboneSpec(
|
||||
ctor=models.resnet50,
|
||||
weights_default=models.ResNet50_Weights.DEFAULT,
|
||||
strip=_strip_resnet,
|
||||
blocks=_blocks_resnet,
|
||||
),
|
||||
"densenet121": BackboneSpec(
|
||||
ctor=models.densenet121,
|
||||
weights_default=models.DenseNet121_Weights.DEFAULT,
|
||||
strip=_strip_densenet,
|
||||
blocks=_blocks_densenet,
|
||||
),
|
||||
"vgg16": BackboneSpec(
|
||||
ctor=models.vgg16,
|
||||
weights_default=models.VGG16_Weights.DEFAULT,
|
||||
strip=_strip_vgg,
|
||||
blocks=_blocks_vgg,
|
||||
),
|
||||
"mobilenet_v2": BackboneSpec(
|
||||
ctor=models.mobilenet_v2,
|
||||
weights_default=models.MobileNet_V2_Weights.DEFAULT,
|
||||
strip=_strip_mobilenet_v2,
|
||||
blocks=_blocks_mobilenet_v2,
|
||||
),
|
||||
"inception_v3": BackboneSpec(
|
||||
ctor=models.inception_v3,
|
||||
weights_default=models.Inception_V3_Weights.DEFAULT,
|
||||
strip=_strip_inception_v3,
|
||||
blocks=_blocks_inception_v3,
|
||||
),
|
||||
"refugelike": BackboneSpec(
|
||||
ctor=models.resnet50,
|
||||
weights_default=None,
|
||||
strip=_strip_resnet,
|
||||
blocks=_blocks_resnet,
|
||||
),
|
||||
"refuge_densenet": BackboneSpec(
|
||||
ctor=models.densenet121,
|
||||
weights_default=None,
|
||||
strip=_strip_densenet,
|
||||
blocks=_blocks_densenet,
|
||||
),
|
||||
"refuge_efficient_b0": BackboneSpec(
|
||||
ctor=models.efficientnet_b0,
|
||||
weights_default=None,
|
||||
strip=_strip_efficientnet_b0,
|
||||
blocks=_blocks_efficientnet_b0,
|
||||
),
|
||||
"refuge_efficient_b7": BackboneSpec(
|
||||
ctor=models.efficientnet_b7,
|
||||
weights_default=None,
|
||||
strip=_strip_efficientnet_b0,
|
||||
blocks=_blocks_efficientnet_b0,
|
||||
),
|
||||
# Xception isn’t in torchvision
|
||||
}
|
||||
|
||||
def list_names() -> List[str]:
|
||||
return list(BACKBONES.keys())
|
||||
|
||||
|
||||
def load_backbone_weights(key: str, model: nn.Module) -> None:
|
||||
if key == "refugelike":
|
||||
path = REFUGELIKE_BACKBONE_PATH
|
||||
elif key == "refuge_densenet":
|
||||
path = REFUGE_DENSENET_PATH
|
||||
elif key == "refuge_efficient_b0":
|
||||
path = REFUGE_EFFICIENT_B0_PATH
|
||||
elif key == "refuge_efficient_b7":
|
||||
path = REFUGE_EFFICIENT_B7_PATH
|
||||
else:
|
||||
return
|
||||
|
||||
if not path.exists():
|
||||
raise FileNotFoundError(
|
||||
"Custom REFUGE backbone not found at "
|
||||
f"{path}. Export it via refuge_build.py --export-backbone first."
|
||||
)
|
||||
state = torch.load(path, map_location="cpu")
|
||||
model.load_state_dict(state, strict=False)
|
||||
@@ -1,93 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from classes.v2.SE_attention import SEBlock, SEGateLogger
|
||||
|
||||
|
||||
class Bridge(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
img_dim,
|
||||
meta_dim,
|
||||
num_classes,
|
||||
fusion_dim=256,
|
||||
mode="fused",
|
||||
use_se: bool = True,
|
||||
se_reduction: int = 16,
|
||||
se_pre_norm: bool = True,
|
||||
):
|
||||
super().__init__()
|
||||
self.mode = mode
|
||||
self.use_se = use_se
|
||||
|
||||
# project towers to equal width
|
||||
self.W_img = nn.Linear(img_dim, fusion_dim)
|
||||
self.W_md = nn.Linear(meta_dim, fusion_dim)
|
||||
|
||||
# optional: layernorm before SE
|
||||
self.ln_img = nn.LayerNorm(fusion_dim) if se_pre_norm else nn.Identity()
|
||||
self.ln_md = nn.LayerNorm(fusion_dim) if se_pre_norm else nn.Identity()
|
||||
|
||||
# SE gate on the fused vector
|
||||
self.se = SEBlock(fusion_dim, reduction=se_reduction, residual=True) if use_se else None
|
||||
self.se_log = SEGateLogger(enabled=use_se, track_channels=False, dim=fusion_dim)
|
||||
|
||||
# heads
|
||||
self.classifier_fused = nn.Sequential(
|
||||
nn.ReLU(),
|
||||
nn.Dropout(0.5),
|
||||
nn.Linear(fusion_dim, num_classes),
|
||||
)
|
||||
self.classifier_img = nn.Linear(img_dim, num_classes)
|
||||
self.classifier_md = nn.Linear(meta_dim, num_classes)
|
||||
|
||||
def reset_se_stats(self):
|
||||
"""Call at epoch start."""
|
||||
if getattr(self, "se_log", None):
|
||||
self.se_log.reset()
|
||||
|
||||
def get_se_stats(self, reset: bool = True):
|
||||
"""Call after eval. Returns dict or None."""
|
||||
if getattr(self, "se_log", None) and self.se_log.enabled:
|
||||
return self.se_log.get(reset=reset)
|
||||
return None
|
||||
|
||||
def forward(self, img_feats, md_feats):
|
||||
out_img = None if self.mode == "metadata_only" else self.classifier_img(img_feats)
|
||||
out_md = None if self.mode == "image_only" else self.classifier_md(md_feats)
|
||||
|
||||
if self.mode == "fused":
|
||||
hi = self.ln_img(self.W_img(img_feats)) # image features
|
||||
hm = self.ln_md(self.W_md(md_feats)) # metadata features
|
||||
fused = hi * hm # elementwise product
|
||||
# apply SE gates
|
||||
if self.se is not None:
|
||||
fused, gates = self.se(fused)
|
||||
if self.se_log.enabled:
|
||||
self.se_log.accumulate(gates)
|
||||
|
||||
if self.se is not None and self.training and self.se_log.enabled:
|
||||
if not hasattr(self, "_dbg_seen"):
|
||||
self._dbg_seen = 0
|
||||
if self._dbg_seen < 3: # print only a few times
|
||||
print("[SE] gate mean this batch:", gates.mean().item())
|
||||
self._dbg_seen += 1
|
||||
out_f = self.classifier_fused(fused)
|
||||
return out_f, out_img, out_md
|
||||
# if ablation modes:
|
||||
if self.mode == "image_only":
|
||||
return out_img, out_img, None
|
||||
if self.mode == "metadata_only":
|
||||
return out_md, None, out_md
|
||||
|
||||
|
||||
class VoteBridge(nn.Module):
|
||||
def __init__(self, num_classes):
|
||||
super().__init__()
|
||||
self.vote_combiner = nn.Linear(num_classes * 2, num_classes) # two sets of logits
|
||||
|
||||
def forward(self, out_img, out_md):
|
||||
votes = torch.cat([out_img, out_md], dim=1)
|
||||
return self.vote_combiner(votes)
|
||||
@@ -1,276 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, Iterable, List, Optional
|
||||
|
||||
import json
|
||||
|
||||
from classes.v2.papila_data import PapilaData
|
||||
|
||||
|
||||
@dataclass
|
||||
class ImportSpec:
|
||||
id: str
|
||||
class_name: str
|
||||
params: Dict[str, Any]
|
||||
|
||||
|
||||
@dataclass
|
||||
class DataSourceSpec:
|
||||
node_id: str
|
||||
label: str
|
||||
output_type: str
|
||||
source: Optional[Dict[str, Any]]
|
||||
source_ref: Optional[Dict[str, Any]]
|
||||
|
||||
|
||||
@dataclass
|
||||
class TransformSpec:
|
||||
node_id: str
|
||||
label: str
|
||||
transform_type: str
|
||||
params: Dict[str, Any]
|
||||
|
||||
|
||||
@dataclass
|
||||
class LoaderSpec:
|
||||
node_id: str
|
||||
label: str
|
||||
input_type: str
|
||||
input_index: str
|
||||
input_key: str
|
||||
output_key: str
|
||||
transforms: List[TransformSpec]
|
||||
data_source: Optional[DataSourceSpec]
|
||||
|
||||
|
||||
@dataclass
|
||||
class TowerSpec:
|
||||
node_id: str
|
||||
label: str
|
||||
tower_type: str
|
||||
params: Dict[str, Any]
|
||||
|
||||
|
||||
@dataclass
|
||||
class BridgeSpec:
|
||||
node_id: str
|
||||
label: str
|
||||
method: str
|
||||
params: Dict[str, Any]
|
||||
|
||||
|
||||
@dataclass
|
||||
class ClassifierSpec:
|
||||
node_id: str
|
||||
label: str
|
||||
|
||||
|
||||
@dataclass
|
||||
class ConfigAssembly:
|
||||
raw: Dict[str, Any]
|
||||
imports: Dict[str, ImportSpec]
|
||||
data_sources: Dict[str, DataSourceSpec]
|
||||
transforms: Dict[str, TransformSpec]
|
||||
loaders: Dict[str, LoaderSpec]
|
||||
towers: Dict[str, TowerSpec]
|
||||
bridges: Dict[str, BridgeSpec]
|
||||
classifiers: Dict[str, ClassifierSpec]
|
||||
|
||||
|
||||
def load_config(path: Path) -> Dict[str, Any]:
|
||||
payload = json.loads(Path(path).read_text())
|
||||
if not isinstance(payload, dict):
|
||||
raise ValueError("Config JSON must be an object.")
|
||||
return payload
|
||||
|
||||
|
||||
def assemble_config(path: Path) -> ConfigAssembly:
|
||||
config = load_config(path)
|
||||
meta = config.get("meta", {})
|
||||
imports = _build_imports(meta.get("imports", []))
|
||||
nodes = {node["id"]: node for node in config.get("nodes", [])}
|
||||
edges = config.get("edges", [])
|
||||
|
||||
data_sources: Dict[str, DataSourceSpec] = {}
|
||||
transforms: Dict[str, TransformSpec] = {}
|
||||
loaders: Dict[str, LoaderSpec] = {}
|
||||
towers: Dict[str, TowerSpec] = {}
|
||||
bridges: Dict[str, BridgeSpec] = {}
|
||||
classifiers: Dict[str, ClassifierSpec] = {}
|
||||
|
||||
for node in nodes.values():
|
||||
ntype = node.get("type")
|
||||
if ntype == "data":
|
||||
data_sources[node["id"]] = DataSourceSpec(
|
||||
node_id=node["id"],
|
||||
label=node.get("label", ""),
|
||||
output_type=node.get("outputType", ""),
|
||||
source=node.get("source"),
|
||||
source_ref=node.get("sourceRef"),
|
||||
)
|
||||
elif ntype == "transform":
|
||||
transforms[node["id"]] = TransformSpec(
|
||||
node_id=node["id"],
|
||||
label=node.get("label", ""),
|
||||
transform_type=node.get("transformType", ""),
|
||||
params=_extract_transform_params(node),
|
||||
)
|
||||
elif ntype == "loader":
|
||||
loaders[node["id"]] = LoaderSpec(
|
||||
node_id=node["id"],
|
||||
label=node.get("label", ""),
|
||||
input_type=node.get("inputType", ""),
|
||||
input_index=node.get("inputIndex", ""),
|
||||
input_key=node.get("inputKey", ""),
|
||||
output_key=node.get("outputKey", ""),
|
||||
transforms=[],
|
||||
data_source=None,
|
||||
)
|
||||
elif ntype in ("image_tower", "metadata_tower"):
|
||||
towers[node["id"]] = TowerSpec(
|
||||
node_id=node["id"],
|
||||
label=node.get("label", ""),
|
||||
tower_type=node.get("towerType", "image" if ntype == "image_tower" else "metadata"),
|
||||
params=_extract_tower_params(node),
|
||||
)
|
||||
elif ntype == "bridge":
|
||||
bridges[node["id"]] = BridgeSpec(
|
||||
node_id=node["id"],
|
||||
label=node.get("label", ""),
|
||||
method=node.get("bridgeMethod", "fusion"),
|
||||
params=_extract_bridge_params(node),
|
||||
)
|
||||
elif ntype == "classifier":
|
||||
classifiers[node["id"]] = ClassifierSpec(
|
||||
node_id=node["id"],
|
||||
label=node.get("label", ""),
|
||||
)
|
||||
|
||||
# attach transforms + data sources to loaders by walking upstream
|
||||
for loader_id, loader in loaders.items():
|
||||
chain = _upstream_chain(loader_id, nodes, edges)
|
||||
for node_id in reversed(chain):
|
||||
if node_id in transforms:
|
||||
loader.transforms.append(transforms[node_id])
|
||||
if node_id in data_sources:
|
||||
loader.data_source = data_sources[node_id]
|
||||
|
||||
return ConfigAssembly(
|
||||
raw=config,
|
||||
imports=imports,
|
||||
data_sources=data_sources,
|
||||
transforms=transforms,
|
||||
loaders=loaders,
|
||||
towers=towers,
|
||||
bridges=bridges,
|
||||
classifiers=classifiers,
|
||||
)
|
||||
|
||||
|
||||
def resolve_imports(assembly: ConfigAssembly) -> Dict[str, Any]:
|
||||
resolved: Dict[str, Any] = {}
|
||||
for import_id, spec in assembly.imports.items():
|
||||
if spec.class_name == "PapilaData":
|
||||
params = spec.params
|
||||
resolved[import_id] = PapilaData.from_dirs(
|
||||
image_dir=params.get("image_dir", "Papila/FundusImages"),
|
||||
clinical_dir=params.get("clinical_dir", "Papila/ClinicalData"),
|
||||
label_col=params.get("label_col", "Diagnosis"),
|
||||
cat_cols=params.get("cat_cols", ["Gender", "Phakic/Pseudophakic"]),
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Unsupported import class {spec.class_name!r}")
|
||||
return resolved
|
||||
|
||||
|
||||
def _build_imports(entries: Iterable[Dict[str, Any]]) -> Dict[str, ImportSpec]:
|
||||
specs: Dict[str, ImportSpec] = {}
|
||||
for entry in entries or []:
|
||||
import_id = entry.get("id")
|
||||
if not import_id:
|
||||
continue
|
||||
specs[import_id] = ImportSpec(
|
||||
id=import_id,
|
||||
class_name=entry.get("className", ""),
|
||||
params=entry.get("params", {}) or {},
|
||||
)
|
||||
return specs
|
||||
|
||||
|
||||
def _extract_transform_params(node: Dict[str, Any]) -> Dict[str, Any]:
|
||||
return {
|
||||
"transformType": node.get("transformType"),
|
||||
"roiMaskSource": node.get("roiMaskSource"),
|
||||
"roiScale": node.get("roiScale"),
|
||||
"roiTargetSize": node.get("roiTargetSize"),
|
||||
"roiFallback": node.get("roiFallback"),
|
||||
"centerCropSize": node.get("centerCropSize"),
|
||||
"jitterHFlip": node.get("jitterHFlip"),
|
||||
"jitterVFlip": node.get("jitterVFlip"),
|
||||
"jitterRotation": node.get("jitterRotation"),
|
||||
"jitterColorEnabled": node.get("jitterColorEnabled"),
|
||||
"jitterColor": node.get("jitterColor"),
|
||||
"resizeSize": node.get("resizeSize"),
|
||||
}
|
||||
|
||||
|
||||
def _extract_tower_params(node: Dict[str, Any]) -> Dict[str, Any]:
|
||||
if node.get("towerType") == "metadata":
|
||||
return {
|
||||
"hidden_dim": node.get("mdHiddenDim"),
|
||||
"dropout": node.get("mdDropout"),
|
||||
"use_se": node.get("mdUseSe"),
|
||||
"se_reduction": node.get("mdSeReduction"),
|
||||
"se_pre_norm": node.get("mdSePreNorm"),
|
||||
"freeze_ratio": node.get("mdFreezeRatio"),
|
||||
}
|
||||
return {
|
||||
"backbone": node.get("imageBackbone"),
|
||||
"freeze_ratio": node.get("imageFreezeRatio"),
|
||||
"augment": node.get("imageAugment"),
|
||||
"geometry_dim": node.get("imageGeometryDim"),
|
||||
"use_se": node.get("imageUseSe"),
|
||||
"se_reduction": node.get("imageSeReduction"),
|
||||
"se_pre_norm": node.get("imageSePreNorm"),
|
||||
}
|
||||
|
||||
|
||||
def _extract_bridge_params(node: Dict[str, Any]) -> Dict[str, Any]:
|
||||
return {
|
||||
"fusion_dim": node.get("bridgeFusionDim"),
|
||||
"use_se": node.get("bridgeUseSe"),
|
||||
"se_reduction": node.get("bridgeSeReduction"),
|
||||
"se_pre_norm": node.get("bridgeSePreNorm"),
|
||||
}
|
||||
|
||||
|
||||
def _edge_from(edge: Dict[str, Any]) -> Optional[str]:
|
||||
return edge.get("from") or edge.get("source")
|
||||
|
||||
|
||||
def _edge_to(edge: Dict[str, Any]) -> Optional[str]:
|
||||
return edge.get("to") or edge.get("target")
|
||||
|
||||
|
||||
def _upstream_chain(start_id: str, nodes: Dict[str, Dict[str, Any]], edges: List[Dict[str, Any]]) -> List[str]:
|
||||
chain: List[str] = []
|
||||
visited = set()
|
||||
current = start_id
|
||||
while True:
|
||||
if current in visited:
|
||||
break
|
||||
visited.add(current)
|
||||
incoming = [edge for edge in edges if _edge_to(edge) == current]
|
||||
if not incoming:
|
||||
break
|
||||
# prefer first incoming edge for now
|
||||
current = _edge_from(incoming[0])
|
||||
if not current:
|
||||
break
|
||||
chain.append(current)
|
||||
node = nodes.get(current)
|
||||
if node and node.get("type") == "data":
|
||||
break
|
||||
return chain
|
||||
@@ -1,418 +0,0 @@
|
||||
"""Optic-disc image croppers and preprocessor factory for V2."""
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
from typing import Dict, Optional, Tuple
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import torch
|
||||
from PIL import Image, ImageDraw
|
||||
from torchvision import transforms
|
||||
|
||||
from classes.v2.geometry_features import compute_geometry_features, disc_cup_from_mask_image
|
||||
from classes.v2.unet_segmenter import UNetSegmenter
|
||||
|
||||
|
||||
def _geometry_from_mask(mask: np.ndarray, scale: float) -> Dict:
|
||||
mask = np.asarray(mask) > 0
|
||||
coords = np.argwhere(mask)
|
||||
if coords.size == 0:
|
||||
raise RuntimeError("Empty mask; cannot derive geometry")
|
||||
ys, xs = coords[:, 0], coords[:, 1]
|
||||
centre_x = float(xs.mean())
|
||||
centre_y = float(ys.mean())
|
||||
width = float(xs.max() - xs.min())
|
||||
height = float(ys.max() - ys.min())
|
||||
diameter = max(width, height)
|
||||
radius = diameter / 2.0
|
||||
crop_radius = radius * scale
|
||||
return {
|
||||
"centre_x": centre_x,
|
||||
"centre_y": centre_y,
|
||||
"radius": radius,
|
||||
"crop_radius": crop_radius,
|
||||
"crop_size": crop_radius * 2.0,
|
||||
}
|
||||
|
||||
|
||||
class UNetImageCropper:
|
||||
def __init__(
|
||||
self,
|
||||
manifest_path: Path,
|
||||
weights_path: Path,
|
||||
normalize: str = "per_image",
|
||||
threshold: float = 0.5,
|
||||
tta: bool = False,
|
||||
scale: float = 2.5,
|
||||
target_size: int = 224,
|
||||
cache_dir: Optional[Path] = None,
|
||||
) -> None:
|
||||
self.segmenter = UNetSegmenter(
|
||||
manifest_path=manifest_path,
|
||||
normalize=normalize,
|
||||
)
|
||||
state = torch.load(weights_path, map_location=self.segmenter.device)
|
||||
state_dict = state.get("model", state)
|
||||
self.segmenter.model.load_state_dict(state_dict)
|
||||
self.segmenter.model.to(self.segmenter.device)
|
||||
self.segmenter.model.eval()
|
||||
|
||||
self.threshold = threshold
|
||||
self.tta = tta
|
||||
self.scale = scale
|
||||
self.target_size = target_size
|
||||
self.cache_dir = Path(cache_dir) if cache_dir is not None else None
|
||||
if self.cache_dir is not None:
|
||||
self.cache_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
self.to_tensor = transforms.ToTensor()
|
||||
|
||||
def _cache_path(self, image_path: Path) -> Optional[Path]:
|
||||
if self.cache_dir is None:
|
||||
return None
|
||||
stem = image_path.stem
|
||||
return self.cache_dir / f"{stem}_s{int(self.scale * 100)}.npz"
|
||||
|
||||
def clear_cache(self) -> None:
|
||||
if self.cache_dir is None or not self.cache_dir.exists():
|
||||
return
|
||||
removed = sum(1 for f in self.cache_dir.glob("*.npz") if f.unlink() or True)
|
||||
print(f"[UNetImageCropper] Cleared {removed} cached crop files from {self.cache_dir}")
|
||||
|
||||
def _infer_masks(self, image: Image.Image) -> Optional[Tuple[np.ndarray, np.ndarray]]:
|
||||
resized = self.segmenter.preprocess_image(image)
|
||||
tensor = self.segmenter._normalize_tensor(
|
||||
self.to_tensor(resized).to(self.segmenter.device)
|
||||
).unsqueeze(0)
|
||||
|
||||
with torch.no_grad():
|
||||
logits = self.segmenter.model(tensor)
|
||||
if self.tta:
|
||||
t_h = torch.flip(tensor, dims=[3])
|
||||
log_h = self.segmenter.model(t_h)
|
||||
log_h = torch.flip(log_h, dims=[3])
|
||||
t_v = torch.flip(tensor, dims=[2])
|
||||
log_v = self.segmenter.model(t_v)
|
||||
log_v = torch.flip(log_v, dims=[2])
|
||||
logits = (logits + log_h + log_v) / 3.0
|
||||
probs = torch.sigmoid(logits)[0].cpu().numpy()
|
||||
|
||||
disc_pred = (probs[0] > self.threshold).astype(np.uint8) * 255
|
||||
cup_pred = (probs[1] > self.threshold).astype(np.uint8) * 255
|
||||
disc_img = Image.fromarray(disc_pred, mode="L").resize(image.size, Image.NEAREST)
|
||||
disc_mask = np.array(disc_img, dtype=np.uint8)
|
||||
cup_img = Image.fromarray(cup_pred, mode="L").resize(image.size, Image.NEAREST)
|
||||
cup_mask = (np.array(cup_img, dtype=np.uint8) > 0).astype(np.uint8)
|
||||
cup_mask = (cup_mask > 0) & (disc_mask > 0)
|
||||
cup_mask = cup_mask.astype(np.uint8)
|
||||
disc_mask = (disc_mask > 0).astype(np.uint8)
|
||||
return disc_mask, cup_mask
|
||||
|
||||
def _compute_crop_info(self, image: Image.Image, image_path: Path) -> Optional[dict]:
|
||||
image_path = Path(image_path).resolve()
|
||||
cache_path = self._cache_path(image_path)
|
||||
cached_bounds = None
|
||||
if cache_path is not None and cache_path.exists():
|
||||
data = np.load(cache_path, allow_pickle=False)
|
||||
try:
|
||||
cached_bounds = {
|
||||
"left": float(data["left"]),
|
||||
"upper": float(data["upper"]),
|
||||
"right": float(data["right"]),
|
||||
"lower": float(data["lower"]),
|
||||
}
|
||||
if "features" in data.files:
|
||||
cached_bounds["features"] = data["features"].astype(np.float32)
|
||||
return cached_bounds
|
||||
except KeyError:
|
||||
cached_bounds = None
|
||||
|
||||
masks = self._infer_masks(image)
|
||||
if masks is None:
|
||||
return cached_bounds
|
||||
disc_mask, cup_mask = masks
|
||||
try:
|
||||
geom = _geometry_from_mask(disc_mask, self.scale)
|
||||
except Exception:
|
||||
return cached_bounds
|
||||
cx = geom["centre_x"]
|
||||
cy = geom["centre_y"]
|
||||
r = geom["crop_radius"]
|
||||
left = max(0.0, cx - r)
|
||||
upper = max(0.0, cy - r)
|
||||
right = min(float(image.width), cx + r)
|
||||
lower = min(float(image.height), cy + r)
|
||||
features = compute_geometry_features(disc_mask, cup_mask)
|
||||
|
||||
info = {
|
||||
"left": left,
|
||||
"upper": upper,
|
||||
"right": right,
|
||||
"lower": lower,
|
||||
"features": features,
|
||||
}
|
||||
if cache_path is not None:
|
||||
np.savez(
|
||||
cache_path,
|
||||
left=left,
|
||||
upper=upper,
|
||||
right=right,
|
||||
lower=lower,
|
||||
width=float(image.width),
|
||||
height=float(image.height),
|
||||
scale=self.scale,
|
||||
target_size=self.target_size,
|
||||
features=features,
|
||||
)
|
||||
return info
|
||||
|
||||
def __call__(self, image: Image.Image, image_path: Path) -> Image.Image:
|
||||
info = self._compute_crop_info(image, image_path)
|
||||
if info is None:
|
||||
return image
|
||||
left = info["left"]
|
||||
upper = info["upper"]
|
||||
right = info["right"]
|
||||
lower = info["lower"]
|
||||
if right <= left or lower <= upper:
|
||||
return image
|
||||
crop = image.crop((left, upper, right, lower))
|
||||
return crop.resize((self.target_size, self.target_size), Image.BILINEAR)
|
||||
|
||||
def geometry_features(self, image: Image.Image, image_path: Path) -> Optional[np.ndarray]:
|
||||
info = self._compute_crop_info(image, image_path)
|
||||
if info is None:
|
||||
return None
|
||||
features = info.get("features")
|
||||
if features is None:
|
||||
return None
|
||||
return np.asarray(features, dtype=np.float32)
|
||||
|
||||
|
||||
class ManifestImageCropper:
|
||||
def __init__(
|
||||
self,
|
||||
manifest_path: Path,
|
||||
scale: float = 2.5,
|
||||
target_size: int = 224,
|
||||
cache_dir: Optional[Path] = None,
|
||||
) -> None:
|
||||
self.scale = scale
|
||||
self.target_size = target_size
|
||||
self.cache_dir = Path(cache_dir) if cache_dir is not None else None
|
||||
if self.cache_dir is not None:
|
||||
self.cache_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
df = pd.read_csv(manifest_path)
|
||||
self.entries: Dict[str, dict] = {}
|
||||
for _, row in df.iterrows():
|
||||
img_path = Path(row["image_path"]).resolve()
|
||||
self.entries[str(img_path)] = {
|
||||
"annotation_disc": row.get("annotation_disc"),
|
||||
"annotation_cup": row.get("annotation_cup"),
|
||||
"annotation_type_disc": row.get("annotation_type_disc"),
|
||||
"annotation_type_cup": row.get("annotation_type_cup"),
|
||||
}
|
||||
|
||||
def _cache_path(self, image_path: Path) -> Optional[Path]:
|
||||
if self.cache_dir is None:
|
||||
return None
|
||||
return self.cache_dir / f"{image_path.stem}_s{int(self.scale * 100)}.npz"
|
||||
|
||||
def clear_cache(self) -> None:
|
||||
if self.cache_dir is None or not self.cache_dir.exists():
|
||||
return
|
||||
removed = sum(1 for f in self.cache_dir.glob("*.npz") if f.unlink() or True)
|
||||
print(f"[ManifestImageCropper] Cleared {removed} cached crop files from {self.cache_dir}")
|
||||
|
||||
@staticmethod
|
||||
def _load_contour(path: Path) -> np.ndarray:
|
||||
coords = np.loadtxt(path)
|
||||
if coords.ndim == 1:
|
||||
coords = coords.reshape(-1, 2)
|
||||
return coords
|
||||
|
||||
@staticmethod
|
||||
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)
|
||||
|
||||
def _load_masks(self, entry: dict, image: Image.Image) -> Optional[Tuple[np.ndarray, np.ndarray]]:
|
||||
disc_path = entry.get("annotation_disc")
|
||||
cup_path = entry.get("annotation_cup")
|
||||
disc_type = (entry.get("annotation_type_disc") or "").lower()
|
||||
cup_type = (entry.get("annotation_type_cup") or "").lower()
|
||||
|
||||
disc_mask: Optional[np.ndarray] = None
|
||||
cup_mask: Optional[np.ndarray] = None
|
||||
|
||||
if disc_path and not pd.isna(disc_path):
|
||||
disc_path = Path(disc_path)
|
||||
try:
|
||||
if disc_type == "mask":
|
||||
mask_img = Image.open(disc_path)
|
||||
mask_img = mask_img.resize(image.size, Image.NEAREST)
|
||||
disc_mask, cup_from_mask = disc_cup_from_mask_image(mask_img)
|
||||
if cup_from_mask.sum() > 0:
|
||||
cup_mask = cup_from_mask
|
||||
elif disc_type == "contour":
|
||||
coords = self._load_contour(disc_path)
|
||||
disc_mask = self._contour_to_mask(coords, image.size)
|
||||
except Exception:
|
||||
disc_mask = None
|
||||
|
||||
if cup_mask is None and cup_path and not pd.isna(cup_path):
|
||||
cup_path = Path(cup_path)
|
||||
try:
|
||||
if cup_type == "mask":
|
||||
mask_img = Image.open(cup_path)
|
||||
mask_img = mask_img.resize(image.size, Image.NEAREST)
|
||||
_, cup_mask = disc_cup_from_mask_image(mask_img)
|
||||
elif cup_type == "contour":
|
||||
coords = self._load_contour(cup_path)
|
||||
cup_mask = self._contour_to_mask(coords, image.size)
|
||||
except Exception:
|
||||
cup_mask = None
|
||||
|
||||
if disc_mask is None:
|
||||
return None
|
||||
disc_mask = (disc_mask > 0).astype(np.uint8)
|
||||
if cup_mask is None:
|
||||
cup_mask = np.zeros_like(disc_mask, dtype=np.uint8)
|
||||
cup_mask = ((cup_mask > 0) & (disc_mask > 0)).astype(np.uint8)
|
||||
return disc_mask, cup_mask
|
||||
|
||||
def _compute_crop_info(self, image: Image.Image, image_path: Path) -> Optional[dict]:
|
||||
image_path = Path(image_path).resolve()
|
||||
entry = self.entries.get(str(image_path))
|
||||
if entry is None:
|
||||
return None
|
||||
cache_path = self._cache_path(image_path)
|
||||
cached_bounds = None
|
||||
if cache_path is not None and cache_path.exists():
|
||||
data = np.load(cache_path, allow_pickle=False)
|
||||
try:
|
||||
cached_bounds = {
|
||||
"left": float(data["left"]),
|
||||
"upper": float(data["upper"]),
|
||||
"right": float(data["right"]),
|
||||
"lower": float(data["lower"]),
|
||||
}
|
||||
if "features" in data.files:
|
||||
cached_bounds["features"] = data["features"].astype(np.float32)
|
||||
return cached_bounds
|
||||
except KeyError:
|
||||
cached_bounds = None
|
||||
|
||||
masks = self._load_masks(entry, image)
|
||||
if masks is None:
|
||||
return cached_bounds
|
||||
disc_mask, cup_mask = masks
|
||||
try:
|
||||
geom = _geometry_from_mask(disc_mask, self.scale)
|
||||
except Exception:
|
||||
return cached_bounds
|
||||
cx = geom["centre_x"]
|
||||
cy = geom["centre_y"]
|
||||
r = geom["crop_radius"]
|
||||
left = max(0.0, cx - r)
|
||||
upper = max(0.0, cy - r)
|
||||
right = min(float(image.width), cx + r)
|
||||
lower = min(float(image.height), cy + r)
|
||||
features = compute_geometry_features(disc_mask, cup_mask)
|
||||
|
||||
info = {
|
||||
"left": left,
|
||||
"upper": upper,
|
||||
"right": right,
|
||||
"lower": lower,
|
||||
"features": features,
|
||||
}
|
||||
if cache_path is not None:
|
||||
np.savez(
|
||||
cache_path,
|
||||
left=left,
|
||||
upper=upper,
|
||||
right=right,
|
||||
lower=lower,
|
||||
width=float(image.width),
|
||||
height=float(image.height),
|
||||
scale=self.scale,
|
||||
target_size=self.target_size,
|
||||
features=features,
|
||||
)
|
||||
return info
|
||||
|
||||
def __call__(self, image: Image.Image, image_path: Path) -> Image.Image:
|
||||
info = self._compute_crop_info(image, image_path)
|
||||
if info is None:
|
||||
return image
|
||||
left = info["left"]
|
||||
upper = info["upper"]
|
||||
right = info["right"]
|
||||
lower = info["lower"]
|
||||
if right <= left or lower <= upper:
|
||||
return image
|
||||
crop = image.crop((left, upper, right, lower))
|
||||
return crop.resize((self.target_size, self.target_size), Image.BILINEAR)
|
||||
|
||||
def geometry_features(self, image: Image.Image, image_path: Path) -> Optional[np.ndarray]:
|
||||
info = self._compute_crop_info(image, image_path)
|
||||
if info is None:
|
||||
return None
|
||||
features = info.get("features")
|
||||
if features is None:
|
||||
return None
|
||||
return np.asarray(features, dtype=np.float32)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Factory
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def build_image_preprocessor_from_args(args):
|
||||
"""Construct the correct image cropper from CLI args, or return None."""
|
||||
crop_manifest = getattr(args, "img_crop_manifest", None)
|
||||
crop_weights = getattr(args, "img_crop_weights", None)
|
||||
use_gt = bool(getattr(args, "img_crop_gt", False))
|
||||
if not crop_manifest:
|
||||
return None
|
||||
crop_cache = Path(getattr(args, "img_crop_cache", Path("cache_data/hypertower_crops")))
|
||||
persist_cache = bool(getattr(args, "persist_img_crop_cache", False))
|
||||
if use_gt:
|
||||
pre = ManifestImageCropper(
|
||||
manifest_path=Path(crop_manifest),
|
||||
scale=getattr(args, "img_crop_scale", 2.5),
|
||||
target_size=getattr(args, "img_crop_size", 224),
|
||||
cache_dir=crop_cache,
|
||||
)
|
||||
if not persist_cache:
|
||||
pre.clear_cache()
|
||||
print(f"[V2 modes] GT disc cropper enabled -> cache at {crop_cache}", flush=True)
|
||||
return pre
|
||||
if crop_weights:
|
||||
pre = UNetImageCropper(
|
||||
manifest_path=Path(crop_manifest),
|
||||
weights_path=Path(crop_weights),
|
||||
normalize=getattr(args, "img_crop_normalize", "per_image"),
|
||||
threshold=getattr(args, "img_crop_threshold", 0.5),
|
||||
tta=getattr(args, "img_crop_tta", False),
|
||||
scale=getattr(args, "img_crop_scale", 2.5),
|
||||
target_size=getattr(args, "img_crop_size", 224),
|
||||
cache_dir=crop_cache,
|
||||
)
|
||||
if not persist_cache:
|
||||
pre.clear_cache()
|
||||
print(f"[V2 modes] UNet disc cropper enabled -> cache at {crop_cache}", flush=True)
|
||||
return pre
|
||||
print(
|
||||
"[V2 modes] img_crop_manifest provided but no --img-crop-gt or --img-crop-weights; cropping disabled.",
|
||||
flush=True,
|
||||
)
|
||||
return None
|
||||
@@ -1,241 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
from typing import Callable, Dict, Iterable, List, Optional, Tuple
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
|
||||
|
||||
class DataBundle:
|
||||
"""
|
||||
Generic, torch-free container for metadata and file/label bookkeeping.
|
||||
|
||||
Keeps feature typing, vectorization, and patient-level splits generic.
|
||||
Dataset-specific preprocessing (e.g., eye canonicalization) should live
|
||||
in the dataset builder (e.g., papila_builders in v2).
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
image_dir: str,
|
||||
clinical_dir: Optional[str] = None,
|
||||
label_col: str,
|
||||
patient_col: str = "Patient ID",
|
||||
cat_cols: Optional[Iterable[str]] = None,
|
||||
max_unique_for_cat: int = 4,
|
||||
n_splits: int = 5,
|
||||
random_seed: int = 42,
|
||||
filename_template: str = "RET{pid:03d}{eye}.jpg",
|
||||
image_path_fn: Optional[Callable[[pd.Series], Path]] = None,
|
||||
) -> None:
|
||||
self.image_dir = Path(image_dir)
|
||||
self.label_col = label_col
|
||||
self.patient_col = patient_col
|
||||
self.max_unique_for_cat = max_unique_for_cat
|
||||
self.n_splits = n_splits
|
||||
self.filename_template = filename_template
|
||||
self.image_path_fn = image_path_fn
|
||||
self.clinical_dir = Path(clinical_dir) if clinical_dir else None
|
||||
|
||||
# Internal state
|
||||
self.frames: List[pd.DataFrame] = []
|
||||
self.df: pd.DataFrame = pd.DataFrame()
|
||||
self.scalar_cols: List[str] = []
|
||||
self.cat_cols: List[str] = list(cat_cols) if cat_cols is not None else []
|
||||
self.scalar_stats: Dict[str, Dict[str, float]] = {}
|
||||
self.cat_maps: Dict[str, Dict[object, int]] = {}
|
||||
self.feature_dim: int = 0
|
||||
self.folds: Dict[int, Dict[str, List[object]]] = {}
|
||||
self.random_seed = int(random_seed)
|
||||
|
||||
# ------------------- Public API -------------------
|
||||
def add_df(
|
||||
self,
|
||||
df: pd.DataFrame,
|
||||
*,
|
||||
id_column: Optional[str] = None,
|
||||
exclude_cols: Optional[Iterable[str]] = None,
|
||||
) -> None:
|
||||
"""
|
||||
Add a dataframe and re-run typing, stats, and K-fold indices.
|
||||
QC rules:
|
||||
- Must have patient ID column; if not provided under that name, specify id_column.
|
||||
"""
|
||||
df = df.copy()
|
||||
self._ensure_patient_id(df, id_column)
|
||||
if self.label_col not in df.columns:
|
||||
raise ValueError(f"label_col '{self.label_col}' not found in added dataframe")
|
||||
|
||||
self.frames.append(df)
|
||||
self._refresh_master_df(exclude_cols=exclude_cols)
|
||||
self._infer_or_validate_feature_types(exclude_cols=exclude_cols)
|
||||
self._compute_numeric_stats()
|
||||
self._build_cat_maps()
|
||||
self._compute_feature_dim()
|
||||
self._build_kfold_indices()
|
||||
|
||||
def get_split_ids(self, fold: int) -> Tuple[List[object], List[object]]:
|
||||
rec = self.folds.get(fold)
|
||||
if not rec:
|
||||
raise KeyError(f"Fold {fold} not available. Built folds: {sorted(self.folds.keys())}")
|
||||
return rec["train_ids"], rec["test_ids"]
|
||||
|
||||
def get_split_dfs(self, fold: int) -> Tuple[pd.DataFrame, pd.DataFrame]:
|
||||
train_ids, test_ids = self.get_split_ids(fold)
|
||||
train_df = self.df[self.df[self.patient_col].isin(train_ids)].reset_index(drop=True)
|
||||
test_df = self.df[self.df[self.patient_col].isin(test_ids)].reset_index(drop=True)
|
||||
return train_df, test_df
|
||||
|
||||
def vectorize_row(self, row: pd.Series) -> np.ndarray:
|
||||
"""Return a numpy feature vector (torch-free)."""
|
||||
feats: List[float] = []
|
||||
miss: List[float] = []
|
||||
# numeric
|
||||
for col in self.scalar_cols:
|
||||
v = pd.to_numeric(row.get(col), errors="coerce")
|
||||
if pd.isna(v):
|
||||
miss.append(1.0)
|
||||
v = self.scalar_stats[col]["median"]
|
||||
else:
|
||||
miss.append(0.0)
|
||||
lo = self.scalar_stats[col]["min"]
|
||||
hi = self.scalar_stats[col]["max"]
|
||||
feats.append((float(v) - lo) / (hi - lo) if hi > lo else 0.0)
|
||||
# categorical
|
||||
for col in self.cat_cols:
|
||||
mapping = self.cat_maps[col]
|
||||
one = [0.0] * len(mapping)
|
||||
key = row.get(col)
|
||||
one[mapping.get(key, 0)] = 1.0 # 0 is <UNK>
|
||||
feats.extend(one)
|
||||
# numeric missing flags
|
||||
feats.extend(miss)
|
||||
return np.asarray(feats, dtype=np.float32)
|
||||
|
||||
def get_image_path(self, row: pd.Series) -> Path:
|
||||
if self.image_path_fn is not None:
|
||||
return Path(self.image_path_fn(row))
|
||||
pid = int(row[self.patient_col])
|
||||
eye = row.get("eyeID", "")
|
||||
if eye in ("OS", "OD"):
|
||||
eye_str = eye
|
||||
else:
|
||||
eye_str = str(eye)
|
||||
return self.image_dir / self.filename_template.format(pid=pid, eye=eye_str)
|
||||
|
||||
def encode_metadata(self, row: pd.Series) -> np.ndarray:
|
||||
return self.vectorize_row(row)
|
||||
|
||||
def get_label(self, row: pd.Series) -> int:
|
||||
return int(row[self.label_col])
|
||||
|
||||
# ------------------- Internal helpers -------------------
|
||||
def _ensure_patient_id(self, df: pd.DataFrame, id_column: Optional[str]) -> None:
|
||||
if self.patient_col in df.columns:
|
||||
return
|
||||
if id_column and id_column in df.columns:
|
||||
df.rename(columns={id_column: self.patient_col}, inplace=True)
|
||||
return
|
||||
candidates = [
|
||||
c
|
||||
for c in df.columns
|
||||
if c.lower().replace(" ", "") in {"patientid", "patient", "pid"}
|
||||
]
|
||||
if len(candidates) == 1:
|
||||
df.rename(columns={candidates[0]: self.patient_col}, inplace=True)
|
||||
return
|
||||
raise ValueError(
|
||||
f"A '{self.patient_col}' column is required; provide id_column=... if it has a different name."
|
||||
)
|
||||
|
||||
def _refresh_master_df(self, exclude_cols: Optional[Iterable[str]] = None) -> None:
|
||||
self.df = pd.concat(self.frames, axis=0, ignore_index=True)
|
||||
if exclude_cols:
|
||||
self.df = self.df.drop(columns=[c for c in exclude_cols if c in self.df.columns])
|
||||
|
||||
def _infer_or_validate_feature_types(self, exclude_cols: Optional[Iterable[str]] = None) -> None:
|
||||
excluded = set(exclude_cols or []) | {self.label_col, self.patient_col}
|
||||
feature_candidates = [c for c in self.df.columns if c not in excluded]
|
||||
cats = set(self.cat_cols) if self.cat_cols else set()
|
||||
scalars = set()
|
||||
for c in feature_candidates:
|
||||
if c in cats:
|
||||
continue
|
||||
s = self.df[c]
|
||||
as_num = pd.to_numeric(s, errors="coerce")
|
||||
num_missing = as_num.isna().mean()
|
||||
num_unique = s.dropna().nunique()
|
||||
if as_num.notna().any() and num_missing < 1.0 and num_unique > self.max_unique_for_cat:
|
||||
scalars.add(c)
|
||||
else:
|
||||
if num_unique <= self.max_unique_for_cat or as_num.isna().mean() > 0.0:
|
||||
cats.add(c)
|
||||
else:
|
||||
scalars.add(c)
|
||||
self.cat_cols = sorted(cats)
|
||||
self.scalar_cols = sorted(scalars)
|
||||
|
||||
def _compute_numeric_stats(self) -> None:
|
||||
self.scalar_stats.clear()
|
||||
for col in self.scalar_cols:
|
||||
s = pd.to_numeric(self.df[col], errors="coerce")
|
||||
vals = s.dropna().astype(float).values
|
||||
if vals.size == 0:
|
||||
lo, hi, med = 0.0, 1.0, 0.0
|
||||
else:
|
||||
lo, hi = float(np.min(vals)), float(np.max(vals))
|
||||
med = float(np.median(vals))
|
||||
if hi <= lo:
|
||||
hi = lo + 1.0
|
||||
self.scalar_stats[col] = {"min": lo, "max": hi, "median": med}
|
||||
|
||||
def _build_cat_maps(self) -> None:
|
||||
self.cat_maps.clear()
|
||||
for col in self.cat_cols:
|
||||
cats = [v for v in self.df[col].dropna().unique().tolist()]
|
||||
try:
|
||||
cats = sorted(cats)
|
||||
except Exception:
|
||||
pass
|
||||
mapping = {"<UNK>": 0}
|
||||
for i, v in enumerate(cats, start=1):
|
||||
mapping[v] = i
|
||||
self.cat_maps[col] = mapping
|
||||
|
||||
def _compute_feature_dim(self) -> None:
|
||||
self.feature_dim = len(self.scalar_cols) + sum(len(m) for m in self.cat_maps.values()) + len(self.scalar_cols)
|
||||
|
||||
# ------------------- K-fold on unique patients -------------------
|
||||
def _build_kfold_indices(self) -> None:
|
||||
pats = self.df[self.patient_col].unique().tolist()
|
||||
labels_by_pat: Dict[object, object] = {}
|
||||
for pid, grp in self.df.groupby(self.patient_col):
|
||||
lab = grp[self.label_col].dropna()
|
||||
if len(lab) == 0:
|
||||
labels_by_pat[pid] = 0
|
||||
else:
|
||||
labels_by_pat[pid] = lab.mode().iloc[0]
|
||||
y_pat = np.array([labels_by_pat[p] for p in pats])
|
||||
|
||||
try:
|
||||
from sklearn.model_selection import StratifiedGroupKFold
|
||||
|
||||
sgkf = StratifiedGroupKFold(
|
||||
n_splits=self.n_splits, shuffle=True, random_state=self.random_seed
|
||||
)
|
||||
split_iter = sgkf.split(X=pats, y=y_pat, groups=pats)
|
||||
except Exception:
|
||||
from sklearn.model_selection import StratifiedKFold
|
||||
|
||||
skf = StratifiedKFold(
|
||||
n_splits=self.n_splits, shuffle=True, random_state=self.random_seed
|
||||
)
|
||||
split_iter = skf.split(X=np.zeros(len(pats)), y=y_pat)
|
||||
|
||||
self.folds.clear()
|
||||
for i, (train_idx, test_idx) in enumerate(split_iter):
|
||||
train_ids = [pats[j] for j in train_idx]
|
||||
test_ids = [pats[j] for j in test_idx]
|
||||
self.folds[i] = {"train_ids": train_ids, "test_ids": test_ids}
|
||||
@@ -1,115 +0,0 @@
|
||||
from torch.utils.data import Dataset
|
||||
from PIL import Image
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
|
||||
class ClinicalDataset(Dataset):
|
||||
"""Generic dataset wrapping a DataBundle-like instance.
|
||||
Returns (img_tensor, meta_tensor, label)."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
clinical_data,
|
||||
img_transform,
|
||||
meta_transform=None,
|
||||
image_preprocessor=None,
|
||||
geometry_provider=None,
|
||||
geometry_dim: int = 0,
|
||||
image_cache: "dict | None" = None,
|
||||
):
|
||||
self.clinical = clinical_data
|
||||
self.transform_image = img_transform
|
||||
self.meta_transform = meta_transform or (lambda x: x)
|
||||
self.image_preprocessor = image_preprocessor
|
||||
self.geometry_provider = geometry_provider
|
||||
self.geometry_dim = geometry_dim if geometry_provider is not None else 0
|
||||
self.image_cache = image_cache
|
||||
|
||||
def __len__(self):
|
||||
return len(self.clinical.df)
|
||||
|
||||
def __getitem__(self, idx: int):
|
||||
row = self.clinical.df.iloc[idx]
|
||||
# load & transform image
|
||||
img_path = self.clinical.get_image_path(row)
|
||||
cache_key = str(img_path)
|
||||
if self.image_cache is not None and cache_key in self.image_cache:
|
||||
orig_img = Image.fromarray(self.image_cache[cache_key])
|
||||
else:
|
||||
orig_img = Image.open(img_path).convert("RGB")
|
||||
if self.image_cache is not None:
|
||||
self.image_cache[cache_key] = np.asarray(orig_img, dtype=np.uint8)
|
||||
img = orig_img
|
||||
if self.image_preprocessor is not None:
|
||||
img = self.image_preprocessor(img, img_path)
|
||||
img_t = self.transform_image(img)
|
||||
# encode & transform metadata
|
||||
meta = self.clinical.encode_metadata(row)
|
||||
meta_t = self.meta_transform(meta)
|
||||
# label
|
||||
label = self.clinical.get_label(row)
|
||||
if self.geometry_dim > 0:
|
||||
features = None
|
||||
if self.geometry_provider is not None and hasattr(self.geometry_provider, "geometry_features"):
|
||||
features = self.geometry_provider.geometry_features(orig_img, img_path)
|
||||
if features is None:
|
||||
geom_vec = torch.zeros(self.geometry_dim, dtype=torch.float32)
|
||||
else:
|
||||
features = np.asarray(features, dtype=np.float32)
|
||||
if features.shape[0] != self.geometry_dim:
|
||||
geom_vec = torch.zeros(self.geometry_dim, dtype=torch.float32)
|
||||
else:
|
||||
geom_vec = torch.from_numpy(features)
|
||||
return img_t, meta_t, geom_vec, label
|
||||
return img_t, meta_t, label
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _ClinicalView — shim used by V2HyperTower._run_fold
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
from .data_bundle import DataBundle # noqa: E402
|
||||
|
||||
|
||||
class _ClinicalView:
|
||||
"""Minimal shim so ClinicalDataset can iterate an epoch-specific DataFrame
|
||||
while still delegating encoding/paths/labels to the DataBundle object."""
|
||||
|
||||
def __init__(self, base: DataBundle, df):
|
||||
self.base = base
|
||||
self.df = df
|
||||
|
||||
@property
|
||||
def image_dir(self):
|
||||
return self.base.image_dir
|
||||
|
||||
@property
|
||||
def clinical_dir(self):
|
||||
return self.base.clinical_dir
|
||||
|
||||
@property
|
||||
def id_cols(self):
|
||||
return ("Patient ID", "eyeID")
|
||||
|
||||
@property
|
||||
def label_col(self):
|
||||
return self.base.label_col
|
||||
|
||||
@property
|
||||
def filename_template(self):
|
||||
return getattr(self.base, "filename_template", "RET{pid:03d}{eye}.jpg")
|
||||
|
||||
@property
|
||||
def dim(self):
|
||||
return self.base.feature_dim
|
||||
|
||||
def encode_metadata(self, row):
|
||||
vec = self.base.vectorize_row(row)
|
||||
return torch.as_tensor(vec, dtype=torch.float32)
|
||||
|
||||
def get_image_path(self, row):
|
||||
return self.base.get_image_path(row)
|
||||
|
||||
def get_label(self, row):
|
||||
return int(row[self.base.label_col])
|
||||
@@ -1,119 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Iterable, List, Sequence, Tuple, Union
|
||||
import re
|
||||
|
||||
import pandas as pd
|
||||
|
||||
|
||||
@dataclass
|
||||
class RegexFilter:
|
||||
pattern: str
|
||||
flags: int = 0
|
||||
|
||||
def apply_paths(self, paths: Sequence[str]) -> Tuple[List[str], List[str]]:
|
||||
if not self.pattern:
|
||||
return list(paths), []
|
||||
try:
|
||||
regex = re.compile(self.pattern, self.flags)
|
||||
except re.error as err:
|
||||
return list(paths), [f'Invalid regex "{self.pattern}": {err}']
|
||||
filtered = [p for p in paths if regex.search(p)]
|
||||
return filtered, []
|
||||
|
||||
|
||||
@dataclass
|
||||
class ColumnFilter:
|
||||
column: str
|
||||
operator: str
|
||||
value: str
|
||||
case_insensitive: bool = True
|
||||
|
||||
def apply_df(self, df: pd.DataFrame) -> Tuple[pd.DataFrame, List[str]]:
|
||||
warnings: List[str] = []
|
||||
if not self.column:
|
||||
return df, ["Column filter missing column name."]
|
||||
columns = list(df.columns)
|
||||
col_index = _resolve_column_index(columns, self.column, warnings)
|
||||
if col_index is None:
|
||||
return df, warnings
|
||||
col_name = columns[col_index]
|
||||
if self.value is None or self.value == "":
|
||||
return df, [f'Column filter "{self.column}" missing value.']
|
||||
series = df[col_name]
|
||||
mask = series.apply(
|
||||
lambda cell: compare_cell(
|
||||
cell, self.value, self.operator, case_insensitive=self.case_insensitive
|
||||
)
|
||||
)
|
||||
return df[mask], warnings
|
||||
|
||||
|
||||
FilterSpec = Union[RegexFilter, ColumnFilter]
|
||||
|
||||
|
||||
def apply_regex_filters(paths: Sequence[str], filters: Iterable[RegexFilter]) -> Tuple[List[str], List[str]]:
|
||||
filtered = list(paths)
|
||||
warnings: List[str] = []
|
||||
for filt in filters:
|
||||
filtered, warn = filt.apply_paths(filtered)
|
||||
warnings.extend(warn)
|
||||
return filtered, warnings
|
||||
|
||||
|
||||
def apply_column_filters(df: pd.DataFrame, filters: Iterable[ColumnFilter]) -> Tuple[pd.DataFrame, List[str]]:
|
||||
filtered = df
|
||||
warnings: List[str] = []
|
||||
for filt in filters:
|
||||
filtered, warn = filt.apply_df(filtered)
|
||||
warnings.extend(warn)
|
||||
return filtered, warnings
|
||||
|
||||
|
||||
def compare_cell(cell, raw_value: str, operator: str, case_insensitive: bool = True) -> bool:
|
||||
cell_str = "" if cell is None else str(cell).strip()
|
||||
value_str = "" if raw_value is None else str(raw_value).strip()
|
||||
if case_insensitive:
|
||||
cell_str = cell_str.lower()
|
||||
value_str = value_str.lower()
|
||||
if operator == "=":
|
||||
return cell_str == value_str
|
||||
if operator == "!=":
|
||||
return cell_str != value_str
|
||||
cell_num = _to_float(cell_str)
|
||||
value_num = _to_float(value_str)
|
||||
if cell_num is None or value_num is None:
|
||||
return False
|
||||
if operator == ">":
|
||||
return cell_num > value_num
|
||||
if operator == ">=":
|
||||
return cell_num >= value_num
|
||||
if operator == "<":
|
||||
return cell_num < value_num
|
||||
if operator == "<=":
|
||||
return cell_num <= value_num
|
||||
return False
|
||||
|
||||
|
||||
def _resolve_column_index(columns: Sequence[str], column: str, warnings: List[str]) -> int | None:
|
||||
try:
|
||||
return columns.index(column)
|
||||
except ValueError:
|
||||
lower = column.lower()
|
||||
matches = [idx for idx, col in enumerate(columns) if str(col).lower() == lower]
|
||||
if matches:
|
||||
if len(matches) > 1:
|
||||
warnings.append(
|
||||
f'Column "{column}" matched multiple headers; using "{columns[matches[0]]}".'
|
||||
)
|
||||
return matches[0]
|
||||
warnings.append(f'Column "{column}" not found.')
|
||||
return None
|
||||
|
||||
|
||||
def _to_float(value: str) -> float | None:
|
||||
try:
|
||||
return float(value)
|
||||
except (TypeError, ValueError):
|
||||
return None
|
||||
@@ -1,87 +0,0 @@
|
||||
"""Shared helpers for deriving disc/cup geometry features."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections import Counter
|
||||
from typing import Tuple
|
||||
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
EPS = 1e-6
|
||||
FEATURE_DIM = 5
|
||||
|
||||
|
||||
def disc_cup_from_mask_image(mask_img: Image.Image) -> Tuple[np.ndarray, np.ndarray]:
|
||||
"""Return binary disc/cup masks from a REFUGE-style annotation image."""
|
||||
arr = np.asarray(mask_img)
|
||||
if arr.ndim == 3:
|
||||
h, w, c = arr.shape
|
||||
border = np.concatenate(
|
||||
[arr[0, :, :], arr[-1, :, :], arr[:, 0, :], arr[:, -1, :]],
|
||||
axis=0,
|
||||
)
|
||||
border_counts = Counter(map(tuple, border))
|
||||
bg_color = border_counts.most_common(1)[0][0]
|
||||
flat = arr.reshape(-1, c)
|
||||
colors = Counter(map(tuple, flat))
|
||||
colors.pop(bg_color, None)
|
||||
disc = (~np.all(arr == bg_color, axis=-1)).astype(np.uint8)
|
||||
if colors:
|
||||
cup_color = min(colors.keys(), key=lambda col: sum(col))
|
||||
cup = np.all(arr == cup_color, axis=-1).astype(np.uint8)
|
||||
else:
|
||||
cup = np.zeros((h, w), dtype=np.uint8)
|
||||
else:
|
||||
border = np.concatenate([arr[0, :], arr[-1, :], arr[:, 0], arr[:, -1]])
|
||||
counts = Counter(border.tolist())
|
||||
bg_value = counts.most_common(1)[0][0]
|
||||
disc = (arr != bg_value).astype(np.uint8)
|
||||
fg = arr[arr != bg_value]
|
||||
if fg.size > 0:
|
||||
cup_value = int(np.min(fg))
|
||||
cup = (arr == cup_value).astype(np.uint8)
|
||||
else:
|
||||
cup = np.zeros_like(arr, dtype=np.uint8)
|
||||
cup = (cup > 0) & (disc > 0)
|
||||
return disc.astype(np.uint8), cup.astype(np.uint8)
|
||||
|
||||
|
||||
def compute_geometry_features(disc_mask: np.ndarray, cup_mask: np.ndarray) -> np.ndarray:
|
||||
"""Compute cup/disc geometry descriptors (area, rim, diameter ratios, centre shift)."""
|
||||
disc = (disc_mask > 0).astype(np.float32)
|
||||
cup = (cup_mask > 0).astype(np.float32)
|
||||
|
||||
disc_area = disc.sum()
|
||||
cup_area = cup.sum()
|
||||
area_ratio = cup_area / (disc_area + EPS)
|
||||
rim_ratio = (disc_area - cup_area) / (disc_area + EPS)
|
||||
|
||||
disc_rows = np.any(disc > 0, axis=1)
|
||||
cup_rows = np.any(cup > 0, axis=1)
|
||||
disc_cols = np.any(disc > 0, axis=0)
|
||||
cup_cols = np.any(cup > 0, axis=0)
|
||||
|
||||
disc_height = float(disc_rows.sum())
|
||||
cup_height = float(cup_rows.sum())
|
||||
disc_width = float(disc_cols.sum())
|
||||
cup_width = float(cup_cols.sum())
|
||||
|
||||
vertical_ratio = cup_height / (disc_height + EPS)
|
||||
horizontal_ratio = cup_width / (disc_width + EPS)
|
||||
|
||||
def _centre(mask: np.ndarray) -> Tuple[float, float]:
|
||||
coords = np.argwhere(mask > 0)
|
||||
if coords.size == 0:
|
||||
return 0.5, 0.5
|
||||
ys, xs = coords[:, 0], coords[:, 1]
|
||||
return float(xs.mean()) / mask.shape[1], float(ys.mean()) / mask.shape[0]
|
||||
|
||||
disc_cx, disc_cy = _centre(disc)
|
||||
cup_cx, cup_cy = _centre(cup)
|
||||
centre_shift = float(np.hypot(cup_cx - disc_cx, cup_cy - disc_cy))
|
||||
|
||||
return np.array(
|
||||
[area_ratio, rim_ratio, vertical_ratio, horizontal_ratio, centre_shift],
|
||||
dtype=np.float32,
|
||||
)
|
||||
@@ -1,128 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import csv
|
||||
import json
|
||||
import logging
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
|
||||
DEFAULT_OPTIONAL_EPOCH_COLS = [
|
||||
"pct_fused",
|
||||
"pct_img",
|
||||
"pct_md",
|
||||
"phase",
|
||||
"se_mean",
|
||||
"se_std",
|
||||
"se_pct_lt_0.2",
|
||||
"se_pct_gt_0.8",
|
||||
"holdout_loss",
|
||||
"holdout_acc_fused",
|
||||
"holdout_acc_img",
|
||||
"holdout_acc_md",
|
||||
"holdout_auc_fused",
|
||||
"holdout_auc_img",
|
||||
"holdout_auc_md",
|
||||
"best_monitor",
|
||||
"best_so_far",
|
||||
"best_epoch",
|
||||
"early_best_so_far",
|
||||
"early_bad_epochs",
|
||||
"early_improved",
|
||||
"early_monitor",
|
||||
"holdout_best_monitor",
|
||||
"holdout_best_so_far",
|
||||
"holdout_best_epoch",
|
||||
]
|
||||
|
||||
|
||||
class HypertowerLogger:
|
||||
"""
|
||||
Shared logging utility for V2 tower workflows.
|
||||
- train.log line logging
|
||||
- epoch_log.csv row logging with stable header
|
||||
- lightweight JSON/array artifact helpers
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
run_dir: Path,
|
||||
train_log_path: Optional[Path] = None,
|
||||
epoch_log_path: Optional[Path] = None,
|
||||
logger_name: Optional[str] = None,
|
||||
) -> None:
|
||||
self.run_dir = Path(run_dir).resolve()
|
||||
self.run_dir.mkdir(parents=True, exist_ok=True)
|
||||
self.train_log_path = Path(train_log_path) if train_log_path else (self.run_dir / "train.log")
|
||||
self.epoch_log_path = Path(epoch_log_path) if epoch_log_path else (self.run_dir / "epoch_log.csv")
|
||||
|
||||
self._logger_name = logger_name or f"hypertower.{id(self)}"
|
||||
self.logger = logging.getLogger(self._logger_name)
|
||||
self.logger.setLevel(logging.INFO)
|
||||
self.logger.handlers = []
|
||||
fh = logging.FileHandler(str(self.train_log_path))
|
||||
fh.setFormatter(logging.Formatter("%(asctime)s - %(message)s"))
|
||||
self.logger.addHandler(fh)
|
||||
self.logger.propagate = False
|
||||
|
||||
self._epoch_log_fp = None
|
||||
self._epoch_log_writer = None
|
||||
self._epoch_log_fields: list[str] | None = None
|
||||
|
||||
def info(self, msg: str) -> None:
|
||||
self.logger.info(msg)
|
||||
|
||||
def warning(self, msg: str) -> None:
|
||||
self.logger.warning(msg)
|
||||
|
||||
def error(self, msg: str) -> None:
|
||||
self.logger.error(msg)
|
||||
|
||||
def write_epoch_row(
|
||||
self,
|
||||
row: dict,
|
||||
*,
|
||||
path: str | Path | None = None,
|
||||
optional_cols: Optional[list[str]] = None,
|
||||
) -> None:
|
||||
optional = optional_cols if optional_cols is not None else DEFAULT_OPTIONAL_EPOCH_COLS
|
||||
if self._epoch_log_writer is None:
|
||||
fieldnames = list(dict.fromkeys([*row.keys(), *optional]))
|
||||
target_path = Path(path) if path is not None else self.epoch_log_path
|
||||
target_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
self._epoch_log_fp = open(target_path, "w", newline="", encoding="utf-8")
|
||||
self._epoch_log_writer = csv.DictWriter(self._epoch_log_fp, fieldnames=fieldnames)
|
||||
self._epoch_log_writer.writeheader()
|
||||
self._epoch_log_fields = fieldnames
|
||||
|
||||
assert self._epoch_log_fields is not None
|
||||
assert self._epoch_log_writer is not None
|
||||
assert self._epoch_log_fp is not None
|
||||
for key in self._epoch_log_fields:
|
||||
row.setdefault(key, None)
|
||||
self._epoch_log_writer.writerow({k: row.get(k) for k in self._epoch_log_fields})
|
||||
self._epoch_log_fp.flush()
|
||||
|
||||
def write_json(self, path: str | Path, payload: dict) -> None:
|
||||
target = Path(path)
|
||||
if not target.is_absolute():
|
||||
target = self.run_dir / target
|
||||
target.parent.mkdir(parents=True, exist_ok=True)
|
||||
target.write_text(json.dumps(payload, indent=2), encoding="utf-8")
|
||||
|
||||
def close(self) -> None:
|
||||
if self._epoch_log_fp is not None:
|
||||
try:
|
||||
self._epoch_log_fp.close()
|
||||
except Exception:
|
||||
pass
|
||||
self._epoch_log_fp = None
|
||||
self._epoch_log_writer = None
|
||||
self._epoch_log_fields = None
|
||||
for handler in list(self.logger.handlers):
|
||||
try:
|
||||
handler.close()
|
||||
except Exception:
|
||||
pass
|
||||
self.logger.removeHandler(handler)
|
||||
@@ -1,244 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Callable, Optional
|
||||
|
||||
import torch
|
||||
from torch.utils.data import DataLoader, WeightedRandomSampler
|
||||
|
||||
from .network_manager import LoaderBundle, PatientSplit
|
||||
from .slot_dataset import SlotDataset, slot_collate
|
||||
from .profiles.base import SlotDescriptor, SimpleDatasetProfile
|
||||
|
||||
|
||||
def _default_slot_descriptors(patient_col: str, label_col: str) -> dict[str, SlotDescriptor]:
|
||||
return {
|
||||
"id_1": SlotDescriptor(
|
||||
key="id_1",
|
||||
kind="id",
|
||||
description=f"Patient identifier column ({patient_col})",
|
||||
required=True,
|
||||
shape_hint="scalar",
|
||||
),
|
||||
"eye_id_1": SlotDescriptor(
|
||||
key="eye_id_1",
|
||||
kind="id",
|
||||
description="Eye side identifier (OD/OS)",
|
||||
required=False,
|
||||
shape_hint="scalar",
|
||||
),
|
||||
"label_1": SlotDescriptor(
|
||||
key="label_1",
|
||||
kind="label",
|
||||
description=f"Label column ({label_col})",
|
||||
required=True,
|
||||
shape_hint="scalar",
|
||||
),
|
||||
"image_1": SlotDescriptor(
|
||||
key="image_1",
|
||||
kind="image",
|
||||
description="Primary image slot",
|
||||
required=False,
|
||||
shape_hint="HWC or CHW",
|
||||
),
|
||||
"matrix_1": SlotDescriptor(
|
||||
key="matrix_1",
|
||||
kind="matrix",
|
||||
description="Primary matrix slot",
|
||||
required=False,
|
||||
shape_hint="[feature_dim]",
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
def _row_to_sample(
|
||||
row: Any,
|
||||
*,
|
||||
clinical: Any,
|
||||
patient_col: str,
|
||||
label_col: str,
|
||||
) -> dict[str, Any]:
|
||||
return {
|
||||
"id_1": row[patient_col],
|
||||
"eye_id_1": str(row.get("eyeID", "")),
|
||||
"label_1": row[label_col],
|
||||
"image_1": clinical.get_image_path(row) if hasattr(clinical, "get_image_path") else None,
|
||||
"matrix_1": clinical.vectorize_row(row) if hasattr(clinical, "vectorize_row") else None,
|
||||
}
|
||||
|
||||
|
||||
@dataclass
|
||||
class SlotLoaderFactory:
|
||||
"""
|
||||
Generic loader factory that emits dict batches keyed by slot names.
|
||||
"""
|
||||
|
||||
image_transform: Optional[Callable] = None
|
||||
matrix_transform: Optional[Callable] = None
|
||||
num_workers: int = 0
|
||||
|
||||
def build(
|
||||
self,
|
||||
*,
|
||||
clinical: Any,
|
||||
split: PatientSplit,
|
||||
args: Any,
|
||||
fold: int,
|
||||
profile: Optional[Any] = None,
|
||||
) -> LoaderBundle:
|
||||
batch_size = int(getattr(args, "batch_size", 8))
|
||||
slot_desc = self._resolve_slot_descriptors(clinical=clinical, profile=profile)
|
||||
|
||||
train_samples = self._build_samples(split.train, clinical, profile, slot_desc)
|
||||
val_samples = self._build_samples(split.val, clinical, profile, slot_desc)
|
||||
holdout_samples = (
|
||||
self._build_samples(split.holdout, clinical, profile, slot_desc)
|
||||
if split.holdout is not None
|
||||
else None
|
||||
)
|
||||
|
||||
train_loader = DataLoader(
|
||||
SlotDataset(
|
||||
train_samples,
|
||||
slot_desc,
|
||||
image_transform=self.image_transform,
|
||||
matrix_transform=self.matrix_transform,
|
||||
),
|
||||
batch_size=batch_size,
|
||||
shuffle=True,
|
||||
num_workers=self.num_workers,
|
||||
collate_fn=slot_collate,
|
||||
)
|
||||
val_loader = DataLoader(
|
||||
SlotDataset(
|
||||
val_samples,
|
||||
slot_desc,
|
||||
image_transform=self.image_transform,
|
||||
matrix_transform=self.matrix_transform,
|
||||
),
|
||||
batch_size=batch_size,
|
||||
shuffle=False,
|
||||
num_workers=self.num_workers,
|
||||
collate_fn=slot_collate,
|
||||
)
|
||||
holdout_loader = None
|
||||
if holdout_samples is not None:
|
||||
holdout_loader = DataLoader(
|
||||
SlotDataset(
|
||||
holdout_samples,
|
||||
slot_desc,
|
||||
image_transform=self.image_transform,
|
||||
matrix_transform=self.matrix_transform,
|
||||
),
|
||||
batch_size=batch_size,
|
||||
shuffle=False,
|
||||
num_workers=self.num_workers,
|
||||
collate_fn=slot_collate,
|
||||
)
|
||||
return LoaderBundle(train=train_loader, val=val_loader, holdout=holdout_loader)
|
||||
|
||||
@staticmethod
|
||||
def _resolve_slot_descriptors(
|
||||
*,
|
||||
clinical: Any,
|
||||
profile: Optional[Any],
|
||||
) -> dict[str, SlotDescriptor]:
|
||||
if profile is not None and hasattr(profile, "slot_descriptors"):
|
||||
return profile.slot_descriptors()
|
||||
patient_col = getattr(clinical, "patient_col", "Patient ID")
|
||||
label_col = getattr(clinical, "label_col", "Diagnosis")
|
||||
return _default_slot_descriptors(patient_col, label_col)
|
||||
|
||||
@staticmethod
|
||||
def _build_samples(
|
||||
df,
|
||||
clinical: Any,
|
||||
profile: Optional[Any],
|
||||
slot_desc: dict[str, SlotDescriptor],
|
||||
) -> list[dict[str, Any]]:
|
||||
if df is None or df.empty:
|
||||
return []
|
||||
if profile is not None and hasattr(profile, "build_samples"):
|
||||
return profile.build_samples(df=df, clinical=clinical)
|
||||
|
||||
patient_col = getattr(profile, "patient_col", None) if profile is not None else None
|
||||
label_col = getattr(profile, "label_col", None) if profile is not None else None
|
||||
pcol = patient_col or "Patient ID"
|
||||
lcol = label_col or getattr(clinical, "label_col", "Diagnosis")
|
||||
samples = []
|
||||
for _, row in df.iterrows():
|
||||
sample = _row_to_sample(row, clinical=clinical, patient_col=pcol, label_col=lcol)
|
||||
for key in slot_desc.keys():
|
||||
sample.setdefault(key, None)
|
||||
samples.append(sample)
|
||||
return samples
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# V2 filter / loader helpers (used by V2HyperTower._run_fold)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def filter_eye_samples(samples: list[dict]) -> list[dict]:
|
||||
"""Keep any single-eye sample with a valid image, matrix, and label."""
|
||||
return [
|
||||
s for s in samples
|
||||
if s.get("image_1") is not None
|
||||
and s.get("matrix_1") is not None
|
||||
and s.get("label_1") is not None
|
||||
]
|
||||
|
||||
|
||||
def filter_bilateral_samples(samples: list[dict]) -> list[dict]:
|
||||
"""Keep only patient-level samples where both eyes are fully present."""
|
||||
return [
|
||||
s for s in samples
|
||||
if s.get("image_1") is not None
|
||||
and s.get("matrix_1") is not None
|
||||
and s.get("image_2") is not None
|
||||
and s.get("matrix_2") is not None
|
||||
and s.get("label_1") is not None
|
||||
]
|
||||
|
||||
|
||||
def make_loader(
|
||||
samples: list[dict],
|
||||
slots: dict,
|
||||
*,
|
||||
image_transform,
|
||||
image_preprocessor=None,
|
||||
image_cache=None,
|
||||
batch_size: int,
|
||||
shuffle: bool,
|
||||
num_workers: int,
|
||||
sampler: Optional[WeightedRandomSampler] = None,
|
||||
) -> DataLoader:
|
||||
ds = SlotDataset(
|
||||
samples,
|
||||
slots,
|
||||
image_transform=image_transform,
|
||||
image_preprocessor=image_preprocessor,
|
||||
image_cache=image_cache,
|
||||
)
|
||||
return DataLoader(
|
||||
ds,
|
||||
batch_size=batch_size,
|
||||
shuffle=(shuffle if sampler is None else False),
|
||||
sampler=sampler,
|
||||
num_workers=num_workers,
|
||||
collate_fn=slot_collate,
|
||||
)
|
||||
|
||||
|
||||
def build_balanced_sampler(samples: list[dict], label_key: str = "label_1") -> WeightedRandomSampler:
|
||||
"""Return a WeightedRandomSampler that equalises class frequency for training."""
|
||||
from collections import Counter
|
||||
labels = [s[label_key] for s in samples]
|
||||
counts = Counter(labels)
|
||||
weights = [1.0 / counts[lbl] for lbl in labels]
|
||||
return WeightedRandomSampler(weights, num_samples=len(weights), replacement=True)
|
||||
|
||||
|
||||
def to_label_tensor(labels, device: torch.device) -> torch.Tensor:
|
||||
if torch.is_tensor(labels):
|
||||
return labels.to(device=device, dtype=torch.long)
|
||||
return torch.as_tensor(labels, dtype=torch.long, device=device)
|
||||
@@ -1,243 +0,0 @@
|
||||
"""Metric computation, calibration, and threshold/bias tuning for V2."""
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Optional
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from sklearn.metrics import (
|
||||
cohen_kappa_score,
|
||||
f1_score,
|
||||
matthews_corrcoef,
|
||||
recall_score,
|
||||
roc_auc_score,
|
||||
roc_curve,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Loss
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def focal_loss(
|
||||
logits: torch.Tensor,
|
||||
targets: torch.Tensor,
|
||||
gamma: float = 0.0,
|
||||
weight: Optional[torch.Tensor] = None,
|
||||
reduction: str = "mean",
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Standard focal loss wrapper. When gamma=0 it reduces to cross entropy.
|
||||
weight should be per-class weights (same semantics as CrossEntropyLoss).
|
||||
"""
|
||||
if gamma <= 0:
|
||||
return F.cross_entropy(logits, targets, weight=weight, reduction=reduction)
|
||||
|
||||
log_probs = F.log_softmax(logits, dim=1)
|
||||
probs = log_probs.exp()
|
||||
|
||||
targets = targets.long().view(-1, 1)
|
||||
logpt = log_probs.gather(1, targets)
|
||||
pt = probs.gather(1, targets)
|
||||
|
||||
focal_factor = (1.0 - pt).clamp_min(0.0) ** gamma
|
||||
loss = -focal_factor * logpt
|
||||
|
||||
if weight is not None:
|
||||
class_weight = weight.gather(0, targets.view(-1))
|
||||
loss = loss * class_weight.view(-1, 1)
|
||||
|
||||
loss = loss.view(-1)
|
||||
if reduction == "sum":
|
||||
return loss.sum()
|
||||
if reduction == "mean":
|
||||
return loss.mean()
|
||||
return loss
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Basic array scoring
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _score_arrays(y_true: np.ndarray, probs: np.ndarray, num_classes: int):
|
||||
"""Returns (acc, auc, n)."""
|
||||
if y_true.size == 0:
|
||||
return float("nan"), float("nan"), 0
|
||||
acc = float((probs.argmax(1) == y_true).mean())
|
||||
try:
|
||||
auc = (
|
||||
float(roc_auc_score(y_true, probs[:, 1]))
|
||||
if num_classes == 2
|
||||
else float(roc_auc_score(y_true, probs, multi_class="ovr", average="macro"))
|
||||
)
|
||||
except Exception:
|
||||
auc = float("nan")
|
||||
return acc, auc, int(len(y_true))
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Calibration
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def compute_ece(y_true: np.ndarray, probs: np.ndarray, n_bins: int = 10) -> float:
|
||||
"""Expected Calibration Error: weighted mean of |confidence - accuracy| per bin."""
|
||||
if y_true.size == 0:
|
||||
return float("nan")
|
||||
confidences = probs.max(axis=1)
|
||||
predictions = probs.argmax(axis=1)
|
||||
bin_edges = np.linspace(0.0, 1.0, n_bins + 1)
|
||||
ece = 0.0
|
||||
n = len(y_true)
|
||||
for i, (lo, hi) in enumerate(zip(bin_edges[:-1], bin_edges[1:])):
|
||||
mask = (confidences >= lo) & (
|
||||
confidences <= hi if i == n_bins - 1 else confidences < hi
|
||||
)
|
||||
if not mask.any():
|
||||
continue
|
||||
bin_acc = float((predictions[mask] == y_true[mask]).mean())
|
||||
bin_conf = float(confidences[mask].mean())
|
||||
ece += float(mask.sum()) / n * abs(bin_conf - bin_acc)
|
||||
return float(ece)
|
||||
|
||||
|
||||
def compute_extended_metrics(
|
||||
y_true: np.ndarray,
|
||||
probs: np.ndarray,
|
||||
num_classes: int,
|
||||
n_bins: int = 10,
|
||||
preds_override: Optional[np.ndarray] = None,
|
||||
) -> dict:
|
||||
nan = float("nan")
|
||||
if y_true.size == 0:
|
||||
return dict(
|
||||
kappa=nan, mcc=nan, macro_f1=nan,
|
||||
per_class_recall=np.full(num_classes, nan), ece=nan,
|
||||
)
|
||||
preds = preds_override if preds_override is not None else probs.argmax(axis=1)
|
||||
try:
|
||||
kappa = float(cohen_kappa_score(y_true, preds))
|
||||
except Exception:
|
||||
kappa = nan
|
||||
try:
|
||||
mcc = float(matthews_corrcoef(y_true, preds))
|
||||
except Exception:
|
||||
mcc = nan
|
||||
try:
|
||||
macro_f1 = float(f1_score(y_true, preds, average="macro", zero_division=0))
|
||||
except Exception:
|
||||
macro_f1 = nan
|
||||
try:
|
||||
pcr = recall_score(
|
||||
y_true, preds, average=None,
|
||||
labels=list(range(num_classes)), zero_division=0,
|
||||
).astype(float)
|
||||
except Exception:
|
||||
pcr = np.full(num_classes, nan)
|
||||
ece = compute_ece(y_true, probs, n_bins=n_bins)
|
||||
return dict(kappa=kappa, mcc=mcc, macro_f1=macro_f1, per_class_recall=pcr, ece=ece)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Threshold / bias tuning
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def tune_binary_threshold(y_true: np.ndarray, p1: np.ndarray) -> float:
|
||||
"""Pick threshold via Youden's J (sensitivity + specificity − 1).
|
||||
|
||||
This is class-distribution independent, unlike maximising raw accuracy,
|
||||
which is biased toward the majority class on imbalanced validation sets.
|
||||
Falls back to 0.5 if both classes are not present.
|
||||
"""
|
||||
if y_true.size == 0 or len(np.unique(y_true)) < 2:
|
||||
return 0.5
|
||||
fpr, tpr, thresholds = roc_curve(y_true, p1)
|
||||
j = tpr + (1.0 - fpr) - 1.0
|
||||
return float(thresholds[np.argmax(j)])
|
||||
|
||||
|
||||
def multiclass_acc_with_bias(y_true: np.ndarray, probs: np.ndarray, bias: np.ndarray) -> float:
|
||||
"""Balanced accuracy (mean per-class recall) after applying log-space bias."""
|
||||
if y_true.size == 0:
|
||||
return float("nan")
|
||||
logits = np.log(np.clip(probs, 1e-8, 1.0)) + bias.reshape(1, -1)
|
||||
preds = np.argmax(logits, axis=1)
|
||||
classes = np.unique(y_true)
|
||||
per_class = [(preds[y_true == c] == c).mean() for c in classes]
|
||||
return float(np.mean(per_class))
|
||||
|
||||
|
||||
def tune_multiclass_bias(y_true: np.ndarray, probs: np.ndarray, *, iters: int = 2) -> np.ndarray:
|
||||
"""Grid-search per-class log-space bias to maximise balanced accuracy.
|
||||
|
||||
Balanced accuracy (mean per-class recall) is class-distribution independent,
|
||||
unlike raw accuracy which is biased toward the majority class on imbalanced
|
||||
validation sets.
|
||||
"""
|
||||
if y_true.size == 0 or probs.size == 0:
|
||||
return np.zeros((0,), dtype=float)
|
||||
c = probs.shape[1]
|
||||
bias = np.zeros((c,), dtype=float)
|
||||
grid = np.linspace(-1.0, 1.0, 41)
|
||||
for _ in range(iters):
|
||||
for k in range(c):
|
||||
best_v = bias[k]
|
||||
best_acc = multiclass_acc_with_bias(y_true, probs, bias)
|
||||
old = bias[k]
|
||||
for v in grid:
|
||||
bias[k] = float(v)
|
||||
acc = multiclass_acc_with_bias(y_true, probs, bias)
|
||||
if acc > best_acc or (acc == best_acc and abs(v) < abs(best_v)):
|
||||
best_acc, best_v = acc, float(v)
|
||||
bias[k] = best_v
|
||||
if np.isnan(best_acc):
|
||||
bias[k] = old
|
||||
return bias
|
||||
|
||||
|
||||
def _svf(vec) -> Optional[str]:
|
||||
"""Serialise a float vector to pipe-separated string, or None if empty."""
|
||||
if vec is None:
|
||||
return None
|
||||
arr = np.asarray(vec, dtype=float)
|
||||
if arr.size == 0:
|
||||
return None
|
||||
return "|".join(f"{float(v):.4f}" for v in arr.tolist())
|
||||
|
||||
|
||||
def _tune_and_snap(
|
||||
y: np.ndarray,
|
||||
p: np.ndarray,
|
||||
acc: float,
|
||||
num_classes: int,
|
||||
args,
|
||||
n_bins: int,
|
||||
) -> tuple[dict, float, Optional[np.ndarray], Optional[np.ndarray]]:
|
||||
"""
|
||||
Apply threshold/bias tuning and compute extended metrics.
|
||||
Returns (snap_dict, tuned_auc, threshold, bias).
|
||||
"""
|
||||
thr = 0.5 if num_classes == 2 else float("nan")
|
||||
bias = None
|
||||
ext_preds = None
|
||||
|
||||
if args.tune_binary_threshold and num_classes == 2 and y.size > 0:
|
||||
thr = tune_binary_threshold(y, p[:, 1])
|
||||
ext_preds = (p[:, 1] >= thr).astype(int)
|
||||
acc = float((ext_preds == y).mean())
|
||||
elif args.tune_multiclass_bias and num_classes > 2 and y.size > 0:
|
||||
bias = tune_multiclass_bias(y, p)
|
||||
logits = np.log(np.clip(p, 1e-8, 1.0)) + bias.reshape(1, -1)
|
||||
ext_preds = np.argmax(logits, axis=1)
|
||||
acc = float((ext_preds == y).mean())
|
||||
|
||||
ext = compute_extended_metrics(y, p, num_classes, n_bins=n_bins, preds_override=ext_preds)
|
||||
_, auc, n = _score_arrays(y, p, num_classes)
|
||||
|
||||
snap = dict(
|
||||
auc=auc, acc=acc, n=n,
|
||||
kappa=ext["kappa"], mcc=ext["mcc"], macro_f1=ext["macro_f1"],
|
||||
per_class_recall=ext["per_class_recall"], ece=ext["ece"],
|
||||
threshold=thr, bias=bias,
|
||||
)
|
||||
return snap, auc, thr, bias
|
||||
@@ -1,148 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Callable, Optional
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
from classes.v2.bridges import Bridge, VoteBridge
|
||||
from classes.v2.towers import ImageTower, MDTower
|
||||
|
||||
from .config_builder import ConfigAssembly
|
||||
from .transforms import build_transform_chain
|
||||
|
||||
|
||||
@dataclass
|
||||
class V2ModelBundle:
|
||||
image_tower: Optional[ImageTower]
|
||||
metadata_tower: Optional[MDTower]
|
||||
bridge: Optional[nn.Module]
|
||||
classifier: Optional[nn.Module]
|
||||
image_transform: Optional[Callable]
|
||||
matrix_transform: Optional[Callable]
|
||||
|
||||
|
||||
def build_model_bundle(
|
||||
assembly: ConfigAssembly,
|
||||
clinical: Any,
|
||||
*,
|
||||
device: Optional[torch.device] = None,
|
||||
strict: bool = True,
|
||||
) -> V2ModelBundle:
|
||||
"""
|
||||
Build torch modules and input transforms from a V2 config assembly.
|
||||
"""
|
||||
image_tower_spec = _pick_tower(assembly, "image")
|
||||
md_tower_spec = _pick_tower(assembly, "metadata")
|
||||
bridge_spec = _pick_bridge(assembly)
|
||||
image_loader = _pick_loader(assembly, input_type="image")
|
||||
|
||||
clinical_core = getattr(clinical, "clinical", clinical)
|
||||
num_classes = _infer_num_classes(clinical)
|
||||
|
||||
img_tower = None
|
||||
if image_tower_spec is not None:
|
||||
img_tower = ImageTower(
|
||||
backbone=image_tower_spec.params.get("backbone", "efficientnet_b0"),
|
||||
freeze_ratio=float(image_tower_spec.params.get("freeze_ratio", 0.0) or 0.0),
|
||||
use_se=bool(image_tower_spec.params.get("use_se", False)),
|
||||
se_reduction=int(image_tower_spec.params.get("se_reduction", 16) or 16),
|
||||
se_pre_norm=bool(image_tower_spec.params.get("se_pre_norm", True)),
|
||||
augment=bool(image_tower_spec.params.get("augment", True)),
|
||||
geometry_dim=int(image_tower_spec.params.get("geometry_dim", 0) or 0),
|
||||
)
|
||||
if device is not None:
|
||||
img_tower = img_tower.to(device)
|
||||
|
||||
md_tower = None
|
||||
if md_tower_spec is not None:
|
||||
md_tower = MDTower(
|
||||
clinical_core,
|
||||
hidden_dim=int(md_tower_spec.params.get("hidden_dim", 128) or 128),
|
||||
dropout=float(md_tower_spec.params.get("dropout", 0.1) or 0.1),
|
||||
use_se=bool(md_tower_spec.params.get("use_se", False)),
|
||||
se_reduction=int(md_tower_spec.params.get("se_reduction", 16) or 16),
|
||||
se_pre_norm=bool(md_tower_spec.params.get("se_pre_norm", True)),
|
||||
)
|
||||
if device is not None:
|
||||
md_tower = md_tower.to(device)
|
||||
|
||||
bridge = None
|
||||
if bridge_spec is not None and img_tower is not None and md_tower is not None:
|
||||
if bridge_spec.method == "consensus":
|
||||
bridge = VoteBridge(num_classes=num_classes)
|
||||
else:
|
||||
bridge = Bridge(
|
||||
img_dim=img_tower.out_dim,
|
||||
meta_dim=md_tower.out_dim,
|
||||
num_classes=num_classes,
|
||||
fusion_dim=int(bridge_spec.params.get("fusion_dim", 256) or 256),
|
||||
mode="fused",
|
||||
use_se=bool(bridge_spec.params.get("use_se", True)),
|
||||
se_reduction=int(bridge_spec.params.get("se_reduction", 16) or 16),
|
||||
se_pre_norm=bool(bridge_spec.params.get("se_pre_norm", True)),
|
||||
)
|
||||
if device is not None:
|
||||
bridge = bridge.to(device)
|
||||
|
||||
classifier = None
|
||||
if assembly.classifiers:
|
||||
classifier = nn.Identity()
|
||||
if device is not None:
|
||||
classifier = classifier.to(device)
|
||||
|
||||
image_transform = None
|
||||
if image_loader is not None and image_tower_spec is not None:
|
||||
image_transform = build_transform_chain(
|
||||
image_loader.transforms,
|
||||
backbone_name=image_tower_spec.params.get("backbone", "efficientnet_b0"),
|
||||
augment=bool(image_tower_spec.params.get("augment", True)),
|
||||
strict=strict,
|
||||
)
|
||||
|
||||
return V2ModelBundle(
|
||||
image_tower=img_tower,
|
||||
metadata_tower=md_tower,
|
||||
bridge=bridge,
|
||||
classifier=classifier,
|
||||
image_transform=image_transform,
|
||||
matrix_transform=None,
|
||||
)
|
||||
|
||||
|
||||
def _pick_tower(assembly: ConfigAssembly, tower_type: str):
|
||||
matches = [tower for tower in assembly.towers.values() if tower.tower_type == tower_type]
|
||||
if not matches:
|
||||
return None
|
||||
if len(matches) > 1:
|
||||
raise ValueError(f"Multiple {tower_type} towers found; only one is supported for now.")
|
||||
return matches[0]
|
||||
|
||||
|
||||
def _pick_bridge(assembly: ConfigAssembly):
|
||||
if not assembly.bridges:
|
||||
return None
|
||||
if len(assembly.bridges) > 1:
|
||||
raise ValueError("Multiple bridges found; only one is supported for now.")
|
||||
return next(iter(assembly.bridges.values()))
|
||||
|
||||
|
||||
def _pick_loader(assembly: ConfigAssembly, input_type: str):
|
||||
matches = [loader for loader in assembly.loaders.values() if loader.input_type == input_type]
|
||||
if not matches:
|
||||
return None
|
||||
if len(matches) > 1:
|
||||
raise ValueError(f"Multiple loaders with input_type={input_type!r} found.")
|
||||
return matches[0]
|
||||
|
||||
|
||||
def _infer_num_classes(clinical: Any) -> int:
|
||||
df = getattr(clinical, "df", None)
|
||||
label_col = getattr(clinical, "label_col", None)
|
||||
if df is None and hasattr(clinical, "clinical"):
|
||||
df = clinical.clinical.df
|
||||
label_col = clinical.clinical.label_col
|
||||
if df is None or label_col is None or label_col not in df.columns:
|
||||
return 2
|
||||
return int(df[label_col].dropna().nunique())
|
||||
@@ -1,858 +0,0 @@
|
||||
"""V2 model classes and training/inference helpers."""
|
||||
from __future__ import annotations
|
||||
|
||||
from random import random
|
||||
from typing import Optional
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch import nn
|
||||
from torch.utils.data import DataLoader
|
||||
|
||||
from classes.v2.bridges import Bridge
|
||||
from classes.v2.towers import ImageTower, MDTower
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Model classes
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class SingleEyeHT(nn.Module):
|
||||
"""
|
||||
ImageTower + MDTower + Bridge, trained on eye-level samples.
|
||||
Supports both Classic (eye-level) and Ensemble (patient-level averaging) eval.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
backbone: str,
|
||||
freeze_ratio: float,
|
||||
augment: bool,
|
||||
clinical_data,
|
||||
num_classes: int,
|
||||
md_hidden_dim: int = 128,
|
||||
fusion_dim: int = 256,
|
||||
bridge_mode: str = "fused",
|
||||
):
|
||||
super().__init__()
|
||||
self.img_tower = ImageTower(
|
||||
backbone=backbone,
|
||||
freeze_ratio=freeze_ratio,
|
||||
augment=augment,
|
||||
use_se=False,
|
||||
)
|
||||
self.md_tower = MDTower(
|
||||
clinical_data=clinical_data,
|
||||
hidden_dim=md_hidden_dim,
|
||||
use_se=False,
|
||||
)
|
||||
self.bridge = Bridge(
|
||||
img_dim=self.img_tower.out_dim,
|
||||
meta_dim=self.md_tower.out_dim,
|
||||
num_classes=num_classes,
|
||||
fusion_dim=fusion_dim,
|
||||
mode=bridge_mode,
|
||||
use_se=False,
|
||||
)
|
||||
|
||||
@property
|
||||
def transform(self):
|
||||
return self.img_tower.transform
|
||||
|
||||
def forward(self, x: torch.Tensor, meta: torch.Tensor) -> torch.Tensor:
|
||||
img_feats = None if self.bridge.mode == "metadata_only" else self.img_tower(x)
|
||||
md_feats = None if self.bridge.mode == "image_only" else self.md_tower(meta)
|
||||
out_f, _, _ = self.bridge(img_feats, md_feats)
|
||||
return out_f
|
||||
|
||||
|
||||
class BilateralHT(nn.Module):
|
||||
"""
|
||||
Bilateral mode with joint towers:
|
||||
- shared eye-level towers encode OD/OS independently
|
||||
- joint image and metadata towers combine OD/OS embeddings
|
||||
- standard Bridge fuses joint image + joint metadata embeddings
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
backbone: str,
|
||||
freeze_ratio: float,
|
||||
augment: bool,
|
||||
clinical_data,
|
||||
num_classes: int,
|
||||
md_hidden_dim: int = 128,
|
||||
fusion_dim: int = 256,
|
||||
):
|
||||
super().__init__()
|
||||
self.eye_img_tower = ImageTower(
|
||||
backbone=backbone,
|
||||
freeze_ratio=freeze_ratio,
|
||||
augment=augment,
|
||||
use_se=False,
|
||||
)
|
||||
self.eye_md_tower = MDTower(
|
||||
clinical_data=clinical_data,
|
||||
hidden_dim=md_hidden_dim,
|
||||
use_se=False,
|
||||
)
|
||||
img_dim = self.eye_img_tower.out_dim
|
||||
md_dim = self.eye_md_tower.out_dim
|
||||
self.joint_img = nn.Sequential(
|
||||
nn.Linear(2 * img_dim, fusion_dim),
|
||||
nn.LayerNorm(fusion_dim),
|
||||
nn.ReLU(),
|
||||
nn.Dropout(0.3),
|
||||
nn.Linear(fusion_dim, img_dim),
|
||||
)
|
||||
self.joint_md = nn.Sequential(
|
||||
nn.Linear(2 * md_dim, fusion_dim),
|
||||
nn.LayerNorm(fusion_dim),
|
||||
nn.ReLU(),
|
||||
nn.Dropout(0.3),
|
||||
nn.Linear(fusion_dim, md_dim),
|
||||
)
|
||||
self.bridge = Bridge(
|
||||
img_dim=img_dim,
|
||||
meta_dim=md_dim,
|
||||
num_classes=num_classes,
|
||||
fusion_dim=fusion_dim,
|
||||
mode="fused",
|
||||
use_se=False,
|
||||
)
|
||||
# Auxiliary heads for tower warmup / BCD tower steps.
|
||||
self.aux_img = nn.Linear(img_dim, num_classes)
|
||||
self.aux_md = nn.Linear(md_dim, num_classes)
|
||||
|
||||
@property
|
||||
def transform(self):
|
||||
return self.eye_img_tower.transform
|
||||
|
||||
def encode_joint(
|
||||
self,
|
||||
x_od: torch.Tensor,
|
||||
meta_od: torch.Tensor,
|
||||
x_os: torch.Tensor,
|
||||
meta_os: torch.Tensor,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
img_od = self.eye_img_tower(x_od)
|
||||
md_od = self.eye_md_tower(meta_od)
|
||||
img_os = self.eye_img_tower(x_os)
|
||||
md_os = self.eye_md_tower(meta_os)
|
||||
joint_img = self.joint_img(torch.cat([img_od, img_os], dim=1))
|
||||
joint_md = self.joint_md(torch.cat([md_od, md_os], dim=1))
|
||||
return joint_img, joint_md
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x_od: torch.Tensor,
|
||||
meta_od: torch.Tensor,
|
||||
x_os: torch.Tensor,
|
||||
meta_os: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
joint_img, joint_md = self.encode_joint(x_od, meta_od, x_os, meta_os)
|
||||
out_f, _, _ = self.bridge(joint_img, joint_md)
|
||||
return out_f
|
||||
|
||||
|
||||
class FusedEnsembleHT(nn.Module):
|
||||
"""
|
||||
SingleEyeHT base with a per-eye attention scorer for bilateral fusion.
|
||||
|
||||
The base model is trained eye-level (identical to ensemble mode).
|
||||
After base training completes, the base is frozen and only the
|
||||
eye_scorer is trained on bilateral (patient-level) samples.
|
||||
|
||||
At inference, eye_scorer is applied independently to each eye's logit
|
||||
vector to produce a scalar attention score. Softmax over the two scores
|
||||
gives attention weights; the final logit is a weighted sum:
|
||||
|
||||
score_od = eye_scorer(logit_od) # [B, 1]
|
||||
score_os = eye_scorer(logit_os) # [B, 1]
|
||||
alpha = softmax([score_od, score_os]) # [B, 2], sums to 1
|
||||
out = alpha[:,0:1]*logit_od + alpha[:,1:2]*logit_os
|
||||
|
||||
Because eye_scorer is applied to each eye with the same weights, the
|
||||
mechanism is permutation-equivariant — there is no left/right positional
|
||||
bias. Through training on bilateral labels the scorer learns to give high
|
||||
scores to logits that point strongly toward the GC class, creating the
|
||||
desired asymmetry: a confidently GC eye dominates the patient prediction
|
||||
more than a comparably confident healthy eye would.
|
||||
"""
|
||||
|
||||
def __init__(self, base: SingleEyeHT, num_classes: int):
|
||||
super().__init__()
|
||||
self.base = base
|
||||
# Applied independently to each eye's logit → scalar attention score.
|
||||
# Learns the GC-direction in logit space from bilateral labels.
|
||||
self.eye_scorer = nn.Linear(num_classes, 1, bias=True)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x_od: torch.Tensor,
|
||||
meta_od: torch.Tensor,
|
||||
x_os: torch.Tensor,
|
||||
meta_os: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
logit_od = self.base(x_od, meta_od) # [B, C]
|
||||
logit_os = self.base(x_os, meta_os) # [B, C]
|
||||
scores = torch.cat([self.eye_scorer(logit_od),
|
||||
self.eye_scorer(logit_os)], dim=1) # [B, 2]
|
||||
alpha = torch.softmax(scores, dim=1) # [B, 2]
|
||||
return alpha[:, 0:1] * logit_od + alpha[:, 1:2] * logit_os # [B, C]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Phase control
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _set_requires_grad(module: nn.Module, enabled: bool) -> None:
|
||||
for p in module.parameters():
|
||||
p.requires_grad = enabled
|
||||
|
||||
|
||||
def _set_single_phase(model: SingleEyeHT, phase: str) -> None:
|
||||
bridge_mode = model.bridge.mode
|
||||
# Ablation modes have no fusion bridge; fused_warmup is meaningless — treat as tower_warmup
|
||||
if bridge_mode in ("image_only", "metadata_only") and phase == "fused_warmup":
|
||||
phase = "tower_warmup"
|
||||
if phase == "md_warmup":
|
||||
_set_requires_grad(model.img_tower, False)
|
||||
_set_requires_grad(model.md_tower, True)
|
||||
_set_requires_grad(model.bridge.classifier_img, False)
|
||||
_set_requires_grad(model.bridge.classifier_md, True)
|
||||
_set_requires_grad(model.bridge.W_img, False)
|
||||
_set_requires_grad(model.bridge.W_md, False)
|
||||
_set_requires_grad(model.bridge.classifier_fused, False)
|
||||
return
|
||||
if phase == "tower_warmup":
|
||||
_set_requires_grad(model.img_tower, bridge_mode != "metadata_only")
|
||||
_set_requires_grad(model.md_tower, bridge_mode != "image_only")
|
||||
_set_requires_grad(model.bridge.classifier_img, bridge_mode != "metadata_only")
|
||||
_set_requires_grad(model.bridge.classifier_md, bridge_mode != "image_only")
|
||||
_set_requires_grad(model.bridge.W_img, False)
|
||||
_set_requires_grad(model.bridge.W_md, False)
|
||||
_set_requires_grad(model.bridge.classifier_fused, False)
|
||||
return
|
||||
if phase == "fused_warmup":
|
||||
_set_requires_grad(model.img_tower, False)
|
||||
_set_requires_grad(model.md_tower, False)
|
||||
_set_requires_grad(model.bridge.classifier_img, False)
|
||||
_set_requires_grad(model.bridge.classifier_md, False)
|
||||
_set_requires_grad(model.bridge.W_img, True)
|
||||
_set_requires_grad(model.bridge.W_md, True)
|
||||
_set_requires_grad(model.bridge.classifier_fused, True)
|
||||
return
|
||||
_set_requires_grad(model, True)
|
||||
|
||||
|
||||
def _set_bilateral_phase(model: BilateralHT, phase: str) -> None:
|
||||
if phase == "tower_warmup":
|
||||
_set_requires_grad(model.eye_img_tower, True)
|
||||
_set_requires_grad(model.eye_md_tower, True)
|
||||
_set_requires_grad(model.joint_img, True)
|
||||
_set_requires_grad(model.joint_md, True)
|
||||
_set_requires_grad(model.aux_img, True)
|
||||
_set_requires_grad(model.aux_md, True)
|
||||
_set_requires_grad(model.bridge, False)
|
||||
return
|
||||
if phase == "fused_warmup":
|
||||
_set_requires_grad(model.eye_img_tower, False)
|
||||
_set_requires_grad(model.eye_md_tower, False)
|
||||
_set_requires_grad(model.joint_img, False)
|
||||
_set_requires_grad(model.joint_md, False)
|
||||
_set_requires_grad(model.aux_img, False)
|
||||
_set_requires_grad(model.aux_md, False)
|
||||
_set_requires_grad(model.bridge, True)
|
||||
return
|
||||
_set_requires_grad(model, True)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Training helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def train_single_epoch(
|
||||
model: SingleEyeHT,
|
||||
loader: DataLoader,
|
||||
opt,
|
||||
device: torch.device,
|
||||
*,
|
||||
phase: str,
|
||||
bcd_prob: float = 0.5,
|
||||
tower_loss_mode: str = "bcd",
|
||||
) -> tuple[float, float]:
|
||||
model.train()
|
||||
_set_single_phase(model, phase)
|
||||
total_loss = total_correct = total_n = 0
|
||||
for batch in loader:
|
||||
x = batch.get("image_1")
|
||||
m = batch.get("matrix_1")
|
||||
y = batch.get("label_1")
|
||||
if phase == "md_warmup":
|
||||
if not torch.is_tensor(m):
|
||||
continue
|
||||
m = m.to(device)
|
||||
y = _to_label_tensor(y, device)
|
||||
md_feats = model.md_tower(m)
|
||||
logits = model.bridge.classifier_md(md_feats)
|
||||
loss = F.cross_entropy(logits, y)
|
||||
opt.zero_grad(); loss.backward(); opt.step()
|
||||
bs = y.shape[0]
|
||||
total_loss += float(loss.item()) * bs
|
||||
total_correct += int((logits.argmax(1) == y).sum())
|
||||
total_n += bs
|
||||
continue
|
||||
if not torch.is_tensor(x) or not torch.is_tensor(m):
|
||||
continue
|
||||
x = x.to(device)
|
||||
m = m.to(device)
|
||||
y = _to_label_tensor(y, device)
|
||||
bridge_mode = model.bridge.mode
|
||||
|
||||
img_feats = None if bridge_mode == "metadata_only" else model.img_tower(x)
|
||||
md_feats = None if bridge_mode == "image_only" else model.md_tower(m)
|
||||
|
||||
if phase == "tower_warmup":
|
||||
if bridge_mode == "metadata_only":
|
||||
logits = model.bridge.classifier_md(md_feats)
|
||||
loss = F.cross_entropy(logits, y)
|
||||
elif bridge_mode == "image_only":
|
||||
logits = model.bridge.classifier_img(img_feats)
|
||||
loss = F.cross_entropy(logits, y)
|
||||
else:
|
||||
logits_i = model.bridge.classifier_img(img_feats)
|
||||
logits_m = model.bridge.classifier_md(md_feats)
|
||||
loss = 0.5 * (F.cross_entropy(logits_i, y) + F.cross_entropy(logits_m, y))
|
||||
logits = 0.5 * (F.softmax(logits_i, dim=1) + F.softmax(logits_m, dim=1))
|
||||
elif phase == "fused_warmup":
|
||||
logits, _, _ = model.bridge(img_feats, md_feats)
|
||||
loss = F.cross_entropy(logits, y)
|
||||
else:
|
||||
if bridge_mode == "metadata_only":
|
||||
logits = model.bridge.classifier_md(md_feats)
|
||||
loss = F.cross_entropy(logits, y)
|
||||
elif bridge_mode == "image_only":
|
||||
logits = model.bridge.classifier_img(img_feats)
|
||||
loss = F.cross_entropy(logits, y)
|
||||
elif tower_loss_mode == "all":
|
||||
loss_i = F.cross_entropy(model.bridge.classifier_img(img_feats), y)
|
||||
loss_m = F.cross_entropy(model.bridge.classifier_md(md_feats), y)
|
||||
logits, _, _ = model.bridge(img_feats, md_feats)
|
||||
loss = F.cross_entropy(logits, y) + loss_i + loss_m
|
||||
elif random() < bcd_prob:
|
||||
if random() < 0.5:
|
||||
logits = model.bridge.classifier_img(img_feats)
|
||||
else:
|
||||
logits = model.bridge.classifier_md(md_feats)
|
||||
loss = F.cross_entropy(logits, y)
|
||||
else:
|
||||
logits, _, _ = model.bridge(img_feats, md_feats)
|
||||
loss = F.cross_entropy(logits, y)
|
||||
|
||||
opt.zero_grad()
|
||||
loss.backward()
|
||||
opt.step()
|
||||
bs = y.shape[0]
|
||||
total_loss += float(loss.item()) * bs
|
||||
total_correct += int((logits.argmax(1) == y).sum())
|
||||
total_n += bs
|
||||
return (
|
||||
total_loss / total_n if total_n else float("nan"),
|
||||
total_correct / total_n if total_n else float("nan"),
|
||||
)
|
||||
|
||||
|
||||
def train_bilateral_epoch(
|
||||
model: BilateralHT,
|
||||
loader: DataLoader,
|
||||
opt,
|
||||
device: torch.device,
|
||||
*,
|
||||
phase: str,
|
||||
bcd_prob: float = 0.5,
|
||||
tower_loss_mode: str = "bcd",
|
||||
) -> tuple[float, float]:
|
||||
model.train()
|
||||
_set_bilateral_phase(model, phase)
|
||||
total_loss = total_correct = total_n = 0
|
||||
for batch in loader:
|
||||
x1 = batch.get("image_1")
|
||||
m1 = batch.get("matrix_1")
|
||||
x2 = batch.get("image_2")
|
||||
m2 = batch.get("matrix_2")
|
||||
y = batch.get("label_1")
|
||||
if not (torch.is_tensor(x1) and torch.is_tensor(m1) and torch.is_tensor(x2) and torch.is_tensor(m2)):
|
||||
continue
|
||||
x1 = x1.to(device); m1 = m1.to(device)
|
||||
x2 = x2.to(device); m2 = m2.to(device)
|
||||
y = _to_label_tensor(y, device)
|
||||
joint_img, joint_md = model.encode_joint(x1, m1, x2, m2)
|
||||
|
||||
if phase == "tower_warmup":
|
||||
logits_i = model.aux_img(joint_img)
|
||||
logits_m = model.aux_md(joint_md)
|
||||
loss = 0.5 * (F.cross_entropy(logits_i, y) + F.cross_entropy(logits_m, y))
|
||||
logits = 0.5 * (F.softmax(logits_i, dim=1) + F.softmax(logits_m, dim=1))
|
||||
elif phase == "fused_warmup":
|
||||
logits, _, _ = model.bridge(joint_img, joint_md)
|
||||
loss = F.cross_entropy(logits, y)
|
||||
else:
|
||||
if tower_loss_mode == "all":
|
||||
loss_i = F.cross_entropy(model.aux_img(joint_img), y)
|
||||
loss_m = F.cross_entropy(model.aux_md(joint_md), y)
|
||||
logits, _, _ = model.bridge(joint_img, joint_md)
|
||||
loss = F.cross_entropy(logits, y) + loss_i + loss_m
|
||||
elif random() < bcd_prob:
|
||||
if random() < 0.5:
|
||||
logits = model.aux_img(joint_img)
|
||||
else:
|
||||
logits = model.aux_md(joint_md)
|
||||
loss = F.cross_entropy(logits, y)
|
||||
else:
|
||||
logits, _, _ = model.bridge(joint_img, joint_md)
|
||||
loss = F.cross_entropy(logits, y)
|
||||
|
||||
opt.zero_grad()
|
||||
loss.backward()
|
||||
opt.step()
|
||||
bs = y.shape[0]
|
||||
total_loss += float(loss.item()) * bs
|
||||
total_correct += int((logits.argmax(1) == y).sum())
|
||||
total_n += bs
|
||||
return (
|
||||
total_loss / total_n if total_n else float("nan"),
|
||||
total_correct / total_n if total_n else float("nan"),
|
||||
)
|
||||
|
||||
|
||||
def train_fusion_epoch(
|
||||
model: FusedEnsembleHT,
|
||||
loader: DataLoader,
|
||||
opt,
|
||||
device: torch.device,
|
||||
) -> tuple[float, float]:
|
||||
"""Train only the fusion head; the base SingleEyeHT is frozen in eval mode."""
|
||||
model.base.eval()
|
||||
model.eye_scorer.train()
|
||||
total_loss = total_correct = total_n = 0
|
||||
for batch in loader:
|
||||
x1 = batch.get("image_1"); m1 = batch.get("matrix_1")
|
||||
x2 = batch.get("image_2"); m2 = batch.get("matrix_2")
|
||||
y = batch.get("label_1")
|
||||
if not (torch.is_tensor(x1) and torch.is_tensor(m1) and
|
||||
torch.is_tensor(x2) and torch.is_tensor(m2)):
|
||||
continue
|
||||
y_t = _to_label_tensor(y, device)
|
||||
out = model(x1.to(device), m1.to(device), x2.to(device), m2.to(device))
|
||||
loss = F.cross_entropy(out, y_t)
|
||||
opt.zero_grad()
|
||||
loss.backward()
|
||||
opt.step()
|
||||
bs = y_t.shape[0]
|
||||
total_loss += float(loss.item()) * bs
|
||||
total_correct += int((out.argmax(1) == y_t).sum())
|
||||
total_n += bs
|
||||
return (
|
||||
total_loss / total_n if total_n else float("nan"),
|
||||
total_correct / total_n if total_n else float("nan"),
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Inference helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _to_label_tensor(labels, device: torch.device) -> torch.Tensor:
|
||||
if torch.is_tensor(labels):
|
||||
return labels.to(device=device, dtype=torch.long)
|
||||
return torch.as_tensor(labels, dtype=torch.long, device=device)
|
||||
|
||||
|
||||
def collect_probs_classic(
|
||||
model: SingleEyeHT,
|
||||
loader: DataLoader,
|
||||
device: torch.device,
|
||||
) -> tuple[np.ndarray, np.ndarray]:
|
||||
"""
|
||||
Classic eye-level eval using the bilateral val loader.
|
||||
OD and OS are treated as independent samples (both contribute to the
|
||||
arrays with the same patient label). Returns (y_true [2N], probs [2N, C]).
|
||||
"""
|
||||
model.eval()
|
||||
y_chunks, p_chunks = [], []
|
||||
with torch.no_grad():
|
||||
for batch in loader:
|
||||
x1 = batch.get("image_1"); m1 = batch.get("matrix_1")
|
||||
x2 = batch.get("image_2"); m2 = batch.get("matrix_2")
|
||||
y = batch.get("label_1")
|
||||
if not (torch.is_tensor(x1) and torch.is_tensor(m1) and torch.is_tensor(x2) and torch.is_tensor(m2)):
|
||||
continue
|
||||
y_t = _to_label_tensor(y, device)
|
||||
p_od = F.softmax(model(x1.to(device), m1.to(device)), dim=1)
|
||||
p_os = F.softmax(model(x2.to(device), m2.to(device)), dim=1)
|
||||
y_np = y_t.cpu().numpy()
|
||||
y_chunks += [y_np, y_np]
|
||||
p_chunks += [p_od.cpu().numpy(), p_os.cpu().numpy()]
|
||||
if not y_chunks:
|
||||
return np.array([], dtype=np.int64), np.zeros((0, 0), dtype=np.float32)
|
||||
return np.concatenate(y_chunks), np.concatenate(p_chunks, axis=0)
|
||||
|
||||
|
||||
def collect_probs_ensemble_pereye(
|
||||
model: "SingleEyeHT",
|
||||
loader: DataLoader,
|
||||
device: torch.device,
|
||||
*,
|
||||
return_ids: bool = False,
|
||||
):
|
||||
"""
|
||||
Per-patient, per-eye probs for all 3 heads from a bilateral loader (ensemble mode).
|
||||
|
||||
OD corresponds to image_1/matrix_1; OS to image_2/matrix_2.
|
||||
Arrays are in patient order (not interleaved at sample level).
|
||||
|
||||
Returns:
|
||||
(y, pf_od, pi_od, pm_od, pf_os, pi_os, pm_os)
|
||||
or, when return_ids=True:
|
||||
(y, pf_od, pi_od, pm_od, pf_os, pi_os, pm_os, patient_ids)
|
||||
|
||||
Patient-level averaged ensemble probs can be recovered as:
|
||||
p_en = 0.5 * (pf_od + pf_os)
|
||||
"""
|
||||
model.eval()
|
||||
y_chunks: list = []
|
||||
pf_od_c, pi_od_c, pm_od_c = [], [], []
|
||||
pf_os_c, pi_os_c, pm_os_c = [], [], []
|
||||
id_chunks: list[str] = []
|
||||
|
||||
with torch.no_grad():
|
||||
for batch in loader:
|
||||
x1 = batch.get("image_1"); m1 = batch.get("matrix_1")
|
||||
x2 = batch.get("image_2"); m2 = batch.get("matrix_2")
|
||||
y = batch.get("label_1")
|
||||
if not (torch.is_tensor(x1) and torch.is_tensor(m1) and
|
||||
torch.is_tensor(x2) and torch.is_tensor(m2)):
|
||||
continue
|
||||
y_t = _to_label_tensor(y, device)
|
||||
|
||||
def _fwd(x, m):
|
||||
img_feats = None if model.bridge.mode == "metadata_only" else model.img_tower(x.to(device))
|
||||
md_feats = None if model.bridge.mode == "image_only" else model.md_tower(m.to(device))
|
||||
out_f, out_i, out_m = model.bridge(img_feats, md_feats)
|
||||
pf = F.softmax(out_f, dim=1)
|
||||
pi = F.softmax(out_i, dim=1) if out_i is not None else pf
|
||||
pm = F.softmax(out_m, dim=1) if out_m is not None else pf
|
||||
return pf, pi, pm
|
||||
|
||||
pf_od, pi_od, pm_od = _fwd(x1, m1)
|
||||
pf_os, pi_os, pm_os = _fwd(x2, m2)
|
||||
|
||||
y_chunks.append(y_t.cpu().numpy())
|
||||
pf_od_c.append(pf_od.cpu().numpy()); pi_od_c.append(pi_od.cpu().numpy()); pm_od_c.append(pm_od.cpu().numpy())
|
||||
pf_os_c.append(pf_os.cpu().numpy()); pi_os_c.append(pi_os.cpu().numpy()); pm_os_c.append(pm_os.cpu().numpy())
|
||||
|
||||
if return_ids:
|
||||
ids = batch.get("id_1", [""] * len(y_t))
|
||||
if torch.is_tensor(ids):
|
||||
ids = ids.tolist()
|
||||
id_chunks.extend([str(i) for i in ids])
|
||||
|
||||
if not y_chunks:
|
||||
z = np.zeros((0, 0), dtype=np.float32)
|
||||
empty_i = np.array([], dtype=np.int64)
|
||||
base = (empty_i, z, z, z, z, z, z)
|
||||
return base + (np.array([], dtype=object),) if return_ids else base
|
||||
|
||||
y = np.concatenate(y_chunks)
|
||||
pf_od = np.concatenate(pf_od_c, axis=0); pi_od = np.concatenate(pi_od_c, axis=0); pm_od = np.concatenate(pm_od_c, axis=0)
|
||||
pf_os = np.concatenate(pf_os_c, axis=0); pi_os = np.concatenate(pi_os_c, axis=0); pm_os = np.concatenate(pm_os_c, axis=0)
|
||||
if return_ids:
|
||||
return y, pf_od, pi_od, pm_od, pf_os, pi_os, pm_os, np.array(id_chunks, dtype=object)
|
||||
return y, pf_od, pi_od, pm_od, pf_os, pi_os, pm_os
|
||||
|
||||
|
||||
def collect_probs_ensemble(
|
||||
model: SingleEyeHT,
|
||||
loader: DataLoader,
|
||||
device: torch.device,
|
||||
) -> tuple[np.ndarray, np.ndarray]:
|
||||
"""
|
||||
Patient-level ensemble eval: average OD and OS softmax probabilities.
|
||||
Returns (y_true [N], probs [N, C]).
|
||||
"""
|
||||
model.eval()
|
||||
y_chunks, p_chunks = [], []
|
||||
with torch.no_grad():
|
||||
for batch in loader:
|
||||
x1 = batch.get("image_1"); m1 = batch.get("matrix_1")
|
||||
x2 = batch.get("image_2"); m2 = batch.get("matrix_2")
|
||||
y = batch.get("label_1")
|
||||
if not (torch.is_tensor(x1) and torch.is_tensor(m1) and torch.is_tensor(x2) and torch.is_tensor(m2)):
|
||||
continue
|
||||
y_t = _to_label_tensor(y, device)
|
||||
p_od = F.softmax(model(x1.to(device), m1.to(device)), dim=1)
|
||||
p_os = F.softmax(model(x2.to(device), m2.to(device)), dim=1)
|
||||
p = 0.5 * (p_od + p_os)
|
||||
y_chunks.append(y_t.cpu().numpy())
|
||||
p_chunks.append(p.cpu().numpy())
|
||||
if not y_chunks:
|
||||
return np.array([], dtype=np.int64), np.zeros((0, 0), dtype=np.float32)
|
||||
return np.concatenate(y_chunks), np.concatenate(p_chunks, axis=0)
|
||||
|
||||
|
||||
def collect_probs_bilateral(
|
||||
model: BilateralHT,
|
||||
loader: DataLoader,
|
||||
device: torch.device,
|
||||
) -> tuple[np.ndarray, np.ndarray]:
|
||||
"""Patient-level bilateral eval. Returns (y_true [N], probs [N, C])."""
|
||||
model.eval()
|
||||
y_chunks, p_chunks = [], []
|
||||
with torch.no_grad():
|
||||
for batch in loader:
|
||||
x1 = batch.get("image_1"); m1 = batch.get("matrix_1")
|
||||
x2 = batch.get("image_2"); m2 = batch.get("matrix_2")
|
||||
y = batch.get("label_1")
|
||||
if not (torch.is_tensor(x1) and torch.is_tensor(m1) and torch.is_tensor(x2) and torch.is_tensor(m2)):
|
||||
continue
|
||||
y_t = _to_label_tensor(y, device)
|
||||
p = F.softmax(model(x1.to(device), m1.to(device), x2.to(device), m2.to(device)), dim=1)
|
||||
y_chunks.append(y_t.cpu().numpy())
|
||||
p_chunks.append(p.cpu().numpy())
|
||||
if not y_chunks:
|
||||
return np.array([], dtype=np.int64), np.zeros((0, 0), dtype=np.float32)
|
||||
return np.concatenate(y_chunks), np.concatenate(p_chunks, axis=0)
|
||||
|
||||
|
||||
def collect_probs_fused(
|
||||
model: FusedEnsembleHT,
|
||||
loader: DataLoader,
|
||||
device: torch.device,
|
||||
) -> tuple[np.ndarray, np.ndarray]:
|
||||
"""Patient-level fused-head eval. Returns (y_true [N], probs [N, C])."""
|
||||
model.eval()
|
||||
y_chunks, p_chunks = [], []
|
||||
with torch.no_grad():
|
||||
for batch in loader:
|
||||
x1 = batch.get("image_1"); m1 = batch.get("matrix_1")
|
||||
x2 = batch.get("image_2"); m2 = batch.get("matrix_2")
|
||||
y = batch.get("label_1")
|
||||
if not (torch.is_tensor(x1) and torch.is_tensor(m1) and
|
||||
torch.is_tensor(x2) and torch.is_tensor(m2)):
|
||||
continue
|
||||
y_t = _to_label_tensor(y, device)
|
||||
p = F.softmax(model(x1.to(device), m1.to(device),
|
||||
x2.to(device), m2.to(device)), dim=1)
|
||||
y_chunks.append(y_t.cpu().numpy())
|
||||
p_chunks.append(p.cpu().numpy())
|
||||
if not y_chunks:
|
||||
return np.array([], dtype=np.int64), np.zeros((0, 0), dtype=np.float32)
|
||||
return np.concatenate(y_chunks), np.concatenate(p_chunks, axis=0)
|
||||
|
||||
|
||||
def collect_probs_single_components(
|
||||
model: SingleEyeHT,
|
||||
loader: DataLoader,
|
||||
device: torch.device,
|
||||
*,
|
||||
aggregate_patient: bool,
|
||||
return_logits: bool = False,
|
||||
):
|
||||
"""
|
||||
Collect fused/img/md probabilities (and optionally raw logits) for SingleEyeHT.
|
||||
- aggregate_patient=False: eye-level (OD/OS as independent samples)
|
||||
- aggregate_patient=True : patient-level (average OD/OS per head)
|
||||
- return_logits=False: returns (y, probs_f, probs_i, probs_m)
|
||||
- return_logits=True: returns (y, probs_f, probs_i, probs_m,
|
||||
logits_f, logits_i, logits_m)
|
||||
Note: logits are averaged across eyes when aggregate_patient=True,
|
||||
which is equivalent to averaging in logit space (before softmax).
|
||||
"""
|
||||
model.eval()
|
||||
y_chunks = []
|
||||
pf_chunks, pi_chunks, pm_chunks = [], [], []
|
||||
lf_chunks, li_chunks, lm_chunks = [], [], []
|
||||
with torch.no_grad():
|
||||
for batch in loader:
|
||||
x1 = batch.get("image_1"); m1 = batch.get("matrix_1")
|
||||
x2 = batch.get("image_2"); m2 = batch.get("matrix_2")
|
||||
y = batch.get("label_1")
|
||||
if not (torch.is_tensor(x1) and torch.is_tensor(m1) and torch.is_tensor(x2) and torch.is_tensor(m2)):
|
||||
continue
|
||||
y_t = _to_label_tensor(y, device)
|
||||
|
||||
def _per_eye(x, m):
|
||||
img_feats = None if model.bridge.mode == "metadata_only" else model.img_tower(x.to(device))
|
||||
md_feats = None if model.bridge.mode == "image_only" else model.md_tower(m.to(device))
|
||||
out_f, out_i, out_m = model.bridge(img_feats, md_feats)
|
||||
pf = F.softmax(out_f, dim=1)
|
||||
pi = F.softmax(out_i, dim=1) if out_i is not None else pf
|
||||
pm = F.softmax(out_m, dim=1) if out_m is not None else pf
|
||||
lf = out_f
|
||||
li = out_i if out_i is not None else out_f
|
||||
lm = out_m if out_m is not None else out_f
|
||||
return pf, pi, pm, lf, li, lm
|
||||
|
||||
pf_od, pi_od, pm_od, lf_od, li_od, lm_od = _per_eye(x1, m1)
|
||||
pf_os, pi_os, pm_os, lf_os, li_os, lm_os = _per_eye(x2, m2)
|
||||
|
||||
if aggregate_patient:
|
||||
y_chunks.append(y_t.cpu().numpy())
|
||||
pf_chunks.append((0.5 * (pf_od + pf_os)).cpu().numpy())
|
||||
pi_chunks.append((0.5 * (pi_od + pi_os)).cpu().numpy())
|
||||
pm_chunks.append((0.5 * (pm_od + pm_os)).cpu().numpy())
|
||||
lf_chunks.append((0.5 * (lf_od + lf_os)).cpu().numpy())
|
||||
li_chunks.append((0.5 * (li_od + li_os)).cpu().numpy())
|
||||
lm_chunks.append((0.5 * (lm_od + lm_os)).cpu().numpy())
|
||||
else:
|
||||
y_np = y_t.cpu().numpy()
|
||||
y_chunks += [y_np, y_np]
|
||||
pf_chunks += [pf_od.cpu().numpy(), pf_os.cpu().numpy()]
|
||||
pi_chunks += [pi_od.cpu().numpy(), pi_os.cpu().numpy()]
|
||||
pm_chunks += [pm_od.cpu().numpy(), pm_os.cpu().numpy()]
|
||||
lf_chunks += [lf_od.cpu().numpy(), lf_os.cpu().numpy()]
|
||||
li_chunks += [li_od.cpu().numpy(), li_os.cpu().numpy()]
|
||||
lm_chunks += [lm_od.cpu().numpy(), lm_os.cpu().numpy()]
|
||||
|
||||
if not y_chunks:
|
||||
z = np.zeros((0, 0), dtype=np.float32)
|
||||
if return_logits:
|
||||
return np.array([], dtype=np.int64), z, z, z, z, z, z
|
||||
return np.array([], dtype=np.int64), z, z, z
|
||||
|
||||
y = np.concatenate(y_chunks)
|
||||
pf = np.concatenate(pf_chunks, axis=0)
|
||||
pi = np.concatenate(pi_chunks, axis=0)
|
||||
pm = np.concatenate(pm_chunks, axis=0)
|
||||
if return_logits:
|
||||
lf = np.concatenate(lf_chunks, axis=0)
|
||||
li = np.concatenate(li_chunks, axis=0)
|
||||
lm = np.concatenate(lm_chunks, axis=0)
|
||||
return y, pf, pi, pm, lf, li, lm
|
||||
return y, pf, pi, pm
|
||||
|
||||
|
||||
def collect_probs_eye_level(
|
||||
model: "SingleEyeHT",
|
||||
loader: DataLoader,
|
||||
device: torch.device,
|
||||
*,
|
||||
return_ids: bool = False,
|
||||
):
|
||||
"""
|
||||
Collect fused/img/md probabilities from a single-eye loader (image_1/matrix_1 only).
|
||||
Used for eval-mode passes over the training set.
|
||||
|
||||
Returns (y, probs_f, probs_i, probs_m) or, when return_ids=True,
|
||||
(y, probs_f, probs_i, probs_m, sample_ids) where sample_ids is an
|
||||
array of strings like "2OD", "4OS".
|
||||
"""
|
||||
model.eval()
|
||||
y_chunks, pf_chunks, pi_chunks, pm_chunks, id_chunks = [], [], [], [], []
|
||||
with torch.no_grad():
|
||||
for batch in loader:
|
||||
x = batch.get("image_1")
|
||||
m = batch.get("matrix_1")
|
||||
y = batch.get("label_1")
|
||||
if not (torch.is_tensor(x) and torch.is_tensor(m)):
|
||||
continue
|
||||
y_t = _to_label_tensor(y, device)
|
||||
img_feats = None if model.bridge.mode == "metadata_only" else model.img_tower(x.to(device))
|
||||
md_feats = None if model.bridge.mode == "image_only" else model.md_tower(m.to(device))
|
||||
out_f, out_i, out_m = model.bridge(img_feats, md_feats)
|
||||
pf = F.softmax(out_f, dim=1)
|
||||
pi = F.softmax(out_i, dim=1) if out_i is not None else pf
|
||||
pm = F.softmax(out_m, dim=1) if out_m is not None else pf
|
||||
y_chunks.append(y_t.cpu().numpy())
|
||||
pf_chunks.append(pf.cpu().numpy())
|
||||
pi_chunks.append(pi.cpu().numpy())
|
||||
pm_chunks.append(pm.cpu().numpy())
|
||||
if return_ids:
|
||||
ids = batch.get("id_1", [""] * len(y_t))
|
||||
eyes = batch.get("eye_id_1", [""] * len(y_t))
|
||||
# ids/eyes may be tensors (int) or lists of strings
|
||||
if torch.is_tensor(ids):
|
||||
ids = ids.tolist()
|
||||
if torch.is_tensor(eyes):
|
||||
eyes = eyes.tolist()
|
||||
id_chunks.extend(
|
||||
[f"{pid}{eye}" for pid, eye in zip(ids, eyes)]
|
||||
)
|
||||
|
||||
if not y_chunks:
|
||||
z = np.zeros((0, 0), dtype=np.float32)
|
||||
empty_ids = np.array([], dtype=object)
|
||||
if return_ids:
|
||||
return np.array([], dtype=np.int64), z, z, z, empty_ids
|
||||
return np.array([], dtype=np.int64), z, z, z
|
||||
|
||||
y = np.concatenate(y_chunks)
|
||||
pf = np.concatenate(pf_chunks, axis=0)
|
||||
pi = np.concatenate(pi_chunks, axis=0)
|
||||
pm = np.concatenate(pm_chunks, axis=0)
|
||||
if return_ids:
|
||||
return y, pf, pi, pm, np.array(id_chunks, dtype=object)
|
||||
return y, pf, pi, pm
|
||||
|
||||
|
||||
def collect_probs_bilateral_components(
|
||||
model: BilateralHT,
|
||||
loader: DataLoader,
|
||||
device: torch.device,
|
||||
) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
|
||||
"""Collect fused/img/md probabilities for bilateral joint-tower model."""
|
||||
model.eval()
|
||||
y_chunks = []
|
||||
pf_chunks, pi_chunks, pm_chunks = [], [], []
|
||||
with torch.no_grad():
|
||||
for batch in loader:
|
||||
x1 = batch.get("image_1"); m1 = batch.get("matrix_1")
|
||||
x2 = batch.get("image_2"); m2 = batch.get("matrix_2")
|
||||
y = batch.get("label_1")
|
||||
if not (torch.is_tensor(x1) and torch.is_tensor(m1) and torch.is_tensor(x2) and torch.is_tensor(m2)):
|
||||
continue
|
||||
y_t = _to_label_tensor(y, device)
|
||||
joint_img, joint_md = model.encode_joint(
|
||||
x1.to(device), m1.to(device), x2.to(device), m2.to(device)
|
||||
)
|
||||
out_f, _, _ = model.bridge(joint_img, joint_md)
|
||||
out_i = model.aux_img(joint_img)
|
||||
out_m = model.aux_md(joint_md)
|
||||
y_chunks.append(y_t.cpu().numpy())
|
||||
pf_chunks.append(F.softmax(out_f, dim=1).cpu().numpy())
|
||||
pi_chunks.append(F.softmax(out_i, dim=1).cpu().numpy())
|
||||
pm_chunks.append(F.softmax(out_m, dim=1).cpu().numpy())
|
||||
if not y_chunks:
|
||||
z = np.zeros((0, 0), dtype=np.float32)
|
||||
return np.array([], dtype=np.int64), z, z, z
|
||||
return (
|
||||
np.concatenate(y_chunks),
|
||||
np.concatenate(pf_chunks, axis=0),
|
||||
np.concatenate(pi_chunks, axis=0),
|
||||
np.concatenate(pm_chunks, axis=0),
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# V2ModeComparisonOps — thin class wrapper kept for external import compat
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class V2ModeComparisonOps:
|
||||
"""Namespace wrapper kept for backward-compatibility imports."""
|
||||
|
||||
_set_requires_grad = staticmethod(_set_requires_grad)
|
||||
_set_single_phase = staticmethod(_set_single_phase)
|
||||
_set_bilateral_phase = staticmethod(_set_bilateral_phase)
|
||||
train_single_epoch = staticmethod(train_single_epoch)
|
||||
train_bilateral_epoch = staticmethod(train_bilateral_epoch)
|
||||
collect_probs_classic = staticmethod(collect_probs_classic)
|
||||
collect_probs_ensemble = staticmethod(collect_probs_ensemble)
|
||||
collect_probs_bilateral = staticmethod(collect_probs_bilateral)
|
||||
|
||||
@staticmethod
|
||||
def _to_label_tensor(labels, device: torch.device) -> torch.Tensor:
|
||||
return _to_label_tensor(labels, device)
|
||||
@@ -1,200 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Optional, Protocol
|
||||
|
||||
import pandas as pd
|
||||
|
||||
|
||||
@dataclass
|
||||
class PatientSplit:
|
||||
"""Patient-disjoint split definition for a fold."""
|
||||
|
||||
train: pd.DataFrame
|
||||
val: pd.DataFrame
|
||||
holdout: Optional[pd.DataFrame] = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class LoaderBundle:
|
||||
"""All loaders needed by a training run."""
|
||||
|
||||
train: Any
|
||||
val: Any
|
||||
holdout: Optional[Any] = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class FoldResult:
|
||||
"""Normalized fold output from trainer implementations."""
|
||||
|
||||
fold: int
|
||||
metrics: dict[str, Any]
|
||||
artifacts: dict[str, Any]
|
||||
|
||||
|
||||
class SplitManager(Protocol):
|
||||
def build_plans(
|
||||
self,
|
||||
*,
|
||||
clinical: Any,
|
||||
args: Any,
|
||||
profile: Optional[Any] = None,
|
||||
) -> list[PatientSplit]:
|
||||
...
|
||||
|
||||
|
||||
class GraphFactory(Protocol):
|
||||
def build(
|
||||
self,
|
||||
*,
|
||||
clinical: Any,
|
||||
args: Any,
|
||||
fold: int,
|
||||
profile: Optional[Any] = None,
|
||||
) -> Any:
|
||||
...
|
||||
|
||||
|
||||
class LoaderFactory(Protocol):
|
||||
def build(
|
||||
self,
|
||||
*,
|
||||
clinical: Any,
|
||||
split: PatientSplit,
|
||||
args: Any,
|
||||
fold: int,
|
||||
profile: Optional[Any] = None,
|
||||
) -> LoaderBundle:
|
||||
...
|
||||
|
||||
|
||||
class Trainer(Protocol):
|
||||
def fit(
|
||||
self,
|
||||
*,
|
||||
graph: Any,
|
||||
loaders: LoaderBundle,
|
||||
args: Any,
|
||||
fold: int,
|
||||
profile: Optional[Any] = None,
|
||||
) -> FoldResult:
|
||||
...
|
||||
|
||||
|
||||
class NetworkManager:
|
||||
"""
|
||||
V2 orchestration entrypoint.
|
||||
|
||||
This class is intentionally small and modular:
|
||||
- split policy is delegated to a SplitManager
|
||||
- graph assembly is delegated to a GraphFactory
|
||||
- dataloaders are delegated to a LoaderFactory
|
||||
- train/eval/checkpoint lifecycle is delegated to a Trainer
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
clinical: Any,
|
||||
args: Any,
|
||||
split_manager: SplitManager,
|
||||
graph_factory: GraphFactory,
|
||||
loader_factory: LoaderFactory,
|
||||
trainer: Trainer,
|
||||
profile: Optional[Any] = None,
|
||||
) -> None:
|
||||
self.clinical = clinical
|
||||
self.args = args
|
||||
self.split_manager = split_manager
|
||||
self.graph_factory = graph_factory
|
||||
self.loader_factory = loader_factory
|
||||
self.trainer = trainer
|
||||
self.profile = profile
|
||||
self._split_plans: Optional[list[PatientSplit]] = None
|
||||
|
||||
def run_fold(self, fold: int) -> FoldResult:
|
||||
plans = self._get_split_plans()
|
||||
if fold < 0 or fold >= len(plans):
|
||||
raise IndexError(f"Requested fold {fold} but only {len(plans)} fold plans are available")
|
||||
split = plans[fold]
|
||||
self._validate_patient_disjointness(split)
|
||||
self._validate_labels(split)
|
||||
|
||||
graph = self.graph_factory.build(
|
||||
clinical=self.clinical,
|
||||
args=self.args,
|
||||
fold=fold,
|
||||
profile=self.profile,
|
||||
)
|
||||
loaders = self.loader_factory.build(
|
||||
clinical=self.clinical,
|
||||
split=split,
|
||||
args=self.args,
|
||||
fold=fold,
|
||||
profile=self.profile,
|
||||
)
|
||||
return self.trainer.fit(
|
||||
graph=graph,
|
||||
loaders=loaders,
|
||||
args=self.args,
|
||||
fold=fold,
|
||||
profile=self.profile,
|
||||
)
|
||||
|
||||
def run_all_folds(self, n_splits: Optional[int] = None) -> list[FoldResult]:
|
||||
plans = self._get_split_plans()
|
||||
max_folds = len(plans)
|
||||
if n_splits is None:
|
||||
n = max_folds
|
||||
else:
|
||||
n = int(n_splits)
|
||||
if n < 1:
|
||||
raise ValueError("n_splits must be >= 1")
|
||||
if n > max_folds:
|
||||
raise ValueError(f"Requested {n} folds but only {max_folds} fold plans are available")
|
||||
return [self.run_fold(fold) for fold in range(n)]
|
||||
|
||||
def _get_split_plans(self) -> list[PatientSplit]:
|
||||
if self._split_plans is None:
|
||||
self._split_plans = self.split_manager.build_plans(
|
||||
clinical=self.clinical,
|
||||
args=self.args,
|
||||
profile=self.profile,
|
||||
)
|
||||
if not self._split_plans:
|
||||
raise ValueError("SplitManager returned no fold plans")
|
||||
return self._split_plans
|
||||
|
||||
def _validate_patient_disjointness(self, split: PatientSplit) -> None:
|
||||
train_ids = self._patient_ids(split.train)
|
||||
val_ids = self._patient_ids(split.val)
|
||||
holdout_ids = self._patient_ids(split.holdout) if split.holdout is not None else set()
|
||||
|
||||
if train_ids & val_ids:
|
||||
overlap = sorted(train_ids & val_ids)[:10]
|
||||
raise ValueError(f"Patient leakage between train/val: {overlap}")
|
||||
if train_ids & holdout_ids:
|
||||
overlap = sorted(train_ids & holdout_ids)[:10]
|
||||
raise ValueError(f"Patient leakage between train/holdout: {overlap}")
|
||||
if val_ids & holdout_ids:
|
||||
overlap = sorted(val_ids & holdout_ids)[:10]
|
||||
raise ValueError(f"Patient leakage between val/holdout: {overlap}")
|
||||
|
||||
def _validate_labels(self, split: PatientSplit) -> None:
|
||||
label_col = getattr(self.clinical, "label_col", None)
|
||||
if not label_col:
|
||||
return
|
||||
for name, df in (("train", split.train), ("val", split.val), ("holdout", split.holdout)):
|
||||
if df is None:
|
||||
continue
|
||||
if label_col not in df.columns:
|
||||
raise ValueError(f"{name} split is missing label column {label_col!r}")
|
||||
|
||||
@staticmethod
|
||||
def _patient_ids(df: Optional[pd.DataFrame]) -> set[Any]:
|
||||
if df is None or df.empty:
|
||||
return set()
|
||||
if "Patient ID" not in df.columns:
|
||||
raise ValueError("Split dataframes must include 'Patient ID'")
|
||||
return set(df["Patient ID"].tolist())
|
||||
@@ -1,240 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Callable, Dict, List, Optional
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
|
||||
from classes.v2.data_bundle import DataBundle
|
||||
|
||||
# ---- Pachymetry → IOP correction (per PAPILA Table 3) ----
|
||||
_PACHY_TABLE: Dict[int, int] = {
|
||||
475: +5,
|
||||
485: +4,
|
||||
495: +4,
|
||||
505: +3,
|
||||
515: +2,
|
||||
525: +1,
|
||||
535: +1,
|
||||
545: 0,
|
||||
555: -1,
|
||||
565: -1,
|
||||
575: -2,
|
||||
585: -3,
|
||||
595: -4,
|
||||
605: -4,
|
||||
615: -5,
|
||||
}
|
||||
_PACHY_KEYS = np.array(sorted(_PACHY_TABLE.keys()))
|
||||
|
||||
|
||||
def _nearest_pachy_key(x: float) -> int:
|
||||
idx = int(np.argmin(np.abs(_PACHY_KEYS - float(x))))
|
||||
return int(_PACHY_KEYS[idx])
|
||||
|
||||
|
||||
def _fit_perkins_converter(
|
||||
frames: List[pd.DataFrame], method: str
|
||||
) -> Callable[[float, Optional[float]], float]:
|
||||
"""
|
||||
Fit a Perkins→Pneumatic converter from pooled paired observations across all frames.
|
||||
Returns a callable: converter(perkins_value, pachymetry_value) -> float.
|
||||
Supported methods: "ratio", "ols", "lad", "multi".
|
||||
"""
|
||||
combined = pd.concat(frames, ignore_index=True)
|
||||
paired = combined.dropna(subset=["Pneumatic", "Perkins"])
|
||||
pneumatic = paired["Pneumatic"].values.astype(float)
|
||||
perkins = paired["Perkins"].values.astype(float)
|
||||
|
||||
if len(paired) == 0:
|
||||
raise ValueError("No paired Pneumatic+Perkins observations found; cannot fit converter.")
|
||||
|
||||
if method == "ratio":
|
||||
ratio = float((pneumatic / perkins).mean())
|
||||
def converter_ratio(p: float, pachy: Optional[float] = None) -> float:
|
||||
return p * ratio
|
||||
return converter_ratio
|
||||
|
||||
elif method == "ols":
|
||||
from scipy import stats as _stats
|
||||
slope, intercept, *_ = _stats.linregress(perkins, pneumatic)
|
||||
slope, intercept = float(slope), float(intercept)
|
||||
def converter_ols(p: float, pachy: Optional[float] = None) -> float:
|
||||
return p * slope + intercept
|
||||
return converter_ols
|
||||
|
||||
elif method == "lad":
|
||||
from scipy import stats as _stats
|
||||
from scipy.optimize import minimize as _minimize
|
||||
slope0, intercept0, *_ = _stats.linregress(perkins, pneumatic)
|
||||
def _lad_loss(params):
|
||||
a, b = params
|
||||
return np.abs(pneumatic - (a * perkins + b)).mean()
|
||||
res = _minimize(_lad_loss, x0=[slope0, intercept0], method="Nelder-Mead")
|
||||
slope, intercept = float(res.x[0]), float(res.x[1])
|
||||
def converter_lad(p: float, pachy: Optional[float] = None) -> float:
|
||||
return p * slope + intercept
|
||||
return converter_lad
|
||||
|
||||
elif method == "multi":
|
||||
from numpy.linalg import lstsq as _lstsq
|
||||
paired_multi = combined.dropna(subset=["Pneumatic", "Perkins", "Pachymetry"])
|
||||
if len(paired_multi) == 0:
|
||||
raise ValueError("No paired Pneumatic+Perkins+Pachymetry rows; cannot fit multi method.")
|
||||
pneu = paired_multi["Pneumatic"].values.astype(float)
|
||||
perk = paired_multi["Perkins"].values.astype(float)
|
||||
pachy_vals = paired_multi["Pachymetry"].values.astype(float)
|
||||
X = np.column_stack([perk, pachy_vals, np.ones(len(perk))])
|
||||
coeffs, *_ = _lstsq(X, pneu, rcond=None)
|
||||
slope, pachy_coef, intercept = float(coeffs[0]), float(coeffs[1]), float(coeffs[2])
|
||||
pachy_fallback = float(pachy_vals.mean())
|
||||
def converter_multi(p: float, pachy: Optional[float] = None) -> float:
|
||||
pv = pachy if (pachy is not None and not np.isnan(pachy)) else pachy_fallback
|
||||
return p * slope + pachy_coef * pv + intercept
|
||||
return converter_multi
|
||||
|
||||
else:
|
||||
raise ValueError(f"Unknown iop_corr_method: {method!r}. Choose ratio/ols/lad/multi.")
|
||||
|
||||
|
||||
def _pick_iop(row: pd.Series, converter: Callable) -> float:
|
||||
"""Prefer Pneumatic; convert Perkins to Pneumatic scale if Pneumatic is absent."""
|
||||
pneumatic = row.get("Pneumatic", np.nan)
|
||||
if not pd.isna(pneumatic):
|
||||
return float(pneumatic)
|
||||
perkins = row.get("Perkins", np.nan)
|
||||
if pd.isna(perkins):
|
||||
return np.nan
|
||||
pachy = row.get("Pachymetry", np.nan)
|
||||
return converter(float(perkins), None if pd.isna(pachy) else float(pachy))
|
||||
|
||||
|
||||
def _correct_iop(raw_iop: float, pachy: float) -> float:
|
||||
"""Return corrected IOP using nearest pachymetry bin; if pachy missing, return raw."""
|
||||
if pd.isna(raw_iop):
|
||||
return np.nan
|
||||
if pd.isna(pachy):
|
||||
return float(raw_iop)
|
||||
key = _nearest_pachy_key(float(pachy))
|
||||
return float(raw_iop) + float(_PACHY_TABLE[key])
|
||||
|
||||
|
||||
def _apply_iop_and_drop_md(
|
||||
df: pd.DataFrame,
|
||||
converter: Callable,
|
||||
drop_raw: bool = False,
|
||||
) -> pd.DataFrame:
|
||||
"""Add IOP_raw/IOP_corr and drop source IOP columns + VF_MD if present (in-place safe)."""
|
||||
df["IOP_raw"] = df.apply(lambda row: _pick_iop(row, converter), axis=1)
|
||||
pachy = df.get("Pachymetry", pd.Series(np.nan, index=df.index))
|
||||
df["IOP_corr"] = [
|
||||
_correct_iop(r, p) for r, p in zip(df["IOP_raw"].values, pachy.values)
|
||||
]
|
||||
drop_cols = [c for c in ("Pneumatic", "Perkins", "VF_MD") if c in df.columns]
|
||||
if drop_raw:
|
||||
drop_cols.append("IOP_raw")
|
||||
if drop_cols:
|
||||
df.drop(columns=drop_cols, inplace=True)
|
||||
return df
|
||||
|
||||
|
||||
def _canonicalize_eye_column(df: pd.DataFrame) -> None:
|
||||
if "eyeID" in df.columns:
|
||||
src = "eyeID"
|
||||
else:
|
||||
src = None
|
||||
for c in df.columns:
|
||||
if "eye" in c.lower():
|
||||
src = c
|
||||
break
|
||||
if src is None:
|
||||
df["eyeID"] = "OS"
|
||||
return
|
||||
|
||||
s = df[src]
|
||||
|
||||
def norm(v):
|
||||
if pd.isna(v):
|
||||
return None
|
||||
x = str(v).strip().upper()
|
||||
if x in {"OS", "L", "LEFT", "0"}:
|
||||
return "OS"
|
||||
if x in {"OD", "R", "RIGHT", "1"}:
|
||||
return "OD"
|
||||
try:
|
||||
num = int(float(x))
|
||||
return "OD" if num % 2 == 1 else "OS"
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
mapped = s.map(norm)
|
||||
uniq = {u for u in mapped.dropna().unique().tolist()}
|
||||
if not uniq.issubset({"OS", "OD"}):
|
||||
raise ValueError(f"eyeID must be binary; found values {sorted(uniq)}")
|
||||
df["eyeID"] = mapped.fillna("OS")
|
||||
|
||||
|
||||
def build_papila_data(
|
||||
*,
|
||||
image_dir: str,
|
||||
clinical_dir: str,
|
||||
label_col: str,
|
||||
cat_cols: List[str],
|
||||
n_splits: int = 5,
|
||||
random_seed: int = 42,
|
||||
iop_corr_method: str = "ratio",
|
||||
iop_drop_raw: bool = False,
|
||||
exclude_cols: Optional[List[str]] = None,
|
||||
) -> DataBundle:
|
||||
"""
|
||||
Build a DataBundle for PAPILA with dataset-specific preprocessing:
|
||||
- load OD/OS Excel sheets
|
||||
- normalize Patient ID
|
||||
- canonicalize eyeID
|
||||
- compute IOP_raw / IOP_corr, drop VF_MD
|
||||
- build feature typing & folds
|
||||
"""
|
||||
_exclude = list(exclude_cols) if exclude_cols else []
|
||||
|
||||
# Remove excluded cols from cat_cols too so the bundle doesn't try to encode them
|
||||
effective_cat_cols = [c for c in cat_cols if c not in _exclude]
|
||||
|
||||
bundle = DataBundle(
|
||||
image_dir=image_dir,
|
||||
clinical_dir=clinical_dir,
|
||||
label_col=label_col,
|
||||
patient_col="Patient ID",
|
||||
cat_cols=effective_cat_cols,
|
||||
n_splits=n_splits,
|
||||
random_seed=random_seed,
|
||||
filename_template="RET{pid:03d}{eye}.jpg",
|
||||
)
|
||||
|
||||
od = pd.read_excel(f"{clinical_dir}/patient_data_od.xlsx", header=1)
|
||||
od["eyeID"] = "OD"
|
||||
os = pd.read_excel(f"{clinical_dir}/patient_data_os.xlsx", header=1)
|
||||
os["eyeID"] = "OS"
|
||||
|
||||
for frame in (od, os):
|
||||
if "Patient ID" not in frame.columns and "ID" in frame.columns:
|
||||
frame.rename(columns={"ID": "Patient ID"}, inplace=True)
|
||||
frame["Patient ID"] = frame["Patient ID"].astype(str).str.extract(r"(\d+)")[0].astype(int)
|
||||
_canonicalize_eye_column(frame)
|
||||
|
||||
bundle.add_df(od, id_column="ID", exclude_cols=_exclude or None)
|
||||
bundle.add_df(os, id_column="ID", exclude_cols=_exclude or None)
|
||||
|
||||
converter = _fit_perkins_converter(bundle.frames, method=iop_corr_method)
|
||||
for i in range(len(bundle.frames)):
|
||||
bundle.frames[i] = _apply_iop_and_drop_md(
|
||||
bundle.frames[i], converter=converter, drop_raw=iop_drop_raw
|
||||
)
|
||||
|
||||
bundle._refresh_master_df(exclude_cols=_exclude or None)
|
||||
bundle._infer_or_validate_feature_types(exclude_cols=_exclude or None)
|
||||
bundle._compute_numeric_stats()
|
||||
bundle._build_cat_maps()
|
||||
bundle._compute_feature_dim()
|
||||
bundle._build_kfold_indices()
|
||||
|
||||
return bundle
|
||||
@@ -1,61 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Iterable, Optional
|
||||
|
||||
import pandas as pd
|
||||
|
||||
from classes.v2.data_bundle import DataBundle
|
||||
from classes.v2.papila_builders import build_papila_data
|
||||
|
||||
|
||||
@dataclass
|
||||
class PapilaData:
|
||||
"""
|
||||
V2-friendly wrapper around the DataBundle pipeline.
|
||||
|
||||
Keeps all formatting/normalization behavior from build_papila_clinical,
|
||||
but exposes a minimal surface area for the V2 engine.
|
||||
"""
|
||||
|
||||
clinical: DataBundle
|
||||
patient_col: str = "Patient ID"
|
||||
|
||||
@property
|
||||
def df(self) -> pd.DataFrame:
|
||||
return self.clinical.df
|
||||
|
||||
@property
|
||||
def label_col(self) -> str:
|
||||
return self.clinical.label_col
|
||||
|
||||
@property
|
||||
def feature_dim(self) -> int:
|
||||
return self.clinical.feature_dim
|
||||
|
||||
def get_image_path(self, row: pd.Series):
|
||||
return self.clinical.get_image_path(row)
|
||||
|
||||
def vectorize_row(self, row: pd.Series):
|
||||
return self.clinical.vectorize_row(row)
|
||||
|
||||
@classmethod
|
||||
def from_dirs(
|
||||
cls,
|
||||
*,
|
||||
image_dir: str,
|
||||
clinical_dir: str,
|
||||
label_col: str,
|
||||
cat_cols: Iterable[str],
|
||||
n_splits: int = 5,
|
||||
random_seed: int = 42,
|
||||
) -> "PapilaData":
|
||||
clinical = build_papila_data(
|
||||
image_dir=image_dir,
|
||||
clinical_dir=clinical_dir,
|
||||
label_col=label_col,
|
||||
cat_cols=list(cat_cols),
|
||||
n_splits=n_splits,
|
||||
random_seed=random_seed,
|
||||
)
|
||||
return cls(clinical=clinical)
|
||||
@@ -1,194 +0,0 @@
|
||||
"""PredictionStore — unified per-epoch prediction tensor across all folds.
|
||||
|
||||
Tensor shape: (n_folds, n_epochs, n_samples, n_heads, n_classes)
|
||||
|
||||
The meaning of "sample" depends on tower_mode:
|
||||
single — each eye is a sample; sample_ids like "5OD", "14OS"
|
||||
ensemble — each patient is a sample; sample_ids like "5", "14"
|
||||
fused — same as ensemble
|
||||
bilateral— same as ensemble
|
||||
|
||||
Head names by mode:
|
||||
single : ["fused", "img", "md"]
|
||||
ensemble : ["od_fused", "od_img", "od_md", "os_fused", "os_img", "os_md"]
|
||||
fused : ["od_fused", "od_img", "od_md", "os_fused", "os_img", "os_md", "bilat_fused"]
|
||||
bilateral : ["fused", "img_joint", "md_joint"]
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
from typing import Sequence
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
def head_names_for_mode(tower_mode: str, *, fused_head: bool = False) -> list[str]:
|
||||
"""Return canonical head name list for a given tower_mode."""
|
||||
if tower_mode in ("single", "classic"):
|
||||
return ["fused", "img", "md"]
|
||||
if tower_mode == "ensemble":
|
||||
names = ["od_fused", "od_img", "od_md", "os_fused", "os_img", "os_md"]
|
||||
return names + ["bilat_fused"] if fused_head else names
|
||||
if tower_mode == "bilateral":
|
||||
return ["fused", "img_joint", "md_joint"]
|
||||
raise ValueError(f"Unknown tower_mode: {tower_mode!r}")
|
||||
|
||||
|
||||
class PredictionStore:
|
||||
"""
|
||||
Stores per-epoch predictions for every sample, head, and fold in one tensor.
|
||||
|
||||
Usage
|
||||
-----
|
||||
# Build once before the fold loop:
|
||||
store = PredictionStore(
|
||||
sample_ids=all_eye_or_patient_ids,
|
||||
y_true=all_labels,
|
||||
head_names=head_names_for_mode(tower_mode, fused_head=args.fused_head),
|
||||
n_folds=n_folds,
|
||||
n_epochs=total_epochs,
|
||||
n_classes=num_classes,
|
||||
)
|
||||
|
||||
# Inside each epoch, after collecting probs:
|
||||
store.record(fold, epoch, patient_ids_batch, "od_fused", probs_od)
|
||||
store.set_split(fold, train_ids, "train")
|
||||
store.set_split(fold, val_ids, "val")
|
||||
|
||||
# After all folds:
|
||||
store.save(run_dir / "predictions.npz")
|
||||
|
||||
# Load and query:
|
||||
store = PredictionStore.load("predictions.npz")
|
||||
store.query("5", "od_fused", fold=0) # → (n_epochs, n_classes)
|
||||
store.query("5", "od_fused") # → (n_folds, n_epochs, n_classes)
|
||||
store.get_split("5", fold=0) # → "train"
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
sample_ids: Sequence[str],
|
||||
y_true: Sequence[int],
|
||||
head_names: Sequence[str],
|
||||
n_folds: int,
|
||||
n_epochs: int,
|
||||
n_classes: int,
|
||||
):
|
||||
self.sample_ids = np.array(sample_ids, dtype=object)
|
||||
self.y_true = np.array(y_true, dtype=np.int64)
|
||||
self.head_names = np.array(head_names, dtype=object)
|
||||
self.n_folds = n_folds
|
||||
self.n_epochs = n_epochs
|
||||
self.n_classes = n_classes
|
||||
|
||||
n_samples = len(self.sample_ids)
|
||||
n_heads = len(self.head_names)
|
||||
|
||||
self.probs = np.full(
|
||||
(n_folds, n_epochs, n_samples, n_heads, n_classes),
|
||||
fill_value=np.nan,
|
||||
dtype=np.float32,
|
||||
)
|
||||
self.split = np.full((n_folds, n_samples), fill_value="", dtype=object)
|
||||
|
||||
self._sid_index: dict[str, int] = {str(s): i for i, s in enumerate(self.sample_ids)}
|
||||
self._head_index: dict[str, int] = {str(h): i for i, h in enumerate(self.head_names)}
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Writing
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def record(
|
||||
self,
|
||||
fold: int,
|
||||
epoch: int,
|
||||
sample_ids: Sequence[str],
|
||||
head_name: str,
|
||||
probs: np.ndarray,
|
||||
) -> None:
|
||||
"""Record a batch of predictions for one head.
|
||||
|
||||
Args:
|
||||
fold: 0-indexed fold number
|
||||
epoch: 0-indexed epoch number
|
||||
sample_ids: sequence of sample ID strings (length B)
|
||||
head_name: which head — must be in self.head_names
|
||||
probs: (B, n_classes) probability array
|
||||
"""
|
||||
head_idx = self._head_index.get(head_name)
|
||||
if head_idx is None:
|
||||
return # head not active in this mode — skip silently
|
||||
for i, sid in enumerate(sample_ids):
|
||||
s_idx = self._sid_index.get(str(sid))
|
||||
if s_idx is not None:
|
||||
self.probs[fold, epoch, s_idx, head_idx, :] = probs[i]
|
||||
|
||||
def set_split(
|
||||
self,
|
||||
fold: int,
|
||||
sample_ids: Sequence[str],
|
||||
label: str,
|
||||
) -> None:
|
||||
"""Label a group of samples as 'train', 'val', or 'holdout' for a fold."""
|
||||
for sid in sample_ids:
|
||||
s_idx = self._sid_index.get(str(sid))
|
||||
if s_idx is not None:
|
||||
self.split[fold, s_idx] = label
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Querying
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def query(
|
||||
self,
|
||||
sample_id: str,
|
||||
head_name: str,
|
||||
fold: int | None = None,
|
||||
) -> np.ndarray:
|
||||
"""Return epoch-level predictions for one sample + head.
|
||||
|
||||
Returns:
|
||||
fold=None → (n_folds, n_epochs, n_classes)
|
||||
fold=int → (n_epochs, n_classes)
|
||||
"""
|
||||
s_idx = self._sid_index[str(sample_id)]
|
||||
head_idx = self._head_index[str(head_name)]
|
||||
if fold is None:
|
||||
return self.probs[:, :, s_idx, head_idx, :]
|
||||
return self.probs[fold, :, s_idx, head_idx, :]
|
||||
|
||||
def get_split(self, sample_id: str, fold: int) -> str:
|
||||
"""Return the split label ('train'/'val'/'holdout') for a sample in a fold."""
|
||||
s_idx = self._sid_index[str(sample_id)]
|
||||
return str(self.split[fold, s_idx])
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Persistence
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def save(self, path: str | Path) -> None:
|
||||
np.savez_compressed(
|
||||
path,
|
||||
probs=self.probs,
|
||||
split=self.split,
|
||||
sample_ids=self.sample_ids,
|
||||
y_true=self.y_true,
|
||||
head_names=self.head_names,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def load(cls, path: str | Path) -> "PredictionStore":
|
||||
data = np.load(path, allow_pickle=True)
|
||||
probs = data["probs"]
|
||||
n_folds, n_epochs, _, _, n_classes = probs.shape
|
||||
store = cls(
|
||||
sample_ids=data["sample_ids"].tolist(),
|
||||
y_true=data["y_true"],
|
||||
head_names=data["head_names"].tolist(),
|
||||
n_folds=n_folds,
|
||||
n_epochs=n_epochs,
|
||||
n_classes=n_classes,
|
||||
)
|
||||
store.probs = probs
|
||||
store.split = data["split"]
|
||||
return store
|
||||
@@ -1,10 +0,0 @@
|
||||
from .base import DatasetProfile, SimpleDatasetProfile, SlotDescriptor
|
||||
from .papila import PapilaProfile, build_papila_profile
|
||||
|
||||
__all__ = [
|
||||
"DatasetProfile",
|
||||
"SimpleDatasetProfile",
|
||||
"SlotDescriptor",
|
||||
"PapilaProfile",
|
||||
"build_papila_profile",
|
||||
]
|
||||
@@ -1,53 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Protocol, Any
|
||||
|
||||
import pandas as pd
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class SlotDescriptor:
|
||||
"""
|
||||
Metadata for a generic batch slot key (e.g., image_1, matrix_1).
|
||||
"""
|
||||
|
||||
key: str
|
||||
kind: str
|
||||
description: str
|
||||
required: bool = True
|
||||
shape_hint: str | None = None
|
||||
|
||||
|
||||
class DatasetProfile(Protocol):
|
||||
"""
|
||||
Dataset-specific wiring that stays outside the generic V2 engine.
|
||||
"""
|
||||
|
||||
name: str
|
||||
patient_col: str
|
||||
label_col: str
|
||||
|
||||
def slot_descriptors(self) -> dict[str, SlotDescriptor]:
|
||||
...
|
||||
|
||||
def semantic_aliases(self) -> dict[str, str]:
|
||||
...
|
||||
|
||||
def build_samples(self, *, df: pd.DataFrame, clinical: Any) -> list[dict[str, Any]]:
|
||||
...
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class SimpleDatasetProfile:
|
||||
name: str
|
||||
patient_col: str
|
||||
label_col: str
|
||||
slots: dict[str, SlotDescriptor]
|
||||
aliases: dict[str, str]
|
||||
|
||||
def slot_descriptors(self) -> dict[str, SlotDescriptor]:
|
||||
return dict(self.slots)
|
||||
|
||||
def semantic_aliases(self) -> dict[str, str]:
|
||||
return dict(self.aliases)
|
||||
@@ -1,155 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import pandas as pd
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
from .base import SimpleDatasetProfile, SlotDescriptor
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class PapilaProfile(SimpleDatasetProfile):
|
||||
sample_mode: str = "patient" # "patient" | "eye"
|
||||
|
||||
def build_samples(self, *, df: pd.DataFrame, clinical) -> list[dict[str, object]]:
|
||||
samples: list[dict[str, object]] = []
|
||||
patient_col = self.patient_col
|
||||
label_col = self.label_col
|
||||
|
||||
mode = (self.sample_mode or "patient").lower()
|
||||
if mode not in {"patient", "eye"}:
|
||||
raise ValueError(f"Unsupported sample_mode '{self.sample_mode}'. Expected 'patient' or 'eye'.")
|
||||
|
||||
if mode == "eye":
|
||||
for _, row in df.iterrows():
|
||||
pid = row[patient_col]
|
||||
label = row[label_col]
|
||||
image_1 = clinical.get_image_path(row) if hasattr(clinical, "get_image_path") else None
|
||||
matrix_1 = clinical.vectorize_row(row) if hasattr(clinical, "vectorize_row") else None
|
||||
samples.append(
|
||||
{
|
||||
"id_1": pid,
|
||||
"label_1": label,
|
||||
"image_1": image_1,
|
||||
"matrix_1": matrix_1,
|
||||
}
|
||||
)
|
||||
return samples
|
||||
|
||||
for pid, grp in df.groupby(patient_col):
|
||||
label_series = grp[label_col]
|
||||
if label_series.empty:
|
||||
continue
|
||||
mode_vals = label_series.mode()
|
||||
label = mode_vals.iloc[0] if not mode_vals.empty else label_series.iloc[0]
|
||||
|
||||
def _row_for_eye(eye: str):
|
||||
if "eyeID" not in grp.columns:
|
||||
return None
|
||||
match = grp[grp["eyeID"].astype(str).str.upper() == eye]
|
||||
if match.empty:
|
||||
return None
|
||||
return match.iloc[0]
|
||||
|
||||
row_od = _row_for_eye("OD")
|
||||
row_os = _row_for_eye("OS")
|
||||
row_any = grp.iloc[0]
|
||||
|
||||
image_1 = clinical.get_image_path(row_od) if row_od is not None else None
|
||||
image_2 = clinical.get_image_path(row_os) if row_os is not None else None
|
||||
matrix_1 = clinical.vectorize_row(row_od) if row_od is not None else None
|
||||
matrix_2 = clinical.vectorize_row(row_os) if row_os is not None else None
|
||||
|
||||
if image_1 is None and hasattr(clinical, "get_image_path"):
|
||||
image_1 = clinical.get_image_path(row_any)
|
||||
if matrix_1 is None and hasattr(clinical, "vectorize_row"):
|
||||
matrix_1 = clinical.vectorize_row(row_any)
|
||||
|
||||
samples.append(
|
||||
{
|
||||
"id_1": pid,
|
||||
"label_1": label,
|
||||
"image_1": image_1,
|
||||
"image_2": image_2,
|
||||
"matrix_1": matrix_1,
|
||||
"matrix_2": matrix_2,
|
||||
}
|
||||
)
|
||||
return samples
|
||||
|
||||
|
||||
def build_papila_profile(
|
||||
*,
|
||||
patient_col: str = "Patient ID",
|
||||
label_col: str = "Diagnosis",
|
||||
sample_mode: str = "patient",
|
||||
) -> PapilaProfile:
|
||||
"""
|
||||
PAPILA-specific semantic map for generic V2 slot keys.
|
||||
|
||||
The engine remains slot-based (image_1/image_2/matrix_1/...).
|
||||
PAPILA meaning is captured here so run config stays dataset-local.
|
||||
"""
|
||||
|
||||
slots = {
|
||||
"id_1": SlotDescriptor(
|
||||
key="id_1",
|
||||
kind="id",
|
||||
description=f"Patient identifier column ({patient_col})",
|
||||
required=True,
|
||||
shape_hint="scalar",
|
||||
),
|
||||
"label_1": SlotDescriptor(
|
||||
key="label_1",
|
||||
kind="label",
|
||||
description=f"Diagnosis label column ({label_col})",
|
||||
required=True,
|
||||
shape_hint="scalar",
|
||||
),
|
||||
"image_1": SlotDescriptor(
|
||||
key="image_1",
|
||||
kind="image",
|
||||
description="Fundus image slot 1 (PAPILA: OD / right eye)",
|
||||
required=False,
|
||||
shape_hint="HWC or CHW",
|
||||
),
|
||||
"image_2": SlotDescriptor(
|
||||
key="image_2",
|
||||
kind="image",
|
||||
description="Fundus image slot 2 (PAPILA: OS / left eye)",
|
||||
required=False,
|
||||
shape_hint="HWC or CHW",
|
||||
),
|
||||
"matrix_1": SlotDescriptor(
|
||||
key="matrix_1",
|
||||
kind="matrix",
|
||||
description="Clinical metadata feature vector",
|
||||
required=False,
|
||||
shape_hint="[feature_dim]",
|
||||
),
|
||||
"matrix_2": SlotDescriptor(
|
||||
key="matrix_2",
|
||||
kind="matrix",
|
||||
description="Optional auxiliary tabular vector (reserved for experiments)",
|
||||
required=False,
|
||||
shape_hint="[feature_dim_2]",
|
||||
),
|
||||
}
|
||||
|
||||
aliases = {
|
||||
"id_1": "patient_id",
|
||||
"label_1": "diagnosis",
|
||||
"image_1": "od_fundus",
|
||||
"image_2": "os_fundus",
|
||||
"matrix_1": "clinical_metadata",
|
||||
"matrix_2": "aux_metadata",
|
||||
}
|
||||
|
||||
return PapilaProfile(
|
||||
name="papila",
|
||||
patient_col=patient_col,
|
||||
label_col=label_col,
|
||||
slots=slots,
|
||||
aliases=aliases,
|
||||
sample_mode=sample_mode,
|
||||
)
|
||||
@@ -1,131 +0,0 @@
|
||||
"""Result dataclasses and serialisation helpers for V2 fold outputs."""
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Optional
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Primitive helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _nan() -> float:
|
||||
return float("nan")
|
||||
|
||||
|
||||
def _f(v) -> Optional[float]:
|
||||
"""Round a scalar to 6 dp, return None for nan/None."""
|
||||
if v is None or (isinstance(v, float) and np.isnan(v)):
|
||||
return None
|
||||
return round(float(v), 6)
|
||||
|
||||
|
||||
def _sv(vec) -> Optional[str]:
|
||||
"""Serialise a vector to a pipe-separated string, or None."""
|
||||
if vec is None:
|
||||
return None
|
||||
return "|".join(f"{float(v):.4f}" for v in vec)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Dataclasses
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@dataclass
|
||||
class FoldResult:
|
||||
mode: str
|
||||
fold: int
|
||||
# Epoch where each model hit its peak val AUC
|
||||
best_epoch_single: int # SingleEyeHT — selected by ensemble val AUC
|
||||
best_epoch_bilat: int # BilateralHT — selected by bilateral val AUC
|
||||
# Classic (eye-level eval of SingleEyeHT; n = 2 * ensemble_val_n)
|
||||
classic_val_auc: float
|
||||
classic_val_acc: float
|
||||
classic_val_kappa: float
|
||||
classic_val_mcc: float
|
||||
classic_val_f1: float
|
||||
classic_val_recall: Optional[str]
|
||||
classic_val_ece: float
|
||||
classic_val_threshold: float
|
||||
classic_val_bias: Optional[str]
|
||||
classic_val_n: int
|
||||
# Ensemble (patient-level eval of same SingleEyeHT)
|
||||
ensemble_val_auc: float
|
||||
ensemble_val_acc: float
|
||||
ensemble_val_kappa: float
|
||||
ensemble_val_mcc: float
|
||||
ensemble_val_f1: float
|
||||
ensemble_val_recall: Optional[str]
|
||||
ensemble_val_ece: float
|
||||
ensemble_val_threshold: float
|
||||
ensemble_val_bias: Optional[str]
|
||||
ensemble_val_n: int
|
||||
# Bilateral (BilateralHT patient-level)
|
||||
bilat_val_auc: float
|
||||
bilat_val_acc: float
|
||||
bilat_val_kappa: float
|
||||
bilat_val_mcc: float
|
||||
bilat_val_f1: float
|
||||
bilat_val_recall: Optional[str]
|
||||
bilat_val_ece: float
|
||||
bilat_val_threshold: float
|
||||
bilat_val_bias: Optional[str]
|
||||
bilat_val_n: int
|
||||
# Holdout metrics (evaluated at best val epoch; nan if no holdout)
|
||||
classic_holdout_auc: float
|
||||
classic_holdout_acc: float
|
||||
ensemble_holdout_auc: float
|
||||
ensemble_holdout_acc: float
|
||||
bilat_holdout_auc: float
|
||||
bilat_holdout_acc: float
|
||||
holdout_n: int # number of holdout bilateral samples
|
||||
# Training sample counts
|
||||
single_train_n: int
|
||||
bilat_train_n: int
|
||||
# Fused head (ensemble + --fused-head; nan / None if --fused-head not used)
|
||||
fused_val_auc: float = float("nan")
|
||||
fused_val_acc: float = float("nan")
|
||||
fused_val_kappa: float = float("nan")
|
||||
fused_val_mcc: float = float("nan")
|
||||
fused_val_f1: float = float("nan")
|
||||
fused_val_recall: Optional[str] = None
|
||||
fused_val_ece: float = float("nan")
|
||||
fused_val_threshold: float = float("nan")
|
||||
fused_val_bias: Optional[str] = None
|
||||
fused_val_n: int = 0
|
||||
fused_holdout_auc: float = float("nan")
|
||||
fused_holdout_acc: float = float("nan")
|
||||
|
||||
|
||||
@dataclass
|
||||
class FoldArtifacts:
|
||||
y_true_classic: Optional[np.ndarray]
|
||||
probs_classic: Optional[np.ndarray]
|
||||
y_true_ensemble: Optional[np.ndarray]
|
||||
probs_ensemble: Optional[np.ndarray]
|
||||
y_true_bilat: Optional[np.ndarray]
|
||||
probs_bilat: Optional[np.ndarray]
|
||||
y_true_fused: Optional[np.ndarray] = None
|
||||
probs_fused: Optional[np.ndarray] = None
|
||||
probs_ensemble_img: Optional[np.ndarray] = None
|
||||
probs_ensemble_md: Optional[np.ndarray] = None
|
||||
probs_classic_img: Optional[np.ndarray] = None
|
||||
probs_classic_md: Optional[np.ndarray] = None
|
||||
# per-eye (pre-averaged) versions for ensemble mode
|
||||
y_true_ensemble_pereye: Optional[np.ndarray] = None
|
||||
probs_ensemble_pereye: Optional[np.ndarray] = None
|
||||
probs_ensemble_img_pereye: Optional[np.ndarray] = None
|
||||
probs_ensemble_md_pereye: Optional[np.ndarray] = None
|
||||
# raw logits (before softmax) — patient-level
|
||||
logits_ensemble: Optional[np.ndarray] = None
|
||||
logits_ensemble_img: Optional[np.ndarray] = None
|
||||
logits_ensemble_md: Optional[np.ndarray] = None
|
||||
logits_classic: Optional[np.ndarray] = None
|
||||
logits_classic_img: Optional[np.ndarray] = None
|
||||
logits_classic_md: Optional[np.ndarray] = None
|
||||
# raw logits — per-eye
|
||||
logits_ensemble_pereye: Optional[np.ndarray] = None
|
||||
logits_ensemble_img_pereye: Optional[np.ndarray] = None
|
||||
logits_ensemble_md_pereye: Optional[np.ndarray] = None
|
||||
@@ -1,159 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||
from typing import Any, Callable, Optional
|
||||
|
||||
from pathlib import Path
|
||||
from PIL import Image
|
||||
import numpy as np
|
||||
import torch
|
||||
from torch.utils.data import Dataset
|
||||
from torchvision import transforms
|
||||
|
||||
from .profiles.base import SlotDescriptor
|
||||
|
||||
|
||||
def slot_collate(batch: list[dict[str, Any]]) -> dict[str, Any]:
|
||||
if not batch:
|
||||
return {}
|
||||
keys = batch[0].keys()
|
||||
out: dict[str, Any] = {}
|
||||
for key in keys:
|
||||
vals = [item.get(key) for item in batch]
|
||||
if all(isinstance(v, torch.Tensor) for v in vals):
|
||||
try:
|
||||
out[key] = torch.stack(vals, dim=0)
|
||||
except Exception:
|
||||
out[key] = vals
|
||||
else:
|
||||
out[key] = vals
|
||||
return out
|
||||
|
||||
|
||||
class SlotDataset(Dataset):
|
||||
"""
|
||||
Dataset that yields dicts of slot-keyed values.
|
||||
|
||||
Sample records are expected to be dicts with keys matching slot descriptors.
|
||||
Image slots accept filesystem paths; matrix slots accept array-like values.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
samples: list[dict[str, Any]],
|
||||
slot_descriptors: dict[str, SlotDescriptor],
|
||||
*,
|
||||
image_transform: Optional[Callable[[Image.Image], torch.Tensor]] = None,
|
||||
matrix_transform: Optional[Callable[[Any], torch.Tensor]] = None,
|
||||
image_preprocessor: Optional[Callable[..., Image.Image]] = None,
|
||||
image_cache: Optional[dict[str, np.ndarray]] = None,
|
||||
) -> None:
|
||||
self.samples = samples
|
||||
self.slot_descriptors = slot_descriptors
|
||||
self.image_transform = image_transform or transforms.ToTensor()
|
||||
self.matrix_transform = matrix_transform or self._default_matrix_transform
|
||||
self.image_preprocessor = image_preprocessor
|
||||
self.image_cache = image_cache
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self.samples)
|
||||
|
||||
def __getitem__(self, idx: int) -> dict[str, Any]:
|
||||
record = self.samples[idx]
|
||||
out: dict[str, Any] = {}
|
||||
for key, desc in self.slot_descriptors.items():
|
||||
val = record.get(key)
|
||||
if desc.kind == "image":
|
||||
out[key] = self._load_image(val, required=desc.required)
|
||||
elif desc.kind == "matrix":
|
||||
out[key] = self._load_matrix(val, required=desc.required)
|
||||
else:
|
||||
out[key] = val
|
||||
return out
|
||||
|
||||
def _load_image(self, value: Any, *, required: bool) -> Optional[torch.Tensor]:
|
||||
if value is None:
|
||||
if required:
|
||||
raise ValueError("Missing required image slot")
|
||||
return None
|
||||
path = Path(value)
|
||||
cache_key = str(value)
|
||||
|
||||
if self.image_cache is not None:
|
||||
cached = self.image_cache.get(cache_key)
|
||||
if cached is not None:
|
||||
return self.image_transform(Image.fromarray(cached, mode="RGB"))
|
||||
|
||||
img = Image.open(path).convert("RGB")
|
||||
if self.image_preprocessor is not None:
|
||||
try:
|
||||
img = self.image_preprocessor(img, path)
|
||||
except TypeError:
|
||||
img = self.image_preprocessor(img)
|
||||
|
||||
if self.image_cache is not None:
|
||||
self.image_cache[cache_key] = np.asarray(img, dtype=np.uint8)
|
||||
|
||||
return self.image_transform(img)
|
||||
|
||||
def prebuild_image_cache(self, cache_workers: int = 0) -> None:
|
||||
"""Pre-populate image_cache for all samples in this dataset."""
|
||||
if self.image_cache is None:
|
||||
return
|
||||
paths = list({
|
||||
str(record[key])
|
||||
for record in self.samples
|
||||
for key, desc in self.slot_descriptors.items()
|
||||
if desc.kind == "image" and record.get(key) is not None
|
||||
})
|
||||
to_warm = [p for p in paths if p not in self.image_cache]
|
||||
if not to_warm:
|
||||
return
|
||||
print(
|
||||
f"[image_cache] warming {len(to_warm)} images "
|
||||
f"({len(paths) - len(to_warm)} already cached)",
|
||||
flush=True,
|
||||
)
|
||||
|
||||
def _warm_one(path_str: str) -> None:
|
||||
if path_str in self.image_cache:
|
||||
return
|
||||
p = Path(path_str)
|
||||
img = Image.open(p).convert("RGB")
|
||||
if self.image_preprocessor is not None:
|
||||
try:
|
||||
img = self.image_preprocessor(img, p)
|
||||
except TypeError:
|
||||
img = self.image_preprocessor(img)
|
||||
self.image_cache[path_str] = np.asarray(img, dtype=np.uint8)
|
||||
|
||||
try:
|
||||
from tqdm import tqdm
|
||||
except ImportError:
|
||||
tqdm = None
|
||||
|
||||
if cache_workers <= 1:
|
||||
it = tqdm(to_warm, desc="Warm image cache", unit="img") if tqdm else to_warm
|
||||
for path_str in it:
|
||||
_warm_one(path_str)
|
||||
else:
|
||||
with ThreadPoolExecutor(max_workers=cache_workers) as ex:
|
||||
futures = {ex.submit(_warm_one, p): p for p in to_warm}
|
||||
it = tqdm(as_completed(futures), total=len(futures), desc="Warm image cache", unit="img") if tqdm else as_completed(futures)
|
||||
for fut in it:
|
||||
fut.result()
|
||||
|
||||
def _load_matrix(self, value: Any, *, required: bool) -> Optional[torch.Tensor]:
|
||||
if value is None:
|
||||
if required:
|
||||
raise ValueError("Missing required matrix slot")
|
||||
return None
|
||||
return self.matrix_transform(value)
|
||||
|
||||
@staticmethod
|
||||
def _default_matrix_transform(value: Any) -> torch.Tensor:
|
||||
if isinstance(value, torch.Tensor):
|
||||
return value.float()
|
||||
if isinstance(value, np.ndarray):
|
||||
return torch.from_numpy(value.astype(np.float32, copy=False))
|
||||
return torch.as_tensor(value, dtype=torch.float32)
|
||||
@@ -1,197 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Iterable, Optional
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
from sklearn.model_selection import KFold, StratifiedKFold
|
||||
|
||||
from .network_manager import PatientSplit
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class SplitPlan:
|
||||
train_patient_ids: set[Any]
|
||||
val_patient_ids: set[Any]
|
||||
holdout_patient_ids: set[Any]
|
||||
|
||||
|
||||
def build_patient_split_plans(
|
||||
patient_ids: Iterable[Any],
|
||||
patient_labels: Iterable[Any],
|
||||
*,
|
||||
n_splits: int,
|
||||
seed: int,
|
||||
holdout_per_class: int = 0,
|
||||
holdout_seed: int = 123,
|
||||
) -> list[SplitPlan]:
|
||||
"""
|
||||
Core vector-based splitter.
|
||||
|
||||
Inputs are one row per patient:
|
||||
- patient_ids: unique patient IDs
|
||||
- patient_labels: one label per patient
|
||||
"""
|
||||
ids = np.asarray(list(patient_ids))
|
||||
labels = np.asarray(list(patient_labels))
|
||||
if ids.ndim != 1 or labels.ndim != 1:
|
||||
raise ValueError("patient_ids and patient_labels must be 1D arrays")
|
||||
if ids.size != labels.size:
|
||||
raise ValueError(f"Length mismatch: ids={ids.size}, labels={labels.size}")
|
||||
if ids.size == 0:
|
||||
raise ValueError("No patients available for splitting")
|
||||
if len(set(ids.tolist())) != ids.size:
|
||||
raise ValueError("patient_ids must be unique (one label per patient)")
|
||||
if n_splits < 2:
|
||||
raise ValueError("n_splits must be >= 2")
|
||||
|
||||
holdout_ids: set[Any] = set()
|
||||
if holdout_per_class > 0:
|
||||
rng = np.random.default_rng(holdout_seed)
|
||||
for label in np.unique(labels):
|
||||
idx = np.where(labels == label)[0]
|
||||
if idx.size == 0:
|
||||
continue
|
||||
n = min(holdout_per_class, idx.size)
|
||||
chosen = rng.choice(idx, size=n, replace=False)
|
||||
holdout_ids.update(ids[chosen].tolist())
|
||||
|
||||
keep_mask = ~np.isin(ids, list(holdout_ids))
|
||||
cv_ids = ids[keep_mask]
|
||||
cv_labels = labels[keep_mask]
|
||||
if cv_ids.size < n_splits:
|
||||
raise ValueError(
|
||||
f"Not enough patients ({cv_ids.size}) for n_splits={n_splits} after holdout removal"
|
||||
)
|
||||
|
||||
use_stratified = _can_stratify(cv_labels, n_splits)
|
||||
if use_stratified:
|
||||
splitter = StratifiedKFold(n_splits=n_splits, shuffle=True, random_state=seed)
|
||||
splits = list(splitter.split(cv_ids, cv_labels))
|
||||
else:
|
||||
splitter = KFold(n_splits=n_splits, shuffle=True, random_state=seed)
|
||||
splits = list(splitter.split(cv_ids))
|
||||
|
||||
plans: list[SplitPlan] = []
|
||||
for train_idx, val_idx in splits:
|
||||
plans.append(
|
||||
SplitPlan(
|
||||
train_patient_ids=set(cv_ids[train_idx].tolist()),
|
||||
val_patient_ids=set(cv_ids[val_idx].tolist()),
|
||||
holdout_patient_ids=set(holdout_ids),
|
||||
)
|
||||
)
|
||||
return plans
|
||||
|
||||
|
||||
class PatientFirstSplitManager:
|
||||
"""
|
||||
Patient-level splitter for V2.
|
||||
|
||||
Behavior:
|
||||
- Optional binary filtering happens first (labels in {0,1} only).
|
||||
- Optional holdout is sampled at the patient level (never per-eye rows).
|
||||
- K-fold split is built on remaining patients.
|
||||
- Returned dataframes contain all rows for each selected patient.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
patient_col: str = "Patient ID",
|
||||
label_col: Optional[str] = None,
|
||||
) -> None:
|
||||
self.patient_col = patient_col
|
||||
self.label_col = label_col
|
||||
|
||||
def build_plans(
|
||||
self,
|
||||
*,
|
||||
clinical: Any,
|
||||
args: Any,
|
||||
profile: Optional[Any] = None,
|
||||
) -> list[PatientSplit]:
|
||||
profile_label_col = getattr(profile, "label_col", None) if profile is not None else None
|
||||
profile_patient_col = getattr(profile, "patient_col", None) if profile is not None else None
|
||||
patient_col = profile_patient_col or self.patient_col
|
||||
label_col = self.label_col or profile_label_col or getattr(clinical, "label_col", None)
|
||||
if label_col is None:
|
||||
raise ValueError("Could not resolve label column from SplitManager or clinical.label_col")
|
||||
|
||||
if not hasattr(clinical, "df"):
|
||||
raise ValueError("Clinical object must expose a dataframe at .df")
|
||||
df_full = clinical.df.copy()
|
||||
self._validate_columns(df_full, label_col, patient_col=patient_col)
|
||||
|
||||
eval_mode = str(getattr(args, "eval_mode", "multiclass")).lower()
|
||||
if eval_mode == "binary":
|
||||
df_full = df_full[df_full[label_col].isin([0, 1])].reset_index(drop=True)
|
||||
|
||||
holdout_per_class = int(getattr(args, "holdout_per_class", 0) or 0)
|
||||
holdout_seed = int(getattr(args, "holdout_seed", 123))
|
||||
n_splits = int(getattr(args, "n_splits", 5))
|
||||
fold_seed = int(getattr(args, "fold_seed", 42))
|
||||
|
||||
patient_table = self._patient_label_table(df_full, label_col, patient_col=patient_col)
|
||||
plans = build_patient_split_plans(
|
||||
patient_ids=patient_table[patient_col].to_numpy(),
|
||||
patient_labels=patient_table["_label"].to_numpy(),
|
||||
n_splits=n_splits,
|
||||
seed=fold_seed,
|
||||
holdout_per_class=holdout_per_class,
|
||||
holdout_seed=holdout_seed,
|
||||
)
|
||||
|
||||
out: list[PatientSplit] = []
|
||||
for plan in plans:
|
||||
train_df = (
|
||||
df_full[df_full[patient_col].isin(plan.train_patient_ids)]
|
||||
.reset_index(drop=True)
|
||||
)
|
||||
val_df = (
|
||||
df_full[df_full[patient_col].isin(plan.val_patient_ids)]
|
||||
.reset_index(drop=True)
|
||||
)
|
||||
holdout_df = None
|
||||
if plan.holdout_patient_ids:
|
||||
holdout_df = (
|
||||
df_full[df_full[patient_col].isin(plan.holdout_patient_ids)]
|
||||
.reset_index(drop=True)
|
||||
)
|
||||
out.append(PatientSplit(train=train_df, val=val_df, holdout=holdout_df))
|
||||
return out
|
||||
|
||||
def _validate_columns(self, df: pd.DataFrame, label_col: str, patient_col: Optional[str] = None) -> None:
|
||||
pcol = patient_col or self.patient_col
|
||||
if pcol not in df.columns:
|
||||
raise ValueError(f"Missing required patient column: {pcol!r}")
|
||||
if label_col not in df.columns:
|
||||
raise ValueError(f"Missing required label column: {label_col!r}")
|
||||
|
||||
def _patient_label_table(
|
||||
self,
|
||||
df: pd.DataFrame,
|
||||
label_col: str,
|
||||
patient_col: Optional[str] = None,
|
||||
) -> pd.DataFrame:
|
||||
pcol = patient_col or self.patient_col
|
||||
grouped = (
|
||||
df.groupby(pcol, as_index=False)[label_col]
|
||||
.agg(lambda x: x.mode().iloc[0] if not x.mode().empty else x.iloc[0])
|
||||
.rename(columns={label_col: "_label"})
|
||||
.sort_values(pcol)
|
||||
.reset_index(drop=True)
|
||||
)
|
||||
if grouped.empty:
|
||||
raise ValueError("No patients available for splitting")
|
||||
return grouped
|
||||
|
||||
|
||||
def _can_stratify(labels: np.ndarray, n_splits: int) -> bool:
|
||||
if labels.size == 0:
|
||||
return False
|
||||
unique, counts = np.unique(labels, return_counts=True)
|
||||
if len(unique) < 2:
|
||||
return False
|
||||
return bool(np.all(counts >= n_splits))
|
||||
@@ -1,279 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
from typing import Optional
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
from torchvision import transforms
|
||||
|
||||
from classes.v2.backbones import BACKBONES, list_names, load_backbone_weights
|
||||
from classes.v2.SE_attention import SEBlock
|
||||
from classes.v2.data_bundle import DataBundle
|
||||
|
||||
|
||||
def build_backbone(name: str, freeze_ratio: float = 0.0, augment: bool = True):
|
||||
"""
|
||||
Operational builder:
|
||||
- instantiate with DEFAULT weights
|
||||
- strip classifier → features
|
||||
- apply ratio-based freezing over coarse blocks
|
||||
- return (model, out_dim, transform)
|
||||
"""
|
||||
key = (name or "").lower()
|
||||
if key not in BACKBONES:
|
||||
raise ValueError(f"Unsupported backbone '{name}'. Valid options: {list_names()}")
|
||||
|
||||
spec = BACKBONES[key]
|
||||
m = spec.ctor(weights=spec.weights_default)
|
||||
out_dim, m = spec.strip(m)
|
||||
load_backbone_weights(key, m)
|
||||
|
||||
# transforms: use the weights’ mean/std, but keep your augmentation pipeline
|
||||
mean = getattr(spec.weights_default, "meta", {}).get("mean", (0.485, 0.456, 0.406))
|
||||
std = getattr(spec.weights_default, "meta", {}).get("std", (0.229, 0.224, 0.225))
|
||||
crop = 299 if key == "inception_v3" else 224
|
||||
|
||||
if augment:
|
||||
transform = transforms.Compose(
|
||||
[
|
||||
transforms.Resize(256),
|
||||
transforms.CenterCrop(crop),
|
||||
transforms.RandomHorizontalFlip(),
|
||||
transforms.RandomVerticalFlip(),
|
||||
transforms.RandomRotation(15),
|
||||
transforms.ColorJitter(0.1, 0.1, 0.1, 0.05),
|
||||
transforms.ToTensor(),
|
||||
transforms.Normalize(mean=mean, std=std),
|
||||
]
|
||||
)
|
||||
else:
|
||||
transform = transforms.Compose(
|
||||
[
|
||||
transforms.Resize(256),
|
||||
transforms.CenterCrop(crop),
|
||||
transforms.ToTensor(),
|
||||
transforms.Normalize(mean=mean, std=std),
|
||||
]
|
||||
)
|
||||
|
||||
# ratio-based freezing: freeze earliest floor(N * freeze_ratio) blocks
|
||||
fr = max(0.0, min(1.0, float(freeze_ratio)))
|
||||
blocks = spec.blocks(m)
|
||||
n = len(blocks)
|
||||
freeze_n = int(math.floor(n * fr))
|
||||
for b in blocks[:freeze_n]:
|
||||
for p in b.parameters():
|
||||
p.requires_grad = False
|
||||
|
||||
return m, out_dim, transform
|
||||
|
||||
|
||||
class ImageTower(nn.Module):
|
||||
"""
|
||||
Vision backbone → pooled features.
|
||||
- backbone: one of list_names() (default 'efficientnet_b0')
|
||||
- always DEFAULT torchvision weights
|
||||
- freeze_ratio ∈ [0,1] freezes earliest floor(N*freeze_ratio) blocks
|
||||
- returns [N, out_dim] features from backbone forward
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
backbone: str = "efficientnet_b0",
|
||||
freeze_ratio: float = 0.0,
|
||||
use_se: bool = False,
|
||||
se_reduction: int = 16,
|
||||
se_pre_norm: bool = True,
|
||||
augment: bool = True,
|
||||
geometry_dim: int = 0,
|
||||
):
|
||||
super().__init__()
|
||||
self.backbone, base_dim, self.transform = build_backbone(
|
||||
backbone, freeze_ratio, augment=augment
|
||||
)
|
||||
self._name = backbone
|
||||
# Keep ordered blocks for dynamic freezing/thawing
|
||||
key = (self._name or "").lower()
|
||||
self._spec = BACKBONES[key]
|
||||
self._blocks = self._spec.blocks(self.backbone)
|
||||
# Optional tower-level SE over the final feature vector
|
||||
self.base_dim = base_dim
|
||||
self.geometry_dim = max(0, int(geometry_dim))
|
||||
self.out_dim = self.base_dim + self.geometry_dim
|
||||
self.tower_ln = nn.LayerNorm(self.base_dim) if se_pre_norm else nn.Identity()
|
||||
self.tower_se = (
|
||||
SEBlock(self.base_dim, reduction=se_reduction, residual=True)
|
||||
if use_se
|
||||
else None
|
||||
)
|
||||
|
||||
def forward(
|
||||
self, x: torch.Tensor, geometry: Optional[torch.Tensor] = None
|
||||
) -> torch.Tensor:
|
||||
y = self.backbone(x)
|
||||
# sanity: pooled features, not logits
|
||||
assert y.dim() == 2 and y.size(1) == self.base_dim, (
|
||||
f"Expected features [N,{self.base_dim}], got {tuple(y.shape)}"
|
||||
)
|
||||
if self.tower_se is not None:
|
||||
y, _ = self.tower_se(self.tower_ln(y))
|
||||
if self.geometry_dim > 0:
|
||||
if geometry is None or geometry.numel() == 0:
|
||||
geom = torch.zeros(
|
||||
y.size(0), self.geometry_dim, device=y.device, dtype=y.dtype
|
||||
)
|
||||
else:
|
||||
if geometry.dim() == 1:
|
||||
geom = geometry.unsqueeze(0)
|
||||
else:
|
||||
geom = geometry
|
||||
geom = geom.to(device=y.device, dtype=y.dtype)
|
||||
if geom.size(0) != y.size(0):
|
||||
raise ValueError(
|
||||
f"Geometry batch size mismatch: {geom.size(0)} vs {y.size(0)}"
|
||||
)
|
||||
if geom.size(1) != self.geometry_dim:
|
||||
raise ValueError(
|
||||
f"Expected geometry dim {self.geometry_dim}, got {geom.size(1)}"
|
||||
)
|
||||
y = torch.cat([y, geom], dim=1)
|
||||
return y
|
||||
|
||||
def set_freeze_ratio(self, ratio: float):
|
||||
"""Dynamically freeze earliest floor(N*ratio) backbone blocks."""
|
||||
r = max(0.0, min(1.0, float(ratio)))
|
||||
n = len(self._blocks)
|
||||
freeze_n = int(math.floor(n * r))
|
||||
# Unfreeze all first
|
||||
for b in self._blocks:
|
||||
for p in b.parameters():
|
||||
p.requires_grad = True
|
||||
# Freeze earliest blocks
|
||||
for b in self._blocks[:freeze_n]:
|
||||
for p in b.parameters():
|
||||
p.requires_grad = False
|
||||
|
||||
|
||||
class SiameseImageTower(nn.Module):
|
||||
"""
|
||||
Shared-weight bilateral image tower.
|
||||
|
||||
Runs OD and OS images through a single shared backbone, then returns
|
||||
cat([f_mean, f_delta]) where:
|
||||
f_mean = (f_od + f_os) / 2 -- shared bilateral representation
|
||||
f_delta = f_od - f_os -- asymmetry, signed OD-relative
|
||||
|
||||
out_dim = 2 * backbone_out_dim
|
||||
|
||||
When x_os is None (single-eye fallback):
|
||||
f_mean = f_od
|
||||
f_delta = zeros
|
||||
so the module degrades gracefully when only one eye is available.
|
||||
|
||||
The shared backbone means both eyes contribute to every gradient update,
|
||||
effectively doubling the training signal for the visual pathway without
|
||||
doubling parameters.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
backbone: str = "efficientnet_b0",
|
||||
freeze_ratio: float = 0.0,
|
||||
use_se: bool = False,
|
||||
se_reduction: int = 16,
|
||||
se_pre_norm: bool = True,
|
||||
augment: bool = True,
|
||||
):
|
||||
super().__init__()
|
||||
self._tower = ImageTower(
|
||||
backbone=backbone,
|
||||
freeze_ratio=freeze_ratio,
|
||||
use_se=use_se,
|
||||
se_reduction=se_reduction,
|
||||
se_pre_norm=se_pre_norm,
|
||||
augment=augment,
|
||||
geometry_dim=0,
|
||||
)
|
||||
self.out_dim = self._tower.out_dim * 2
|
||||
self.transform = self._tower.transform
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x_od: torch.Tensor,
|
||||
x_os: Optional[torch.Tensor] = None,
|
||||
) -> torch.Tensor:
|
||||
f_od = self._tower(x_od)
|
||||
if x_os is None:
|
||||
f_mean = f_od
|
||||
f_delta = torch.zeros_like(f_od)
|
||||
else:
|
||||
f_os = self._tower(x_os)
|
||||
f_mean = (f_od + f_os) * 0.5
|
||||
f_delta = f_od - f_os
|
||||
return torch.cat([f_mean, f_delta], dim=1)
|
||||
|
||||
def set_freeze_ratio(self, ratio: float) -> None:
|
||||
"""Delegates to the shared inner tower."""
|
||||
self._tower.set_freeze_ratio(ratio)
|
||||
|
||||
|
||||
class MDTower(nn.Module):
|
||||
"""MLP over DataBundle.vectorize_row outputs (convert to torch inside tower)."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
clinical_data: DataBundle,
|
||||
hidden_dim: int = 128,
|
||||
dropout: float = 0.1,
|
||||
use_se: bool = False,
|
||||
se_reduction: int = 16,
|
||||
se_pre_norm: bool = True,
|
||||
):
|
||||
super().__init__()
|
||||
self.feature_dim = clinical_data.feature_dim
|
||||
self.out_dim = hidden_dim
|
||||
# two-block MLP so we can optionally freeze/thaw per block
|
||||
self.block0 = nn.Sequential(
|
||||
nn.Linear(self.feature_dim, hidden_dim),
|
||||
nn.LayerNorm(hidden_dim),
|
||||
nn.ReLU(inplace=True),
|
||||
nn.Dropout(dropout),
|
||||
)
|
||||
self.block1 = nn.Sequential(
|
||||
nn.Linear(hidden_dim, hidden_dim),
|
||||
nn.ReLU(inplace=True),
|
||||
)
|
||||
self.net = nn.Sequential(self.block0, self.block1)
|
||||
self.tower_ln = nn.LayerNorm(hidden_dim) if se_pre_norm else nn.Identity()
|
||||
self.tower_se = (
|
||||
SEBlock(hidden_dim, reduction=se_reduction, residual=True)
|
||||
if use_se
|
||||
else None
|
||||
)
|
||||
|
||||
def forward(self, meta_np_or_torch) -> torch.Tensor:
|
||||
if isinstance(meta_np_or_torch, torch.Tensor):
|
||||
x = meta_np_or_torch
|
||||
else:
|
||||
x = torch.as_tensor(meta_np_or_torch, dtype=torch.float32)
|
||||
h = self.net(x)
|
||||
if self.tower_se is not None:
|
||||
h, _ = self.tower_se(self.tower_ln(h))
|
||||
return h
|
||||
|
||||
def set_freeze_ratio(self, ratio: float):
|
||||
"""Optionally freeze earliest blocks of the MLP."""
|
||||
r = max(0.0, min(1.0, float(ratio)))
|
||||
# Unfreeze all
|
||||
for p in self.block0.parameters():
|
||||
p.requires_grad = True
|
||||
for p in self.block1.parameters():
|
||||
p.requires_grad = True
|
||||
# Freeze earliest blocks based on ratio threshold
|
||||
if r >= 0.5:
|
||||
for p in self.block0.parameters():
|
||||
p.requires_grad = False
|
||||
if r >= 1.0:
|
||||
for p in self.block1.parameters():
|
||||
p.requires_grad = False
|
||||
@@ -1,320 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Callable, Iterable, Optional, Tuple, Union
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
from torchvision import transforms
|
||||
|
||||
from classes.v2.backbones import BACKBONES
|
||||
|
||||
|
||||
IMAGENET_MEAN: Tuple[float, float, float] = (0.485, 0.456, 0.406)
|
||||
IMAGENET_STD: Tuple[float, float, float] = (0.229, 0.224, 0.225)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ImageTransformConfig:
|
||||
"""
|
||||
Mirrors the hypertower v1 preprocessing:
|
||||
- Resize(256)
|
||||
- CenterCrop(crop)
|
||||
- Optional augmentations (H/V flip, rotation, color jitter)
|
||||
- ToTensor + Normalize(mean/std)
|
||||
"""
|
||||
|
||||
crop_size: int = 224
|
||||
resize_size: int = 256
|
||||
mean: Tuple[float, float, float] = IMAGENET_MEAN
|
||||
std: Tuple[float, float, float] = IMAGENET_STD
|
||||
augment: bool = True
|
||||
rotation_deg: int = 15
|
||||
color_jitter: Tuple[float, float, float, float] = (0.1, 0.1, 0.1, 0.05)
|
||||
hflip: bool = True
|
||||
vflip: bool = True
|
||||
|
||||
def build(self) -> transforms.Compose:
|
||||
ops = [
|
||||
transforms.Resize(self.resize_size),
|
||||
transforms.CenterCrop(self.crop_size),
|
||||
]
|
||||
if self.augment:
|
||||
if self.hflip:
|
||||
ops.append(transforms.RandomHorizontalFlip())
|
||||
if self.vflip:
|
||||
ops.append(transforms.RandomVerticalFlip())
|
||||
if self.rotation_deg:
|
||||
ops.append(transforms.RandomRotation(self.rotation_deg))
|
||||
if self.color_jitter:
|
||||
ops.append(transforms.ColorJitter(*self.color_jitter))
|
||||
ops.extend(
|
||||
[
|
||||
transforms.ToTensor(),
|
||||
transforms.Normalize(mean=self.mean, std=self.std),
|
||||
]
|
||||
)
|
||||
return transforms.Compose(ops)
|
||||
|
||||
|
||||
def backbone_transform_config(backbone_name: str, augment: bool = True) -> ImageTransformConfig:
|
||||
"""
|
||||
Build a transform config that matches v1 ImageTower/backbone preprocessing.
|
||||
Uses DEFAULT weights mean/std and InceptionV3 crop size when relevant.
|
||||
"""
|
||||
key = (backbone_name or "").lower()
|
||||
if key not in BACKBONES:
|
||||
raise ValueError(f"Unsupported backbone '{backbone_name}'.")
|
||||
spec = BACKBONES[key]
|
||||
mean = getattr(spec.weights_default, "meta", {}).get("mean", IMAGENET_MEAN)
|
||||
std = getattr(spec.weights_default, "meta", {}).get("std", IMAGENET_STD)
|
||||
crop = 299 if key == "inception_v3" else 224
|
||||
return ImageTransformConfig(crop_size=crop, mean=mean, std=std, augment=augment)
|
||||
|
||||
|
||||
def build_backbone_transform(backbone_name: str, augment: bool = True) -> transforms.Compose:
|
||||
return backbone_transform_config(backbone_name, augment=augment).build()
|
||||
|
||||
|
||||
def build_eval_transform(backbone: str) -> transforms.Compose:
|
||||
"""Deterministic eval transform matching backbone normalisation (no augmentation)."""
|
||||
return build_backbone_transform(backbone, augment=False)
|
||||
|
||||
|
||||
def build_imagenet_transform(augment: bool = True, crop_size: int = 224) -> transforms.Compose:
|
||||
return ImageTransformConfig(crop_size=crop_size, augment=augment).build()
|
||||
|
||||
|
||||
@dataclass
|
||||
class ResizeTransform:
|
||||
size: Union[int, Tuple[int, int]] = 256
|
||||
interpolation: int = Image.BILINEAR
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
self._op = transforms.Resize(self.size, interpolation=self.interpolation)
|
||||
|
||||
def __call__(self, image: Image.Image) -> Image.Image:
|
||||
return self._op(image)
|
||||
|
||||
|
||||
@dataclass
|
||||
class CenterCropTransform:
|
||||
size: Union[int, Tuple[int, int]] = 224
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
self._op = transforms.CenterCrop(self.size)
|
||||
|
||||
def __call__(self, image: Image.Image) -> Image.Image:
|
||||
return self._op(image)
|
||||
|
||||
|
||||
class UnetMaskProvider:
|
||||
"""
|
||||
Placeholder for a UNet-powered mask provider.
|
||||
This will be replaced once a UNet tower is wired in.
|
||||
"""
|
||||
|
||||
def __call__(self, image: Image.Image, image_path: Optional[str] = None):
|
||||
raise NotImplementedError("UNet mask provider is not wired yet.")
|
||||
|
||||
|
||||
@dataclass
|
||||
class ROICropTransform:
|
||||
"""
|
||||
Crop an image using a binary mask (GT or UNet).
|
||||
Expects a mask of the same spatial size as the image; nonzero pixels are ROI.
|
||||
"""
|
||||
|
||||
mask_source: str = "gt" # "gt" | "unet"
|
||||
mask_provider: Optional[Callable[[Image.Image, Optional[str]], np.ndarray]] = None
|
||||
scale: float = 2.5
|
||||
target_size: Optional[Tuple[int, int]] = (224, 224)
|
||||
fallback_to_original: bool = True
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
if self.mask_source not in {"gt", "unet"}:
|
||||
raise ValueError(f"mask_source must be 'gt' or 'unet', got '{self.mask_source}'.")
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
image: Image.Image,
|
||||
mask: Optional[Union[np.ndarray, Image.Image]] = None,
|
||||
image_path: Optional[str] = None,
|
||||
) -> Image.Image:
|
||||
resolved_mask = mask
|
||||
if resolved_mask is None and self.mask_provider is not None:
|
||||
resolved_mask = self.mask_provider(image, image_path)
|
||||
if resolved_mask is None:
|
||||
if self.fallback_to_original:
|
||||
return image
|
||||
raise ValueError("ROI crop requested but no mask provided.")
|
||||
|
||||
mask_arr = (
|
||||
np.asarray(resolved_mask)
|
||||
if not isinstance(resolved_mask, Image.Image)
|
||||
else np.array(resolved_mask)
|
||||
)
|
||||
if mask_arr.ndim == 3:
|
||||
mask_arr = mask_arr[..., 0]
|
||||
mask_arr = mask_arr > 0
|
||||
if not np.any(mask_arr):
|
||||
return image if self.fallback_to_original else image
|
||||
|
||||
ys, xs = np.where(mask_arr)
|
||||
y_min, y_max = ys.min(), ys.max()
|
||||
x_min, x_max = xs.min(), xs.max()
|
||||
cx = (x_min + x_max) / 2.0
|
||||
cy = (y_min + y_max) / 2.0
|
||||
width = (x_max - x_min + 1)
|
||||
height = (y_max - y_min + 1)
|
||||
size = max(width, height) * float(self.scale)
|
||||
|
||||
left = int(round(cx - size / 2))
|
||||
right = int(round(cx + size / 2))
|
||||
upper = int(round(cy - size / 2))
|
||||
lower = int(round(cy + size / 2))
|
||||
|
||||
left = max(0, left)
|
||||
upper = max(0, upper)
|
||||
right = min(image.width, right)
|
||||
lower = min(image.height, lower)
|
||||
crop = image.crop((left, upper, right, lower))
|
||||
if self.target_size is not None:
|
||||
crop = crop.resize(self.target_size, Image.BILINEAR)
|
||||
return crop
|
||||
|
||||
|
||||
@dataclass
|
||||
class JitterBundleTransform:
|
||||
"""
|
||||
Augmentations bundle: flips, rotation, color jitter.
|
||||
"""
|
||||
|
||||
hflip: bool = True
|
||||
vflip: bool = True
|
||||
rotation_deg: int = 15
|
||||
color_jitter: Optional[Tuple[float, float, float, float]] = (0.1, 0.1, 0.1, 0.05)
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
ops = []
|
||||
if self.hflip:
|
||||
ops.append(transforms.RandomHorizontalFlip())
|
||||
if self.vflip:
|
||||
ops.append(transforms.RandomVerticalFlip())
|
||||
if self.rotation_deg:
|
||||
ops.append(transforms.RandomRotation(self.rotation_deg))
|
||||
if self.color_jitter:
|
||||
ops.append(transforms.ColorJitter(*self.color_jitter))
|
||||
self._op = transforms.Compose(ops) if ops else None
|
||||
|
||||
def __call__(self, image: Image.Image) -> Image.Image:
|
||||
if self._op is None:
|
||||
return image
|
||||
return self._op(image)
|
||||
|
||||
|
||||
TRANSFORM_REGISTRY = {
|
||||
"resize": ResizeTransform,
|
||||
"roi_crop": ROICropTransform,
|
||||
"center_crop": CenterCropTransform,
|
||||
"jitter_bundle": JitterBundleTransform,
|
||||
}
|
||||
|
||||
|
||||
def _parse_color_jitter(value: Optional[Union[str, Iterable[float]]]) -> Optional[Tuple[float, float, float, float]]:
|
||||
if value is None:
|
||||
return None
|
||||
if isinstance(value, str):
|
||||
parts = [p.strip() for p in value.split(",") if p.strip()]
|
||||
if not parts:
|
||||
return None
|
||||
try:
|
||||
nums = [float(p) for p in parts]
|
||||
except ValueError:
|
||||
return None
|
||||
if len(nums) == 1:
|
||||
return (nums[0], nums[0], nums[0], nums[0])
|
||||
if len(nums) >= 4:
|
||||
return (nums[0], nums[1], nums[2], nums[3])
|
||||
return tuple(nums + [nums[-1]] * (4 - len(nums))) # pad to length 4
|
||||
try:
|
||||
vals = list(value)
|
||||
except TypeError:
|
||||
return None
|
||||
if not vals:
|
||||
return None
|
||||
vals = [float(v) for v in vals]
|
||||
if len(vals) == 1:
|
||||
return (vals[0], vals[0], vals[0], vals[0])
|
||||
if len(vals) >= 4:
|
||||
return (vals[0], vals[1], vals[2], vals[3])
|
||||
return tuple(vals + [vals[-1]] * (4 - len(vals)))
|
||||
|
||||
|
||||
def build_transform_chain(
|
||||
transform_specs: Iterable[object],
|
||||
*,
|
||||
backbone_name: str,
|
||||
augment: bool = True,
|
||||
mask_provider: Optional[Callable[[Image.Image, Optional[str]], np.ndarray]] = None,
|
||||
strict: bool = True,
|
||||
) -> transforms.Compose:
|
||||
"""
|
||||
Build an image transform pipeline from a list of transform specs plus the
|
||||
standard ToTensor + Normalize steps. This mirrors the V1 preprocessing
|
||||
but uses the explicit transform nodes from config.
|
||||
"""
|
||||
ops: list[Callable[[Image.Image], Image.Image]] = []
|
||||
for spec in transform_specs:
|
||||
transform_type = getattr(spec, "transform_type", None)
|
||||
params = getattr(spec, "params", None)
|
||||
if transform_type is None and isinstance(spec, dict):
|
||||
transform_type = spec.get("transformType") or spec.get("transform_type")
|
||||
params = spec
|
||||
params = params or {}
|
||||
|
||||
if transform_type == "resize":
|
||||
size = params.get("resizeSize", 256)
|
||||
ops.append(ResizeTransform(size=size))
|
||||
elif transform_type == "center_crop":
|
||||
size = params.get("centerCropSize", 224)
|
||||
ops.append(CenterCropTransform(size=size))
|
||||
elif transform_type == "jitter_bundle":
|
||||
if not augment:
|
||||
continue
|
||||
jitter = JitterBundleTransform(
|
||||
hflip=bool(params.get("jitterHFlip", True)),
|
||||
vflip=bool(params.get("jitterVFlip", True)),
|
||||
rotation_deg=int(params.get("jitterRotation", 15) or 0),
|
||||
color_jitter=_parse_color_jitter(params.get("jitterColor"))
|
||||
if params.get("jitterColorEnabled", True)
|
||||
else None,
|
||||
)
|
||||
ops.append(jitter)
|
||||
elif transform_type == "roi_crop":
|
||||
roi = ROICropTransform(
|
||||
mask_source=params.get("roiMaskSource", "gt"),
|
||||
mask_provider=mask_provider,
|
||||
scale=float(params.get("roiScale", 2.5)),
|
||||
target_size=(int(params.get("roiTargetSize", 224)), int(params.get("roiTargetSize", 224)))
|
||||
if params.get("roiTargetSize") is not None
|
||||
else None,
|
||||
fallback_to_original=bool(params.get("roiFallback", True)),
|
||||
)
|
||||
if roi.mask_provider is None and roi.mask_source == "unet":
|
||||
if strict:
|
||||
raise ValueError("ROI crop requires a mask provider for 'unet' source.")
|
||||
ops.append(roi)
|
||||
else:
|
||||
if strict:
|
||||
raise ValueError(f"Unsupported transform type: {transform_type!r}")
|
||||
|
||||
# Always end with tensor + normalize, using backbone defaults
|
||||
cfg = backbone_transform_config(backbone_name, augment=augment)
|
||||
ops.extend(
|
||||
[
|
||||
transforms.ToTensor(),
|
||||
transforms.Normalize(mean=cfg.mean, std=cfg.std),
|
||||
]
|
||||
)
|
||||
return transforms.Compose(ops)
|
||||
@@ -1,894 +0,0 @@
|
||||
"""U-Net based optic disc/cup segmenter for REFUGE + Papila."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
import os
|
||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||
from contextlib import suppress
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Iterable, List, Optional, Set, Tuple
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
from PIL import Image, ImageDraw, ImageOps
|
||||
from PIL.Image import Resampling
|
||||
from skimage import measure
|
||||
import torch
|
||||
from torch import nn
|
||||
from torch.utils.data import DataLoader, Dataset
|
||||
from torchvision import transforms
|
||||
from tqdm import tqdm
|
||||
|
||||
|
||||
@dataclass
|
||||
class ManifestEntry:
|
||||
sample_id: str
|
||||
dataset: str
|
||||
image_path: Path
|
||||
annotation_disc: Path
|
||||
annotation_cup: Path
|
||||
annotation_type_disc: str
|
||||
annotation_type_cup: str
|
||||
split: str # train / holdout / etc.
|
||||
|
||||
|
||||
class UNet(nn.Module):
|
||||
def __init__(
|
||||
self, in_channels: int = 3, base_channels: int = 32, out_channels: int = 2
|
||||
):
|
||||
super().__init__()
|
||||
self.enc1 = self._block(in_channels, base_channels)
|
||||
self.enc2 = self._block(base_channels, base_channels * 2)
|
||||
self.enc3 = self._block(base_channels * 2, base_channels * 4)
|
||||
self.enc4 = self._block(base_channels * 4, base_channels * 8)
|
||||
|
||||
self.pool = nn.MaxPool2d(2)
|
||||
self.bottleneck = self._block(base_channels * 8, base_channels * 16)
|
||||
|
||||
self.up4 = nn.ConvTranspose2d(
|
||||
base_channels * 16, base_channels * 8, 2, stride=2
|
||||
)
|
||||
self.dec4 = self._block(base_channels * 16, base_channels * 8)
|
||||
self.up3 = nn.ConvTranspose2d(base_channels * 8, base_channels * 4, 2, stride=2)
|
||||
self.dec3 = self._block(base_channels * 8, base_channels * 4)
|
||||
self.up2 = nn.ConvTranspose2d(base_channels * 4, base_channels * 2, 2, stride=2)
|
||||
self.dec2 = self._block(base_channels * 4, base_channels * 2)
|
||||
self.up1 = nn.ConvTranspose2d(base_channels * 2, base_channels, 2, stride=2)
|
||||
self.dec1 = self._block(base_channels * 2, base_channels)
|
||||
|
||||
self.out_conv = nn.Conv2d(base_channels, out_channels, kernel_size=1)
|
||||
|
||||
@staticmethod
|
||||
def _block(in_ch: int, out_ch: int) -> nn.Module:
|
||||
return nn.Sequential(
|
||||
nn.Conv2d(in_ch, out_ch, kernel_size=3, padding=1, bias=False),
|
||||
nn.BatchNorm2d(out_ch),
|
||||
nn.ReLU(inplace=True),
|
||||
nn.Conv2d(out_ch, out_ch, kernel_size=3, padding=1, bias=False),
|
||||
nn.BatchNorm2d(out_ch),
|
||||
nn.ReLU(inplace=True),
|
||||
)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
e1 = self.enc1(x)
|
||||
e2 = self.enc2(self.pool(e1))
|
||||
e3 = self.enc3(self.pool(e2))
|
||||
e4 = self.enc4(self.pool(e3))
|
||||
b = self.bottleneck(self.pool(e4))
|
||||
|
||||
d4 = self.up4(b)
|
||||
d4 = torch.cat([d4, e4], dim=1)
|
||||
d4 = self.dec4(d4)
|
||||
d3 = self.up3(d4)
|
||||
d3 = torch.cat([d3, e3], dim=1)
|
||||
d3 = self.dec3(d3)
|
||||
d2 = self.up2(d3)
|
||||
d2 = torch.cat([d2, e2], dim=1)
|
||||
d2 = self.dec2(d2)
|
||||
d1 = self.up1(d2)
|
||||
d1 = torch.cat([d1, e1], dim=1)
|
||||
d1 = self.dec1(d1)
|
||||
return self.out_conv(d1)
|
||||
|
||||
|
||||
class SegmentationDataset(Dataset):
|
||||
def __init__(
|
||||
self,
|
||||
entries: List[ManifestEntry],
|
||||
segmenter: "UNetSegmenter",
|
||||
augment: bool,
|
||||
) -> None:
|
||||
self.entries = entries
|
||||
self.segmenter = segmenter
|
||||
self.augment = augment
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self.entries)
|
||||
|
||||
def __getitem__(self, idx: int):
|
||||
entry = self.entries[idx]
|
||||
image = self.segmenter.load_preprocessed_image(entry)
|
||||
disc_mask, cup_mask = self.segmenter.load_masks(entry)
|
||||
|
||||
if self.augment:
|
||||
image = self.segmenter.jitter_image(image)
|
||||
image, disc_mask, cup_mask = self.segmenter.augment_geometric(
|
||||
image, disc_mask, cup_mask
|
||||
)
|
||||
image_tensor = transforms.ToTensor()(image)
|
||||
image_tensor = self.segmenter._normalize_tensor(image_tensor)
|
||||
|
||||
mask = np.stack([disc_mask, cup_mask], axis=0).astype(np.float32)
|
||||
mask_tensor = torch.from_numpy(mask)
|
||||
return image_tensor, mask_tensor
|
||||
|
||||
|
||||
class UNetSegmenter:
|
||||
def __init__(
|
||||
self,
|
||||
manifest_path: Path,
|
||||
device: Optional[str] = None,
|
||||
cup_weight: float = 1.0,
|
||||
disc_weight: float = 1.0,
|
||||
target_size: int = 512,
|
||||
val_ratio: float = 0.1,
|
||||
train_datasets: Optional[Iterable[str]] = None,
|
||||
val_datasets: Optional[Iterable[str]] = None,
|
||||
holdout_datasets: Optional[Iterable[str]] = None,
|
||||
normalize: str = "none",
|
||||
use_stronger_aug: bool = False,
|
||||
mask_cache_dir: Optional[Path] = None,
|
||||
image_cache_dir: Optional[Path] = None,
|
||||
in_memory_cache: bool = False,
|
||||
loader_workers: int = 0,
|
||||
) -> None:
|
||||
self.manifest_path = manifest_path
|
||||
self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
|
||||
self.cup_weight = cup_weight
|
||||
self.disc_weight = disc_weight
|
||||
self.target_size = target_size
|
||||
self.val_ratio = val_ratio
|
||||
self.normalize = (normalize or "none").lower()
|
||||
self.use_stronger_aug = bool(use_stronger_aug)
|
||||
self.mask_cache_dir = Path(mask_cache_dir).resolve() if mask_cache_dir else None
|
||||
if self.mask_cache_dir:
|
||||
self.mask_cache_dir.mkdir(parents=True, exist_ok=True)
|
||||
self.image_cache_dir = Path(image_cache_dir).resolve() if image_cache_dir else None
|
||||
if self.image_cache_dir:
|
||||
self.image_cache_dir.mkdir(parents=True, exist_ok=True)
|
||||
self.in_memory_cache = bool(in_memory_cache)
|
||||
self._mem_image_cache: dict[str, np.ndarray] = {}
|
||||
self._mem_mask_cache: dict[str, Tuple[np.ndarray, np.ndarray]] = {}
|
||||
self.loader_workers = max(0, int(loader_workers))
|
||||
|
||||
self.train_dataset_filter = self._normalize_filter(train_datasets)
|
||||
self.val_dataset_filter = self._normalize_filter(val_datasets)
|
||||
self.holdout_dataset_filter = self._normalize_filter(holdout_datasets)
|
||||
|
||||
self.model = UNet().to(self.device)
|
||||
self._manifest: List[ManifestEntry] = []
|
||||
self.train_entries: List[ManifestEntry] = []
|
||||
self.val_entries: List[ManifestEntry] = []
|
||||
self.holdout_entries: List[ManifestEntry] = []
|
||||
self.read_manifest()
|
||||
|
||||
def prebuild_in_memory_cache(
|
||||
self,
|
||||
*,
|
||||
cache_workers: int = 0,
|
||||
include_train: bool = True,
|
||||
include_val: bool = True,
|
||||
include_holdout: bool = False,
|
||||
) -> None:
|
||||
if not self.in_memory_cache:
|
||||
return
|
||||
selected: List[ManifestEntry] = []
|
||||
if include_train:
|
||||
selected.extend(self.train_entries)
|
||||
if include_val:
|
||||
selected.extend(self.val_entries)
|
||||
if include_holdout:
|
||||
selected.extend(self.holdout_entries)
|
||||
if not selected:
|
||||
return
|
||||
|
||||
# Deduplicate by cache key.
|
||||
dedup = {}
|
||||
for entry in selected:
|
||||
dedup[self._entry_cache_key(entry)] = entry
|
||||
entries = list(dedup.values())
|
||||
workers = max(0, int(cache_workers))
|
||||
print(
|
||||
f"[UNetSegmenter] prebuilding in-memory cache for {len(entries)} samples "
|
||||
f"(cache_workers={workers})",
|
||||
flush=True,
|
||||
)
|
||||
|
||||
def _warm_one(entry: ManifestEntry) -> None:
|
||||
self.load_preprocessed_image(entry)
|
||||
self.load_masks(entry)
|
||||
|
||||
if workers <= 1:
|
||||
for entry in tqdm(entries, desc="Warm cache", unit="sample"):
|
||||
_warm_one(entry)
|
||||
else:
|
||||
with ThreadPoolExecutor(max_workers=workers) as ex:
|
||||
futures = [ex.submit(_warm_one, entry) for entry in entries]
|
||||
for fut in tqdm(as_completed(futures), total=len(futures), desc="Warm cache", unit="sample"):
|
||||
fut.result()
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
def read_manifest(self) -> None:
|
||||
df = pd.read_csv(self.manifest_path)
|
||||
entries: List[ManifestEntry] = []
|
||||
for _, row in df.iterrows():
|
||||
entry = ManifestEntry(
|
||||
sample_id=row["sample_id"],
|
||||
dataset=row["dataset"],
|
||||
image_path=Path(row["image_path"]),
|
||||
annotation_disc=Path(row["annotation_disc"]),
|
||||
annotation_cup=Path(row["annotation_cup"]),
|
||||
annotation_type_disc=row["annotation_type_disc"],
|
||||
annotation_type_cup=row["annotation_type_cup"],
|
||||
split=row["split"],
|
||||
)
|
||||
entries.append(entry)
|
||||
self._manifest = entries
|
||||
self.holdout_entries = [e for e in entries if e.split == "holdout"]
|
||||
if self.holdout_dataset_filter is not None:
|
||||
self.holdout_entries = [
|
||||
e for e in self.holdout_entries if e.dataset in self.holdout_dataset_filter
|
||||
]
|
||||
|
||||
trainable = [e for e in entries if e.split != "holdout"]
|
||||
if self.train_dataset_filter is not None:
|
||||
trainable = [
|
||||
e for e in trainable if e.dataset in self.train_dataset_filter
|
||||
]
|
||||
|
||||
if not trainable:
|
||||
self.val_entries = []
|
||||
self.train_entries = []
|
||||
return
|
||||
|
||||
val_pool = trainable
|
||||
if self.val_dataset_filter is not None:
|
||||
filtered = [e for e in trainable if e.dataset in self.val_dataset_filter]
|
||||
if filtered:
|
||||
val_pool = filtered
|
||||
|
||||
if len(trainable) == 1:
|
||||
val_count = 0
|
||||
else:
|
||||
val_count = max(1, int(len(trainable) * self.val_ratio))
|
||||
val_count = min(val_count, len(val_pool), len(trainable) - 1)
|
||||
|
||||
selected_val: List[ManifestEntry] = []
|
||||
if val_count > 0:
|
||||
selected_val = list(val_pool[:val_count])
|
||||
self.val_entries = selected_val
|
||||
selected_ids = {id(item) for item in selected_val}
|
||||
self.train_entries = [e for e in trainable if id(e) not in selected_ids]
|
||||
|
||||
if not self.train_entries and trainable:
|
||||
# Fallback when filtering removed all train entries (e.g. val_count forced entire set)
|
||||
self.train_entries = trainable
|
||||
self.val_entries = []
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
def preprocess_image(self, image: Image.Image) -> Image.Image:
|
||||
return image.resize((self.target_size, self.target_size), Resampling.BILINEAR)
|
||||
|
||||
def jitter_image(self, image: Image.Image) -> Image.Image:
|
||||
# Photometric jitter only; geometric ops are applied jointly (image+mask)
|
||||
return transforms.ColorJitter(0.1, 0.1, 0.1, 0.05)(image)
|
||||
|
||||
def augment_geometric(
|
||||
self,
|
||||
image: Image.Image,
|
||||
disc_mask: np.ndarray,
|
||||
cup_mask: np.ndarray,
|
||||
) -> tuple[Image.Image, np.ndarray, np.ndarray]:
|
||||
if not self.use_stronger_aug:
|
||||
return image, disc_mask, cup_mask
|
||||
|
||||
img = image
|
||||
disc_pil = Image.fromarray((disc_mask > 0).astype(np.uint8) * 255)
|
||||
cup_pil = Image.fromarray((cup_mask > 0).astype(np.uint8) * 255)
|
||||
|
||||
# Random horizontal flip
|
||||
if np.random.rand() < 0.5:
|
||||
img = ImageOps.mirror(img)
|
||||
disc_pil = ImageOps.mirror(disc_pil)
|
||||
cup_pil = ImageOps.mirror(cup_pil)
|
||||
# Random vertical flip
|
||||
if np.random.rand() < 0.5:
|
||||
img = ImageOps.flip(img)
|
||||
disc_pil = ImageOps.flip(disc_pil)
|
||||
cup_pil = ImageOps.flip(cup_pil)
|
||||
# Random rotation (multiples of 90° to keep masks aligned)
|
||||
rotations = np.random.choice([0, 90, 180, 270])
|
||||
if rotations:
|
||||
img = img.rotate(rotations, expand=False)
|
||||
disc_pil = disc_pil.rotate(rotations, expand=False)
|
||||
cup_pil = cup_pil.rotate(rotations, expand=False)
|
||||
|
||||
disc_mask = (np.array(disc_pil) > 0).astype(np.float32)
|
||||
cup_mask = (np.array(cup_pil) > 0).astype(np.float32)
|
||||
return img, disc_mask, cup_mask
|
||||
|
||||
@staticmethod
|
||||
def _slugify(text: str) -> str:
|
||||
return "".join(ch if ch.isalnum() or ch in ("-", "_") else "_" for ch in text)
|
||||
|
||||
def _entry_cache_key(self, entry: ManifestEntry) -> str:
|
||||
return self._slugify(f"{entry.dataset}_{entry.sample_id}_sz{self.target_size}")
|
||||
|
||||
def _mask_cache_path(self, entry: ManifestEntry) -> Optional[Path]:
|
||||
if self.mask_cache_dir is None:
|
||||
return None
|
||||
slug = self._slugify(f"{entry.dataset}_{entry.sample_id}")
|
||||
fname = f"{slug}_sz{self.target_size}.npz"
|
||||
return self.mask_cache_dir / fname
|
||||
|
||||
def _image_cache_path(self, entry: ManifestEntry) -> Optional[Path]:
|
||||
if self.image_cache_dir is None:
|
||||
return None
|
||||
slug = self._slugify(f"{entry.dataset}_{entry.sample_id}")
|
||||
fname = f"{slug}_img_sz{self.target_size}.npz"
|
||||
return self.image_cache_dir / fname
|
||||
|
||||
def _load_image_cache(self, cache_path: Path) -> Optional[Image.Image]:
|
||||
try:
|
||||
data = np.load(str(cache_path), allow_pickle=False)
|
||||
arr = data["image"].astype(np.uint8, copy=False)
|
||||
if arr.ndim != 3 or arr.shape[2] != 3:
|
||||
return None
|
||||
return Image.fromarray(arr, mode="RGB")
|
||||
except Exception:
|
||||
with suppress(OSError, FileNotFoundError):
|
||||
cache_path.unlink()
|
||||
return None
|
||||
|
||||
def _save_image_cache(self, cache_path: Optional[Path], image: Image.Image) -> None:
|
||||
if cache_path is None:
|
||||
return
|
||||
cache_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
tmp_path = cache_path.with_suffix(cache_path.suffix + ".tmp.npz")
|
||||
try:
|
||||
arr = np.asarray(image, dtype=np.uint8)
|
||||
np.savez_compressed(tmp_path, image=arr)
|
||||
os.replace(tmp_path, cache_path)
|
||||
except Exception:
|
||||
with suppress(OSError, FileNotFoundError):
|
||||
tmp_path.unlink()
|
||||
|
||||
def load_preprocessed_image(self, entry: ManifestEntry) -> Image.Image:
|
||||
key = self._entry_cache_key(entry)
|
||||
if self.in_memory_cache:
|
||||
cached = self._mem_image_cache.get(key)
|
||||
if cached is not None:
|
||||
return Image.fromarray(cached, mode="RGB")
|
||||
cache_path = self._image_cache_path(entry)
|
||||
if cache_path and cache_path.exists():
|
||||
cached = self._load_image_cache(cache_path)
|
||||
if cached is not None:
|
||||
if self.in_memory_cache:
|
||||
self._mem_image_cache[key] = np.asarray(cached, dtype=np.uint8)
|
||||
return cached
|
||||
image = Image.open(entry.image_path).convert("RGB")
|
||||
image = self.preprocess_image(image)
|
||||
if self.in_memory_cache:
|
||||
self._mem_image_cache[key] = np.asarray(image, dtype=np.uint8)
|
||||
self._save_image_cache(cache_path, image)
|
||||
return image
|
||||
|
||||
def _load_mask_cache(self, cache_path: Path) -> Optional[Tuple[np.ndarray, np.ndarray]]:
|
||||
try:
|
||||
data = np.load(str(cache_path), allow_pickle=False)
|
||||
disc = data["disc"].astype(np.float32)
|
||||
cup = data["cup"].astype(np.float32)
|
||||
return disc, cup
|
||||
except Exception:
|
||||
with suppress(OSError, FileNotFoundError):
|
||||
cache_path.unlink()
|
||||
return None
|
||||
|
||||
def _save_mask_cache(
|
||||
self,
|
||||
cache_path: Optional[Path],
|
||||
disc_mask: np.ndarray,
|
||||
cup_mask: np.ndarray,
|
||||
) -> None:
|
||||
if cache_path is None:
|
||||
return
|
||||
cache_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
tmp_path = cache_path.with_suffix(cache_path.suffix + ".tmp.npz")
|
||||
try:
|
||||
np.savez_compressed(
|
||||
tmp_path,
|
||||
disc=disc_mask.astype(np.uint8),
|
||||
cup=cup_mask.astype(np.uint8),
|
||||
)
|
||||
os.replace(tmp_path, cache_path)
|
||||
except Exception:
|
||||
with suppress(OSError, FileNotFoundError):
|
||||
tmp_path.unlink()
|
||||
|
||||
def _normalize_tensor(self, tensor: torch.Tensor) -> torch.Tensor:
|
||||
if self.normalize == "per_image":
|
||||
mean = tensor.mean(dim=(1, 2), keepdim=True)
|
||||
std = tensor.std(dim=(1, 2), keepdim=True).clamp(min=1e-6)
|
||||
return (tensor - mean) / std
|
||||
if self.normalize == "imagenet":
|
||||
mean = torch.tensor([0.485, 0.456, 0.406]).view(-1, 1, 1)
|
||||
std = torch.tensor([0.229, 0.224, 0.225]).view(-1, 1, 1)
|
||||
return (tensor - mean) / std
|
||||
return tensor
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
def extract_masks_from_image(
|
||||
self,
|
||||
mask_path: Path,
|
||||
disc_color: Optional[tuple[int, int, int]] = None,
|
||||
cup_color: Optional[tuple[int, int, int]] = None,
|
||||
) -> Tuple[np.ndarray, Optional[np.ndarray], Tuple[int, int]]:
|
||||
raw = Image.open(mask_path)
|
||||
arr = np.array(raw)
|
||||
if arr.ndim == 2:
|
||||
h, w = arr.shape
|
||||
flat = arr.reshape(-1).astype(np.int64, copy=False)
|
||||
edges = np.concatenate([arr[0, :], arr[-1, :], arr[:, 0], arr[:, -1]], axis=0).astype(np.int64, copy=False)
|
||||
edge_counts = np.bincount(edges, minlength=256)
|
||||
bg_val = int(np.argmax(edge_counts))
|
||||
counts = np.bincount(flat, minlength=256)
|
||||
counts[bg_val] = 0
|
||||
vals = np.where(counts > 0)[0]
|
||||
if vals.size < 1:
|
||||
raise ValueError(f"Mask {mask_path} does not contain discernible labels")
|
||||
# Disc = ALL non-background pixels (full optic disc: rim + cup combined).
|
||||
# Previously this was rim-only, which caused the cup structural prior
|
||||
# (cup & disc) to produce empty cup masks since cup and rim don't overlap.
|
||||
disc_mask = (arr != bg_val).astype(np.uint8)
|
||||
# Cup = the darkest non-background value (0 in REFUGE = inner cup region).
|
||||
# Using min-value rather than frequency avoids swapping when cup area > rim area.
|
||||
cup_val = int(np.min(vals)) if vals.size > 1 else None
|
||||
cup_mask = (arr == cup_val).astype(np.uint8) if cup_val is not None else np.zeros_like(disc_mask, dtype=np.uint8)
|
||||
return disc_mask, cup_mask if cup_mask.any() else None, (w, h)
|
||||
|
||||
image = raw.convert("RGB")
|
||||
arr = np.array(image)
|
||||
h, w, c = arr.shape
|
||||
|
||||
if disc_color is None or cup_color is None:
|
||||
# Fast color discovery via NumPy (avoid Python-level per-pixel tuple counting).
|
||||
edges = np.concatenate(
|
||||
[arr[0, :, :], arr[-1, :, :], arr[:, 0, :], arr[:, -1, :]], axis=0
|
||||
)
|
||||
edge_colors, edge_counts = np.unique(edges.reshape(-1, c), axis=0, return_counts=True)
|
||||
bg_color_np = edge_colors[int(np.argmax(edge_counts))]
|
||||
|
||||
colors_np, counts_np = np.unique(arr.reshape(-1, c), axis=0, return_counts=True)
|
||||
keep = np.any(colors_np != bg_color_np.reshape(1, -1), axis=1)
|
||||
colors_np = colors_np[keep]
|
||||
counts_np = counts_np[keep]
|
||||
if colors_np.shape[0] < 1:
|
||||
raise ValueError(f"Mask {mask_path} does not contain discernible labels")
|
||||
order = np.argsort(-counts_np)
|
||||
colors_np = colors_np[order]
|
||||
disc_color = tuple(int(v) for v in colors_np[0].tolist())
|
||||
cup_color = (
|
||||
tuple(int(v) for v in colors_np[1].tolist())
|
||||
if colors_np.shape[0] > 1
|
||||
else None
|
||||
)
|
||||
|
||||
disc_mask = np.zeros((h, w), dtype=np.uint8)
|
||||
cup_mask = np.zeros((h, w), dtype=np.uint8)
|
||||
|
||||
if disc_color is not None:
|
||||
disc_mask[np.all(arr == disc_color, axis=-1)] = 1
|
||||
if cup_color is not None:
|
||||
cup_mask[np.all(arr == cup_color, axis=-1)] = 1
|
||||
|
||||
return disc_mask, cup_mask if cup_mask.any() else None, (w, h)
|
||||
|
||||
def load_contour_from_file(self, contour_path: Path) -> np.ndarray:
|
||||
# Fast path: contour files are typically CSV or whitespace-delimited x,y pairs.
|
||||
try:
|
||||
arr = np.loadtxt(str(contour_path), delimiter=",", comments="#", dtype=np.float32)
|
||||
except Exception:
|
||||
try:
|
||||
arr = np.loadtxt(str(contour_path), comments="#", dtype=np.float32)
|
||||
except Exception:
|
||||
return np.zeros((0, 2), dtype=np.float32)
|
||||
if arr.size == 0:
|
||||
return np.zeros((0, 2), dtype=np.float32)
|
||||
if arr.ndim == 1:
|
||||
if arr.shape[0] < 2:
|
||||
return np.zeros((0, 2), dtype=np.float32)
|
||||
arr = arr.reshape(1, -1)
|
||||
if arr.shape[1] < 2:
|
||||
return np.zeros((0, 2), dtype=np.float32)
|
||||
return arr[:, :2].astype(np.float32, copy=False)
|
||||
|
||||
def coords_to_mask(
|
||||
self,
|
||||
coords: Optional[np.ndarray],
|
||||
size: Tuple[int, int],
|
||||
) -> np.ndarray:
|
||||
if coords is None or len(coords) == 0:
|
||||
return np.zeros((self.target_size, self.target_size), dtype=np.float32)
|
||||
|
||||
width, height = map(int, size)
|
||||
target_shape = (height, width)
|
||||
arr = np.asarray(coords)
|
||||
if arr.size == 0:
|
||||
return np.zeros((self.target_size, self.target_size), dtype=np.float32)
|
||||
|
||||
if arr.ndim == 2 and arr.shape[-1] != 2:
|
||||
mask = (arr > 0).astype(np.uint8)
|
||||
return self._resize_mask(mask)
|
||||
|
||||
if arr.ndim > 2:
|
||||
arr = arr.reshape(-1, arr.shape[-1])
|
||||
arr = arr.astype(float, copy=False)
|
||||
if arr.shape[-1] != 2:
|
||||
raise ValueError(f"Expected coordinate pairs, got shape {arr.shape}")
|
||||
|
||||
points = [tuple(map(float, pt)) for pt in arr]
|
||||
if len(points) < 3:
|
||||
return np.zeros(target_shape, dtype=np.float32)
|
||||
|
||||
img = Image.new("L", size, 0)
|
||||
draw = ImageDraw.Draw(img)
|
||||
draw.polygon(points, outline=1, fill=1)
|
||||
mask = np.array(img, dtype=np.uint8)
|
||||
return self._resize_mask(mask)
|
||||
|
||||
def _resize_mask(self, mask: np.ndarray) -> np.ndarray:
|
||||
img = Image.fromarray((mask > 0).astype(np.uint8) * 255)
|
||||
img = img.resize((self.target_size, self.target_size), Resampling.NEAREST)
|
||||
return (np.array(img, dtype=np.uint8) > 0).astype(np.float32)
|
||||
|
||||
def load_masks(self, entry: ManifestEntry) -> Tuple[np.ndarray, np.ndarray]:
|
||||
key = self._entry_cache_key(entry)
|
||||
if self.in_memory_cache:
|
||||
cached = self._mem_mask_cache.get(key)
|
||||
if cached is not None:
|
||||
disc_u8, cup_u8 = cached
|
||||
return disc_u8.astype(np.float32), cup_u8.astype(np.float32)
|
||||
cache_path = self._mask_cache_path(entry)
|
||||
if cache_path and cache_path.exists():
|
||||
cached = self._load_mask_cache(cache_path)
|
||||
if cached is not None:
|
||||
if self.in_memory_cache:
|
||||
disc, cup = cached
|
||||
self._mem_mask_cache[key] = (
|
||||
disc.astype(np.uint8),
|
||||
cup.astype(np.uint8),
|
||||
)
|
||||
return cached
|
||||
|
||||
image = Image.open(entry.image_path)
|
||||
size = image.size
|
||||
|
||||
disc_coords = cup_coords = None
|
||||
if entry.annotation_type_disc == "mask":
|
||||
disc_coords, cup_coords_from_disc, size = self.extract_masks_from_image(
|
||||
entry.annotation_disc
|
||||
)
|
||||
if cup_coords_from_disc is not None:
|
||||
cup_coords = cup_coords_from_disc
|
||||
else:
|
||||
disc_coords = self.load_contour_from_file(entry.annotation_disc)
|
||||
|
||||
if entry.annotation_type_cup == "mask":
|
||||
_, cup_coords_from_cup, size_cup = self.extract_masks_from_image(
|
||||
entry.annotation_cup
|
||||
)
|
||||
if cup_coords_from_cup is not None:
|
||||
cup_coords = cup_coords_from_cup
|
||||
if disc_coords is None:
|
||||
disc_coords, _, size = self.extract_masks_from_image(
|
||||
entry.annotation_cup
|
||||
)
|
||||
else:
|
||||
size = size_cup
|
||||
else:
|
||||
cup_coords = self.load_contour_from_file(entry.annotation_cup)
|
||||
|
||||
disc_mask = self.coords_to_mask(disc_coords, size).astype(np.float32)
|
||||
cup_mask = self.coords_to_mask(cup_coords, size).astype(np.float32)
|
||||
if self.in_memory_cache:
|
||||
self._mem_mask_cache[key] = (
|
||||
disc_mask.astype(np.uint8),
|
||||
cup_mask.astype(np.uint8),
|
||||
)
|
||||
self._save_mask_cache(cache_path, disc_mask, cup_mask)
|
||||
return disc_mask, cup_mask
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
def build_loaders(self, batch_size: int = 4, num_workers: int = 0) -> Tuple[DataLoader, DataLoader]:
|
||||
train_ds = SegmentationDataset(self.train_entries, self, augment=True)
|
||||
val_ds = SegmentationDataset(self.val_entries, self, augment=False)
|
||||
train_loader = DataLoader(
|
||||
train_ds, batch_size=batch_size, shuffle=True, num_workers=num_workers, pin_memory=True
|
||||
)
|
||||
val_loader = DataLoader(
|
||||
val_ds, batch_size=batch_size, shuffle=False, num_workers=num_workers, pin_memory=True
|
||||
)
|
||||
return train_loader, val_loader
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
def dice_score(self, preds: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
|
||||
preds = (preds > 0.5).float()
|
||||
intersection = (preds * targets).sum(dim=(2, 3))
|
||||
union = preds.sum(dim=(2, 3)) + targets.sum(dim=(2, 3))
|
||||
dice = (2 * intersection + 1e-6) / (union + 1e-6)
|
||||
return dice.mean(dim=0)
|
||||
|
||||
def train(
|
||||
self,
|
||||
epochs: int = 40,
|
||||
batch_size: int = 4,
|
||||
lr: float = 1e-3,
|
||||
weight_decay: float = 1e-5,
|
||||
checkpoint_dir: Path = Path("models/unet_segmenter"),
|
||||
) -> None:
|
||||
print(
|
||||
f"[UNetSegmenter] training on device={self.device} "
|
||||
f"(epochs={epochs}, batch_size={batch_size}, workers={self.loader_workers})"
|
||||
)
|
||||
train_loader, val_loader = self.build_loaders(batch_size=batch_size, num_workers=self.loader_workers)
|
||||
optimizer = torch.optim.Adam(
|
||||
self.model.parameters(), lr=lr, weight_decay=weight_decay
|
||||
)
|
||||
criterion = nn.BCEWithLogitsLoss()
|
||||
best_dice = -math.inf
|
||||
checkpoint_dir.mkdir(parents=True, exist_ok=True)
|
||||
best_path = checkpoint_dir / "best.pt"
|
||||
|
||||
epoch_bar = tqdm(range(1, epochs + 1), desc="Epochs", unit="epoch")
|
||||
|
||||
for epoch in epoch_bar:
|
||||
self.model.train()
|
||||
batch_bar = tqdm(
|
||||
train_loader,
|
||||
desc=f"Train {epoch}/{epochs}",
|
||||
leave=False,
|
||||
unit="batch",
|
||||
total=len(train_loader),
|
||||
)
|
||||
train_loss_total = 0.0
|
||||
train_samples = 0
|
||||
for images, masks in batch_bar:
|
||||
images = images.to(self.device)
|
||||
masks = masks.to(self.device)
|
||||
optimizer.zero_grad()
|
||||
logits = self.model(images)
|
||||
loss_disc = criterion(logits[:, 0:1], masks[:, 0:1])
|
||||
loss_cup = criterion(logits[:, 1:2], masks[:, 1:2])
|
||||
loss = self.disc_weight * loss_disc + self.cup_weight * loss_cup
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
batch_size = images.size(0)
|
||||
train_loss_total += loss.item() * batch_size
|
||||
train_samples += batch_size
|
||||
|
||||
train_loss = (
|
||||
train_loss_total / train_samples if train_samples else float("nan")
|
||||
)
|
||||
|
||||
self.model.eval()
|
||||
dices = []
|
||||
val_bar = tqdm(
|
||||
val_loader,
|
||||
desc="Validate",
|
||||
leave=False,
|
||||
unit="batch",
|
||||
total=len(val_loader),
|
||||
)
|
||||
with torch.no_grad():
|
||||
for images, masks in val_bar:
|
||||
images = images.to(self.device)
|
||||
masks = masks.to(self.device)
|
||||
logits = self.model(images)
|
||||
probs = torch.sigmoid(logits)
|
||||
dice = self.dice_score(probs, masks)
|
||||
dices.append(dice.cpu())
|
||||
if dices:
|
||||
mean_dice = torch.stack(dices).mean(dim=0)
|
||||
disc_dice = mean_dice[0].item()
|
||||
cup_dice = mean_dice[1].item()
|
||||
weight_sum = self.disc_weight + self.cup_weight
|
||||
score = (
|
||||
(self.disc_weight * disc_dice + self.cup_weight * cup_dice)
|
||||
/ weight_sum
|
||||
if weight_sum
|
||||
else 0.0
|
||||
)
|
||||
epoch_bar.set_postfix(
|
||||
loss=f"{train_loss:.4f}",
|
||||
dice_disc=f"{disc_dice:.3f}",
|
||||
dice_cup=f"{cup_dice:.3f}",
|
||||
dice_w=f"{score:.3f}",
|
||||
)
|
||||
else:
|
||||
disc_dice = cup_dice = 0.0
|
||||
score = 0.0
|
||||
epoch_bar.set_postfix(loss=f"{train_loss:.4f}")
|
||||
|
||||
if score > best_dice:
|
||||
best_dice = score
|
||||
torch.save({"model": self.model.state_dict()}, best_path)
|
||||
|
||||
if best_path.exists():
|
||||
state = torch.load(best_path, map_location=self.device)
|
||||
self.model.load_state_dict(state["model"])
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
def evaluate_holdout(
|
||||
self, output_dir: Path = Path("analysis_data/segmenter_eval")
|
||||
) -> pd.DataFrame:
|
||||
return self.evaluate_dataset(split_filter={"holdout"}, output_dir=output_dir)
|
||||
|
||||
@staticmethod
|
||||
def overlay_masks(
|
||||
image: Image.Image, disc: np.ndarray, cup: np.ndarray
|
||||
) -> Image.Image:
|
||||
overlay = image.copy()
|
||||
disc_img = Image.fromarray((disc * 255).astype(np.uint8))
|
||||
cup_img = Image.fromarray((cup * 255).astype(np.uint8))
|
||||
disc_color = Image.new("RGBA", image.size, (255, 0, 0, 0))
|
||||
cup_color = Image.new("RGBA", image.size, (0, 255, 0, 0))
|
||||
disc_color.paste((255, 0, 0, 100), mask=disc_img)
|
||||
cup_color.paste((0, 255, 0, 100), mask=cup_img)
|
||||
overlay = overlay.convert("RGBA")
|
||||
overlay = Image.alpha_composite(overlay, disc_color)
|
||||
overlay = Image.alpha_composite(overlay, cup_color)
|
||||
return overlay.convert("RGB")
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
@staticmethod
|
||||
def _normalize_filter(values: Optional[Iterable[str]]) -> Optional[Set[str]]:
|
||||
if values is None:
|
||||
return None
|
||||
if isinstance(values, str):
|
||||
return {values}
|
||||
return {str(item) for item in values}
|
||||
|
||||
@staticmethod
|
||||
def _dice_from_masks(pred: np.ndarray, target: np.ndarray) -> float:
|
||||
pred = (pred > 0).astype(np.float32)
|
||||
target = (target > 0).astype(np.float32)
|
||||
intersection = float((pred * target).sum())
|
||||
denom = float(pred.sum() + target.sum())
|
||||
return (2.0 * intersection + 1e-6) / (denom + 1e-6)
|
||||
|
||||
def get_entries(
|
||||
self,
|
||||
dataset_filter: Optional[Iterable[str]] = None,
|
||||
split_filter: Optional[Iterable[str]] = None,
|
||||
) -> List[ManifestEntry]:
|
||||
dataset_set = self._normalize_filter(dataset_filter)
|
||||
split_set = self._normalize_filter(split_filter)
|
||||
entries = self._manifest
|
||||
if dataset_set is not None:
|
||||
entries = [e for e in entries if e.dataset in dataset_set]
|
||||
if split_set is not None:
|
||||
entries = [e for e in entries if e.split in split_set]
|
||||
return list(entries)
|
||||
|
||||
def evaluate_dataset(
|
||||
self,
|
||||
dataset_filter: Optional[Iterable[str]] = None,
|
||||
split_filter: Optional[Iterable[str]] = None,
|
||||
output_dir: Path = Path("analysis_data/segmenter_eval"),
|
||||
save_overlays: bool = True,
|
||||
metrics_path: Optional[Path] = None,
|
||||
threshold: float = 0.5,
|
||||
tta: bool = False,
|
||||
) -> pd.DataFrame:
|
||||
entries = self.get_entries(
|
||||
dataset_filter=dataset_filter, split_filter=split_filter
|
||||
)
|
||||
if not entries:
|
||||
return pd.DataFrame(
|
||||
columns=[
|
||||
"sample_id",
|
||||
"dataset",
|
||||
"split",
|
||||
"dice_disc",
|
||||
"dice_cup",
|
||||
]
|
||||
)
|
||||
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
if metrics_path is None:
|
||||
suffix_parts = []
|
||||
if dataset_filter is not None:
|
||||
suffix_parts.append("-".join(sorted(self._normalize_filter(dataset_filter))))
|
||||
if split_filter is not None:
|
||||
suffix_parts.append("-".join(sorted(self._normalize_filter(split_filter))))
|
||||
suffix = "_".join(part for part in suffix_parts if part)
|
||||
csv_name = f"metrics{'_' + suffix if suffix else ''}.csv"
|
||||
metrics_path = output_dir / csv_name
|
||||
|
||||
records = []
|
||||
self.model.eval()
|
||||
progress = tqdm(
|
||||
entries,
|
||||
desc="Evaluate",
|
||||
unit="sample",
|
||||
leave=False,
|
||||
)
|
||||
for entry in progress:
|
||||
orig_image = Image.open(entry.image_path).convert("RGB")
|
||||
image = self.preprocess_image(orig_image)
|
||||
tensor = transforms.ToTensor()(image)
|
||||
tensor = self._normalize_tensor(tensor)
|
||||
tensor = tensor.unsqueeze(0).to(self.device)
|
||||
with torch.no_grad():
|
||||
logits = self.model(tensor)
|
||||
if tta:
|
||||
t_h = torch.flip(tensor, dims=[3])
|
||||
log_h = self.model(t_h)
|
||||
log_h = torch.flip(log_h, dims=[3])
|
||||
t_v = torch.flip(tensor, dims=[2])
|
||||
log_v = self.model(t_v)
|
||||
log_v = torch.flip(log_v, dims=[2])
|
||||
logits = (logits + log_h + log_v) / 3.0
|
||||
probs = torch.sigmoid(logits)[0].cpu().numpy()
|
||||
|
||||
disc_pred = (probs[0] > threshold).astype(np.uint8)
|
||||
cup_pred = (probs[1] > threshold).astype(np.uint8)
|
||||
# Structural prior: cup within disc
|
||||
cup_pred = (cup_pred > 0) & (disc_pred > 0)
|
||||
cup_pred = cup_pred.astype(np.uint8)
|
||||
|
||||
disc_gt, cup_gt = self.load_masks(entry)
|
||||
disc_gt = disc_gt.astype(np.uint8)
|
||||
cup_gt = cup_gt.astype(np.uint8)
|
||||
|
||||
dice_disc = self._dice_from_masks(disc_pred, disc_gt)
|
||||
dice_cup = self._dice_from_masks(cup_pred, cup_gt)
|
||||
|
||||
records.append(
|
||||
{
|
||||
"sample_id": entry.sample_id,
|
||||
"dataset": entry.dataset,
|
||||
"split": entry.split,
|
||||
"dice_disc": dice_disc,
|
||||
"dice_cup": dice_cup,
|
||||
}
|
||||
)
|
||||
|
||||
progress.set_postfix(
|
||||
dice_disc=f"{dice_disc:.3f}", dice_cup=f"{dice_cup:.3f}"
|
||||
)
|
||||
|
||||
if save_overlays:
|
||||
overlay_gt = self.overlay_masks(image, disc_gt, cup_gt)
|
||||
overlay_pred = self.overlay_masks(image, disc_pred, cup_pred)
|
||||
combined = Image.new("RGB", (image.width * 2, image.height))
|
||||
combined.paste(overlay_gt, (0, 0))
|
||||
combined.paste(overlay_pred, (image.width, 0))
|
||||
combined.save(output_dir / f"{entry.sample_id}_eval.png")
|
||||
|
||||
metrics_df = pd.DataFrame(records)
|
||||
summary = metrics_df[["dice_disc", "dice_cup"]].mean()
|
||||
summary_row = {
|
||||
"sample_id": "__mean__",
|
||||
"dataset": "summary",
|
||||
"split": "summary",
|
||||
"dice_disc": summary["dice_disc"],
|
||||
"dice_cup": summary["dice_cup"],
|
||||
}
|
||||
metrics_with_summary = pd.concat(
|
||||
[metrics_df, pd.DataFrame([summary_row])], ignore_index=True
|
||||
)
|
||||
metrics_with_summary.to_csv(metrics_path, index=False)
|
||||
return metrics_with_summary
|
||||
@@ -1,47 +0,0 @@
|
||||
"""General-purpose utilities for the V2 hypertower pipeline."""
|
||||
from __future__ import annotations
|
||||
|
||||
import random as pyrandom
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import torch
|
||||
|
||||
|
||||
def seed_everything(seed: int) -> None:
|
||||
pyrandom.seed(seed)
|
||||
np.random.seed(seed)
|
||||
torch.manual_seed(seed)
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.manual_seed_all(seed)
|
||||
|
||||
|
||||
def choose_device(device_arg: str | None) -> torch.device:
|
||||
if device_arg and device_arg != "auto":
|
||||
return torch.device(device_arg)
|
||||
return torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
|
||||
|
||||
def _drop_mixed_label_patients(df: pd.DataFrame, *, patient_col: str, label_col: str):
|
||||
"""Remove patients whose rows carry conflicting labels. Returns (clean_df, mixed_pids)."""
|
||||
per_patient = (
|
||||
df.groupby(patient_col)[label_col]
|
||||
.agg(lambda s: set(pd.to_numeric(s, errors="coerce").dropna().astype(int).tolist()))
|
||||
)
|
||||
mixed = [pid for pid, labels in per_patient.items() if len(labels) > 1]
|
||||
if not mixed:
|
||||
return df, []
|
||||
return df[~df[patient_col].isin(mixed)].reset_index(drop=True), mixed
|
||||
|
||||
|
||||
def _relabel_mixed_patients_to_max(df: pd.DataFrame, *, patient_col: str, label_col: str):
|
||||
"""Set all rows for each patient to that patient's max observed label."""
|
||||
out = df.copy()
|
||||
labels = pd.to_numeric(out[label_col], errors="coerce")
|
||||
patient_max = labels.groupby(out[patient_col]).transform("max")
|
||||
changed_rows = int((labels != patient_max).fillna(False).sum())
|
||||
out[label_col] = patient_max.astype(int)
|
||||
per_patient_unique = out.groupby(patient_col)[label_col].nunique(dropna=True)
|
||||
still_mixed = per_patient_unique[per_patient_unique > 1].index.tolist()
|
||||
return out.reset_index(drop=True), changed_rows, still_mixed
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,32 @@
|
||||
# Future Directions
|
||||
|
||||
Deferred ideas that are out of scope for the current manuscript but worth revisiting once the paper is out. Nothing here changes headline numbers; the fused HyperTower already beats the Kovalyk-Borodyak reproduction comfortably. These are things we noticed while building `binocular_analog/` that could tighten the image-only or bilateral tower further.
|
||||
|
||||
## Image tower / bilateral fusion (from Kovalyk repro comparison)
|
||||
|
||||
Context: their end-to-end binocular ResNet-50 reproduction lands at ~0.86 bilateral AUC on PAPILA, which is above what our own `img_solo` bilateral image tower gets. The gap is not evidence that their architecture is better; it points at four training choices we haven't yet copied. Ranked by expected effect:
|
||||
|
||||
1. **Backbone: plain ImageNet V1 ResNet-50 instead of REFUGE-pretrained `refugelike`.**
|
||||
In the Kovalyk repro we saw V1 > V2 on PAPILA (opposite of ImageNet-accuracy prediction). Same logic could apply to REFUGE pretraining: the REFUGE checkpoint may over-specialize early features to a different fundus camera / cohort, and plain ImageNet may transfer more cleanly. Try `backbone: "resnet50"` on `img_solo` with everything else held fixed.
|
||||
|
||||
2. **Backpropagate the patient-level loss through the backbone.**
|
||||
In `img_solo`, the `hb` stage does not set `train_towers: true`, so the patient-level loss can only reshape the 4096→256 HyperBridge projection and the head. Kovalyk's BinoCNN is jointly trained end-to-end. Enable tower training on the `hb` stage (probably with a smaller LR to avoid destabilising the eye-level warmup) and see if the bilateral AUC moves.
|
||||
|
||||
3. **Wider patient-level fusion.**
|
||||
HyperBridge `embedding_mlp` compresses 4096→256. Kovalyk's head goes 4096→512→2 with a real hidden layer. Try `hidden_dim: 512` and stack a small MLP head; cheap ablation.
|
||||
|
||||
4. **Label smoothing + explicit class weighting in the head loss.**
|
||||
Kovalyk uses `CrossEntropyLoss(weight=inv_freq, label_smoothing=0.14)`. Ours is BCE-with-logits, no smoothing, no class weight. On a 3.5:1 imbalance at this dataset size, adding smoothing (~0.1) and a pos_weight tends to stabilise the minimizer. Cheap to try.
|
||||
|
||||
**NOT on the list:** disc-region ROI cropping (GT contour bbox or U-Net-derived bbox). Kovalyk pre-crops to a 299×299 square ROI, but our own crop-sweep experiments (`v4/results/experiments/backbone_replication/gtcrop_refugelike`, `unetcrop_refugelike`) showed cropping hurts on PAPILA. Do not add it back as an image-tower improvement.
|
||||
|
||||
## Compute / infrastructure (already documented)
|
||||
|
||||
See `memory/project_future_v4_compute_fixes.md` for the two deferred v4 speedups:
|
||||
- Lazy BCD tower forward (~33% nt savings)
|
||||
- Frozen-prefix embedding cache for hb/val (foundation for a future feature-cache service)
|
||||
|
||||
## Reproduction extensions (nice-to-have, not blocking)
|
||||
|
||||
- Fill in the remaining 6 CNN configs from Kovalyk-Borodyak (VGG16, InceptionV3, MobileNetV2, mono + bino each) if we ever want to report a fuller reproduction table rather than a single ResNet-50 headline.
|
||||
- Try the paper's freeze-count under a proper Ray Tune HPO context; our 0.86-0.88 vs their 0.764 is the persistent unexplained residual, and the freeze-mode ablation did not close it.
|
||||
@@ -1,293 +0,0 @@
|
||||
{
|
||||
"nodes": [
|
||||
{
|
||||
"id": "node_1770737490980_3a7n",
|
||||
"type": "data",
|
||||
"label": "Papila Images",
|
||||
"inputKey": "",
|
||||
"outputKey": "",
|
||||
"inputType": "",
|
||||
"inputIndex": "",
|
||||
"outputType": "image",
|
||||
"source": {
|
||||
"mode": "directory",
|
||||
"path": "FundusImages",
|
||||
"count": null
|
||||
},
|
||||
"selectedPath": null,
|
||||
"sourceRef": {
|
||||
"importId": "import_1770737490980_00lw",
|
||||
"slot": "image_dir",
|
||||
"mode": "directory"
|
||||
},
|
||||
"x": 160,
|
||||
"y": 760
|
||||
},
|
||||
{
|
||||
"id": "node_1770737490980_az3a",
|
||||
"type": "data",
|
||||
"label": "Papila Metadata",
|
||||
"inputKey": "",
|
||||
"outputKey": "",
|
||||
"inputType": "",
|
||||
"inputIndex": "",
|
||||
"outputType": "matrix",
|
||||
"source": {
|
||||
"mode": "dataframe",
|
||||
"name": "clinical.df",
|
||||
"rows": 488,
|
||||
"cols": 15
|
||||
},
|
||||
"selectedPath": null,
|
||||
"sourceRef": {
|
||||
"importId": "import_1770737490980_00lw",
|
||||
"slot": "clinical_df",
|
||||
"mode": "dataframe"
|
||||
},
|
||||
"x": 340,
|
||||
"y": 760
|
||||
},
|
||||
{
|
||||
"id": "node_1770737490980_xnkx",
|
||||
"type": "transform",
|
||||
"label": "Resize",
|
||||
"inputKey": "",
|
||||
"outputKey": "",
|
||||
"inputType": "image",
|
||||
"inputIndex": "1",
|
||||
"outputType": "image",
|
||||
"source": null,
|
||||
"selectedPath": null,
|
||||
"transformType": "resize",
|
||||
"roiMaskSource": "gt",
|
||||
"roiScale": 2.5,
|
||||
"roiTargetSize": 224,
|
||||
"roiFallback": true,
|
||||
"centerCropSize": 224,
|
||||
"jitterHFlip": true,
|
||||
"jitterVFlip": true,
|
||||
"jitterRotation": 15,
|
||||
"jitterColorEnabled": true,
|
||||
"jitterColor": "0.1,0.1,0.1,0.05",
|
||||
"resizeSize": 256,
|
||||
"x": 155.29958733477324,
|
||||
"y": 622.2110067583806
|
||||
},
|
||||
{
|
||||
"id": "node_1770737490981_bjwe",
|
||||
"type": "transform",
|
||||
"label": "Center Crop",
|
||||
"inputKey": "",
|
||||
"outputKey": "",
|
||||
"inputType": "image",
|
||||
"inputIndex": "1",
|
||||
"outputType": "image",
|
||||
"source": null,
|
||||
"selectedPath": null,
|
||||
"transformType": "center_crop",
|
||||
"roiMaskSource": "gt",
|
||||
"roiScale": 2.5,
|
||||
"roiTargetSize": 224,
|
||||
"roiFallback": true,
|
||||
"centerCropSize": 224,
|
||||
"jitterHFlip": true,
|
||||
"jitterVFlip": true,
|
||||
"jitterRotation": 15,
|
||||
"jitterColorEnabled": true,
|
||||
"jitterColor": "0.1,0.1,0.1,0.05",
|
||||
"resizeSize": 256,
|
||||
"x": 139.05693245845418,
|
||||
"y": 550.529525335558
|
||||
},
|
||||
{
|
||||
"id": "node_1770737490981_ybod",
|
||||
"type": "transform",
|
||||
"label": "Jitter",
|
||||
"inputKey": "",
|
||||
"outputKey": "",
|
||||
"inputType": "image",
|
||||
"inputIndex": "1",
|
||||
"outputType": "image",
|
||||
"source": null,
|
||||
"selectedPath": null,
|
||||
"transformType": "jitter_bundle",
|
||||
"roiMaskSource": "gt",
|
||||
"roiScale": 2.5,
|
||||
"roiTargetSize": 224,
|
||||
"roiFallback": true,
|
||||
"centerCropSize": 224,
|
||||
"jitterHFlip": true,
|
||||
"jitterVFlip": true,
|
||||
"jitterRotation": 15,
|
||||
"jitterColorEnabled": true,
|
||||
"jitterColor": "0.1,0.1,0.1,0.05",
|
||||
"resizeSize": 256,
|
||||
"x": 128.68982314179192,
|
||||
"y": 467.09707170591474
|
||||
},
|
||||
{
|
||||
"id": "node_1770737490981_ogmq",
|
||||
"type": "loader",
|
||||
"label": "Image Loader",
|
||||
"inputKey": "image_1",
|
||||
"outputKey": "image_1",
|
||||
"inputType": "image",
|
||||
"inputIndex": "1",
|
||||
"outputType": "image",
|
||||
"source": null,
|
||||
"selectedPath": null,
|
||||
"x": 160,
|
||||
"y": 370
|
||||
},
|
||||
{
|
||||
"id": "node_1770737490981_i8qg",
|
||||
"type": "loader",
|
||||
"label": "Metadata Loader",
|
||||
"inputKey": "matrix_1",
|
||||
"outputKey": "matrix_1",
|
||||
"inputType": "matrix",
|
||||
"inputIndex": "1",
|
||||
"outputType": "matrix",
|
||||
"source": null,
|
||||
"selectedPath": null,
|
||||
"x": 382.30379326203655,
|
||||
"y": 374.70041266522685
|
||||
},
|
||||
{
|
||||
"id": "node_1770737490981_f1rl",
|
||||
"type": "image_tower",
|
||||
"label": "Tower 1",
|
||||
"inputKey": "",
|
||||
"outputKey": "",
|
||||
"inputType": "image",
|
||||
"inputIndex": "1",
|
||||
"outputType": "image",
|
||||
"source": null,
|
||||
"selectedPath": null,
|
||||
"towerType": "image",
|
||||
"imageBackbone": "efficientnet_b0",
|
||||
"imageFreezeRatio": 0,
|
||||
"imageAugment": true,
|
||||
"imageGeometryDim": 0,
|
||||
"imageUseSe": false,
|
||||
"imageSeReduction": 16,
|
||||
"imageSePreNorm": true,
|
||||
"x": 105.94513543739873,
|
||||
"y": 210.62237129546043
|
||||
},
|
||||
{
|
||||
"id": "node_1770737490982_53hy",
|
||||
"type": "metadata_tower",
|
||||
"label": "Tower 2",
|
||||
"inputKey": "",
|
||||
"outputKey": "",
|
||||
"inputType": "image",
|
||||
"inputIndex": "1",
|
||||
"outputType": "image",
|
||||
"source": null,
|
||||
"selectedPath": null,
|
||||
"towerType": "metadata",
|
||||
"mdHiddenDim": 128,
|
||||
"mdDropout": 0.1,
|
||||
"mdUseSe": false,
|
||||
"mdSeReduction": 16,
|
||||
"mdSePreNorm": true,
|
||||
"mdFreezeRatio": 0,
|
||||
"x": 446.7257671413001,
|
||||
"y": 189.90718376822076
|
||||
},
|
||||
{
|
||||
"id": "node_1770737490982_xmc1",
|
||||
"type": "bridge",
|
||||
"label": "Bridge",
|
||||
"inputKey": "",
|
||||
"outputKey": "",
|
||||
"inputType": "image",
|
||||
"inputIndex": "1",
|
||||
"outputType": "image",
|
||||
"source": null,
|
||||
"selectedPath": null,
|
||||
"bridgeMethod": "fusion",
|
||||
"bridgeFusionDim": 256,
|
||||
"bridgeUseSe": true,
|
||||
"bridgeSeReduction": 16,
|
||||
"bridgeSePreNorm": true,
|
||||
"x": 311.5886354629202,
|
||||
"y": 9.377640626755708
|
||||
},
|
||||
{
|
||||
"id": "node_1770737490982_e9g4",
|
||||
"type": "classifier",
|
||||
"label": "Classifier",
|
||||
"inputKey": "",
|
||||
"outputKey": "",
|
||||
"inputType": "image",
|
||||
"inputIndex": "1",
|
||||
"outputType": "image",
|
||||
"source": null,
|
||||
"selectedPath": null,
|
||||
"x": 308.0632962358768,
|
||||
"y": -139.42405168449085
|
||||
}
|
||||
],
|
||||
"edges": [
|
||||
{
|
||||
"from": "node_1770737490980_3a7n",
|
||||
"to": "node_1770737490980_xnkx"
|
||||
},
|
||||
{
|
||||
"from": "node_1770737490980_xnkx",
|
||||
"to": "node_1770737490981_bjwe"
|
||||
},
|
||||
{
|
||||
"from": "node_1770737490981_bjwe",
|
||||
"to": "node_1770737490981_ybod"
|
||||
},
|
||||
{
|
||||
"from": "node_1770737490981_ybod",
|
||||
"to": "node_1770737490981_ogmq"
|
||||
},
|
||||
{
|
||||
"from": "node_1770737490980_az3a",
|
||||
"to": "node_1770737490981_i8qg"
|
||||
},
|
||||
{
|
||||
"from": "node_1770737490981_ogmq",
|
||||
"to": "node_1770737490981_f1rl"
|
||||
},
|
||||
{
|
||||
"from": "node_1770737490981_i8qg",
|
||||
"to": "node_1770737490982_53hy"
|
||||
},
|
||||
{
|
||||
"from": "node_1770737490981_f1rl",
|
||||
"to": "node_1770737490982_xmc1"
|
||||
},
|
||||
{
|
||||
"from": "node_1770737490982_53hy",
|
||||
"to": "node_1770737490982_xmc1"
|
||||
},
|
||||
{
|
||||
"from": "node_1770737490982_xmc1",
|
||||
"to": "node_1770737490982_e9g4"
|
||||
}
|
||||
],
|
||||
"meta": {
|
||||
"generated_at": "2026-02-10T15:32:18.476Z",
|
||||
"imports": [
|
||||
{
|
||||
"id": "import_1770737490980_00lw",
|
||||
"className": "PapilaData",
|
||||
"params": {
|
||||
"image_dir": "FundusImages",
|
||||
"clinical_dir": "ClinicalData",
|
||||
"label_col": "Diagnosis",
|
||||
"cat_cols": [
|
||||
"Gender",
|
||||
"Phakic/Pseudophakic"
|
||||
]
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
}
|
||||
@@ -1,351 +0,0 @@
|
||||
"""
|
||||
portable_versions/image_loader.py
|
||||
==================================
|
||||
A self-contained image loader with in-memory caching, an optional
|
||||
preprocessing pipeline (e.g. disc cropping), and composable augmentations.
|
||||
|
||||
Returns plain NumPy arrays — works with PyTorch, TensorFlow, JAX, or
|
||||
anything else that can consume an ndarray.
|
||||
|
||||
Dependencies: Pillow, numpy (nothing else)
|
||||
|
||||
Quickstart
|
||||
----------
|
||||
from portable_versions.image_loader import ImageLoader, RandomHorizontalFlip, RandomRotation, ColorJitter
|
||||
|
||||
# 1. Build the loader (once per run)
|
||||
loader = ImageLoader(
|
||||
target_size=(200, 200),
|
||||
normalize=True, # float32 in [0, 1] with ImageNet mean/std
|
||||
cache=True, # each image decoded from disk only once
|
||||
workers=4, # parallel cache warm-up threads
|
||||
preprocessor=my_crop_fn, # optional callable(PIL.Image) -> PIL.Image
|
||||
)
|
||||
|
||||
# 2. Attach augmentations (applied randomly and independently per call)
|
||||
loader.augmentation = [
|
||||
RandomHorizontalFlip(p=0.5),
|
||||
RandomRotation(degrees=15),
|
||||
ColorJitter(brightness=0.2, contrast=0.2, saturation=0.1, hue=0.05),
|
||||
]
|
||||
|
||||
# 3. Warm the cache up front (optional but fast)
|
||||
loader.warm(all_paths)
|
||||
|
||||
# 4. Fetch images by path list — call as many times as you like
|
||||
# Returns ndarray of shape (N, H, W, 3), dtype float32
|
||||
imgs = loader.get_img(train_paths)
|
||||
|
||||
# For TensorFlow:
|
||||
import tensorflow as tf
|
||||
tensor = tf.constant(imgs) # (N, H, W, 3)
|
||||
|
||||
# For PyTorch:
|
||||
import torch
|
||||
tensor = torch.from_numpy(imgs).permute(0, 3, 1, 2) # (N, C, H, W)
|
||||
|
||||
|
||||
Augmentations reference
|
||||
-----------------------
|
||||
All augmentation classes live in this file and depend only on PIL + numpy.
|
||||
|
||||
RandomHorizontalFlip(p=0.5)
|
||||
RandomVerticalFlip(p=0.5)
|
||||
RandomRotation(degrees=15)
|
||||
ColorJitter(brightness=0.2, contrast=0.2, saturation=0.1, hue=0.05)
|
||||
RandomGrayscale(p=0.1)
|
||||
|
||||
You can also pass any callable(PIL.Image.Image) -> PIL.Image.Image as an
|
||||
augmentation step.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import random
|
||||
import threading
|
||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||
from pathlib import Path
|
||||
from typing import Callable, Iterable, List, Optional, Tuple, Union
|
||||
|
||||
import numpy as np
|
||||
from PIL import Image, ImageEnhance, ImageOps
|
||||
|
||||
# ImageNet channel statistics (RGB)
|
||||
_IMAGENET_MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32)
|
||||
_IMAGENET_STD = np.array([0.229, 0.224, 0.225], dtype=np.float32)
|
||||
|
||||
PathLike = Union[str, Path]
|
||||
|
||||
|
||||
def _call_preprocessor(
|
||||
fn: Callable[..., Image.Image],
|
||||
img: Image.Image,
|
||||
path: Path,
|
||||
) -> Image.Image:
|
||||
"""Call preprocessor as fn(img, path) if it accepts two args, else fn(img)."""
|
||||
try:
|
||||
return fn(img, path)
|
||||
except TypeError:
|
||||
return fn(img)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Core loader
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class ImageLoader:
|
||||
"""
|
||||
Preprocessing pipeline + in-memory cache + augmentation, returning NumPy.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
target_size : (height, width)
|
||||
Output spatial dimensions. Applied after ``preprocessor`` (if any).
|
||||
Ignored when a ``preprocessor`` already resizes to the right size.
|
||||
normalize : bool
|
||||
When True, output is float32 with ImageNet mean/std subtraction.
|
||||
When False, output is uint8 in [0, 255].
|
||||
cache : bool
|
||||
Store decoded+preprocessed images in RAM so each file is read from
|
||||
disk at most once. The cache persists across ``get_img`` calls.
|
||||
workers : int
|
||||
Thread count for ``warm()``. 0 or 1 = single-threaded.
|
||||
preprocessor : callable, optional
|
||||
Called as ``preprocessor(img: PIL.Image) -> PIL.Image`` before
|
||||
resizing and caching. Use this for disc cropping, padding, etc.
|
||||
augmentation : list of callables
|
||||
Each element is called as ``fn(img: PIL.Image) -> PIL.Image``.
|
||||
Applied **after** cache retrieval, so augmentations are NOT cached —
|
||||
they are re-sampled independently on every ``get_img`` call.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
target_size: Tuple[int, int] = (224, 224),
|
||||
*,
|
||||
normalize: bool = True,
|
||||
cache: bool = True,
|
||||
workers: int = 4,
|
||||
preprocessor: Optional[Callable[[Image.Image], Image.Image]] = None,
|
||||
) -> None:
|
||||
self.target_size = target_size
|
||||
self.normalize = normalize
|
||||
self.workers = workers
|
||||
self.preprocessor = preprocessor
|
||||
self.augmentation: List[Callable[[Image.Image], Image.Image]] = []
|
||||
|
||||
self._cache: Optional[dict[str, np.ndarray]] = {} if cache else None
|
||||
self._lock = threading.Lock()
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Public
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def warm(self, paths: Iterable[PathLike]) -> None:
|
||||
"""
|
||||
Pre-load all *paths* into the cache in parallel.
|
||||
|
||||
Already-cached paths are skipped, so calling ``warm`` multiple
|
||||
times (e.g. once per fold) is safe and only loads new images.
|
||||
"""
|
||||
if self._cache is None:
|
||||
return
|
||||
|
||||
paths = [str(p) for p in paths]
|
||||
to_warm = [p for p in paths if p not in self._cache]
|
||||
if not to_warm:
|
||||
return
|
||||
|
||||
already = len(paths) - len(to_warm)
|
||||
print(
|
||||
f"[ImageLoader] warming {len(to_warm)} images"
|
||||
+ (f" ({already} already cached)" if already else ""),
|
||||
flush=True,
|
||||
)
|
||||
|
||||
def _load_one(path_str: str) -> None:
|
||||
arr = self._decode(path_str)
|
||||
with self._lock:
|
||||
self._cache.setdefault(path_str, arr)
|
||||
|
||||
if self.workers <= 1:
|
||||
for p in to_warm:
|
||||
_load_one(p)
|
||||
else:
|
||||
with ThreadPoolExecutor(max_workers=self.workers) as ex:
|
||||
futures = {ex.submit(_load_one, p): p for p in to_warm}
|
||||
for fut in as_completed(futures):
|
||||
fut.result()
|
||||
|
||||
def get_img(
|
||||
self,
|
||||
paths: Iterable[PathLike],
|
||||
augment: bool = True,
|
||||
) -> np.ndarray:
|
||||
"""
|
||||
Return images for the given paths as a single NumPy array.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
paths : iterable of path-like
|
||||
File paths to load. If the cache is enabled and a path has
|
||||
been warmed (or loaded before), it is served from RAM.
|
||||
augment : bool
|
||||
Apply ``self.augmentation`` pipeline. Set to False at eval time.
|
||||
|
||||
Returns
|
||||
-------
|
||||
np.ndarray, shape (N, H, W, 3)
|
||||
float32 in [0, 1] (or normalised) if ``self.normalize`` is True,
|
||||
otherwise uint8 in [0, 255].
|
||||
"""
|
||||
imgs = []
|
||||
for p in paths:
|
||||
img = self._get_one(str(p), augment=augment)
|
||||
imgs.append(img)
|
||||
return np.stack(imgs, axis=0)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Internal
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _decode(self, path_str: str) -> np.ndarray:
|
||||
"""Open, preprocess, and resize → uint8 HWC ndarray (for the cache)."""
|
||||
img = Image.open(path_str).convert("RGB")
|
||||
if self.preprocessor is not None:
|
||||
img = _call_preprocessor(self.preprocessor, img, Path(path_str))
|
||||
# Only resize here if preprocessor didn't already produce target_size
|
||||
if img.size != (self.target_size[1], self.target_size[0]):
|
||||
img = img.resize((self.target_size[1], self.target_size[0]), Image.BILINEAR)
|
||||
return np.asarray(img, dtype=np.uint8)
|
||||
|
||||
def _get_one(self, path_str: str, augment: bool) -> np.ndarray:
|
||||
if self._cache is not None:
|
||||
arr = self._cache.get(path_str)
|
||||
if arr is None:
|
||||
arr = self._decode(path_str)
|
||||
with self._lock:
|
||||
self._cache.setdefault(path_str, arr)
|
||||
img = Image.fromarray(arr, mode="RGB")
|
||||
else:
|
||||
img = Image.open(path_str).convert("RGB")
|
||||
if self.preprocessor is not None:
|
||||
img = self.preprocessor(img)
|
||||
if img.size != (self.target_size[1], self.target_size[0]):
|
||||
img = img.resize((self.target_size[1], self.target_size[0]), Image.BILINEAR)
|
||||
|
||||
if augment and self.augmentation:
|
||||
for fn in self.augmentation:
|
||||
img = fn(img)
|
||||
|
||||
arr = np.asarray(img, dtype=np.float32) / 255.0
|
||||
if self.normalize:
|
||||
arr = (arr - _IMAGENET_MEAN) / _IMAGENET_STD
|
||||
else:
|
||||
arr = (arr * 255).clip(0, 255).astype(np.uint8)
|
||||
return arr
|
||||
|
||||
def __len__(self) -> int:
|
||||
"""Number of images currently in the cache."""
|
||||
return len(self._cache) if self._cache is not None else 0
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return (
|
||||
f"ImageLoader(target_size={self.target_size}, "
|
||||
f"normalize={self.normalize}, "
|
||||
f"cached={len(self)}, "
|
||||
f"augmentations={len(self.augmentation)})"
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Augmentation primitives (PIL-only, no torch/tf dependencies)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class RandomHorizontalFlip:
|
||||
"""Flip image left-right with probability *p*."""
|
||||
def __init__(self, p: float = 0.5):
|
||||
self.p = p
|
||||
|
||||
def __call__(self, img: Image.Image) -> Image.Image:
|
||||
return ImageOps.mirror(img) if random.random() < self.p else img
|
||||
|
||||
|
||||
class RandomVerticalFlip:
|
||||
"""Flip image top-bottom with probability *p*."""
|
||||
def __init__(self, p: float = 0.5):
|
||||
self.p = p
|
||||
|
||||
def __call__(self, img: Image.Image) -> Image.Image:
|
||||
return ImageOps.flip(img) if random.random() < self.p else img
|
||||
|
||||
|
||||
class RandomRotation:
|
||||
"""Rotate by a uniformly-sampled angle in [-degrees, +degrees]."""
|
||||
def __init__(self, degrees: float = 15):
|
||||
self.degrees = degrees
|
||||
|
||||
def __call__(self, img: Image.Image) -> Image.Image:
|
||||
angle = random.uniform(-self.degrees, self.degrees)
|
||||
return img.rotate(angle, resample=Image.BILINEAR, expand=False)
|
||||
|
||||
|
||||
class ColorJitter:
|
||||
"""
|
||||
Randomly jitter brightness, contrast, saturation, and hue.
|
||||
|
||||
Each factor is sampled uniformly from [1 - amount, 1 + amount].
|
||||
Hue shift is sampled from [-hue, +hue] (range 0–0.5).
|
||||
Pass 0 for any channel to leave it unchanged.
|
||||
"""
|
||||
def __init__(
|
||||
self,
|
||||
brightness: float = 0.2,
|
||||
contrast: float = 0.2,
|
||||
saturation: float = 0.1,
|
||||
hue: float = 0.05,
|
||||
):
|
||||
self.brightness = brightness
|
||||
self.contrast = contrast
|
||||
self.saturation = saturation
|
||||
self.hue = hue
|
||||
|
||||
def __call__(self, img: Image.Image) -> Image.Image:
|
||||
ops = []
|
||||
if self.brightness:
|
||||
ops.append(("brightness", self.brightness))
|
||||
if self.contrast:
|
||||
ops.append(("contrast", self.contrast))
|
||||
if self.saturation:
|
||||
ops.append(("saturation", self.saturation))
|
||||
if self.hue:
|
||||
ops.append(("hue", self.hue))
|
||||
random.shuffle(ops)
|
||||
|
||||
for kind, amount in ops:
|
||||
factor = random.uniform(1 - amount, 1 + amount)
|
||||
if kind == "brightness":
|
||||
img = ImageEnhance.Brightness(img).enhance(factor)
|
||||
elif kind == "contrast":
|
||||
img = ImageEnhance.Contrast(img).enhance(factor)
|
||||
elif kind == "saturation":
|
||||
img = ImageEnhance.Color(img).enhance(factor)
|
||||
elif kind == "hue":
|
||||
# PIL has no direct hue enhancer — shift via HSV in numpy
|
||||
arr = np.asarray(img.convert("HSV"), dtype=np.int16)
|
||||
shift = int(random.uniform(-self.hue, self.hue) * 255)
|
||||
arr[:, :, 0] = (arr[:, :, 0] + shift) % 256
|
||||
img = Image.fromarray(arr.astype(np.uint8), mode="HSV").convert("RGB")
|
||||
return img
|
||||
|
||||
|
||||
class RandomGrayscale:
|
||||
"""Convert to grayscale (keeping 3 channels) with probability *p*."""
|
||||
def __init__(self, p: float = 0.1):
|
||||
self.p = p
|
||||
|
||||
def __call__(self, img: Image.Image) -> Image.Image:
|
||||
if random.random() < self.p:
|
||||
img = ImageOps.grayscale(img).convert("RGB")
|
||||
return img
|
||||
|
||||
|
||||
@@ -1,270 +0,0 @@
|
||||
"""
|
||||
portable_versions/refuge_mask_adapter.py
|
||||
=========================================
|
||||
Optic-disc cropper for REFUGE (and REFUGE2) fundus images, designed as a
|
||||
drop-in ``preprocessor`` for ``ImageLoader``.
|
||||
|
||||
Given an image and its corresponding segmentation mask, it:
|
||||
1. Extracts the optic disc region from the mask
|
||||
2. Computes a padded bounding box around it
|
||||
3. Crops and resizes the original image
|
||||
|
||||
Dependencies: Pillow, numpy (nothing else)
|
||||
|
||||
Quickstart
|
||||
----------
|
||||
from portable_versions.image_loader import ImageLoader, RandomHorizontalFlip, RandomRotation, ColorJitter
|
||||
from portable_versions.refuge_mask_adapter import RefugeMaskCropper
|
||||
|
||||
cropper = RefugeMaskCropper(
|
||||
mask_dir="REFUGE/Annotations/Training400/Disc_Cup_Masks",
|
||||
scale=1.5, # context around disc (1.0 = tight, 2.0 = lots of context)
|
||||
target_size=(200, 200), # output size — should match ImageLoader target_size
|
||||
mask_suffix=".bmp", # REFUGE1 uses .bmp; REFUGE2 uses .png
|
||||
)
|
||||
|
||||
loader = ImageLoader(
|
||||
target_size=(200, 200),
|
||||
normalize=True,
|
||||
preprocessor=cropper,
|
||||
)
|
||||
loader.augmentation = [
|
||||
RandomHorizontalFlip(),
|
||||
RandomRotation(15),
|
||||
ColorJitter(0.2, 0.2, 0.1, 0.05),
|
||||
]
|
||||
|
||||
imgs = loader.get_img(image_paths, augment=True) # (N, 200, 200, 3)
|
||||
|
||||
REFUGE mask formats
|
||||
-------------------
|
||||
REFUGE1 Grayscale BMP: background=128, disc=255, cup=0
|
||||
REFUGE2 RGB PNG: background detected from image borders, disc/cup by colour
|
||||
|
||||
Both are handled automatically.
|
||||
|
||||
Directory structure assumption
|
||||
------------------------------
|
||||
The cropper looks for the mask with the same stem as the image file, inside
|
||||
``mask_dir``. If your layout differs, pass a custom ``mask_path_fn``:
|
||||
|
||||
cropper = RefugeMaskCropper(
|
||||
mask_path_fn=lambda img_path: img_path.with_suffix(".bmp"),
|
||||
scale=1.5,
|
||||
target_size=(200, 200),
|
||||
)
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from collections import Counter
|
||||
from pathlib import Path
|
||||
from typing import Optional, Tuple
|
||||
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
|
||||
_MASK_DIR_NAMES = {"Disc_Cup_Masks", "Disc_Masks", "Disc_Mask"}
|
||||
_MASK_SUFFIXES = {".bmp", ".png"}
|
||||
|
||||
|
||||
class RefugeMaskCropper:
|
||||
"""
|
||||
Crop a fundus image to the optic disc region using its segmentation mask.
|
||||
|
||||
Pass the REFUGE root directory and the cropper will automatically index
|
||||
all masks underneath it — no need to specify which subdirectory or
|
||||
file extension.
|
||||
|
||||
cropper = RefugeMaskCropper("REFUGE/", scale=1.5)
|
||||
loader = ImageLoader(target_size=(200, 200), preprocessor=cropper)
|
||||
imgs = loader.get_img(test_set) # test_set = any list of image paths
|
||||
|
||||
Parameters
|
||||
----------
|
||||
refuge_root : str or Path
|
||||
Top-level REFUGE directory. All mask files under directories named
|
||||
``Disc_Cup_Masks``, ``Disc_Masks``, or ``Disc_Mask`` are indexed
|
||||
automatically (supports both .bmp and .png).
|
||||
scale : float
|
||||
Padding multiplier applied to the disc radius.
|
||||
1.0 = tight crop, 1.5 = moderate context, 2.5 = lots of context.
|
||||
target_size : (height, width)
|
||||
Output size after cropping. Should match ``ImageLoader.target_size``.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
refuge_root: str | Path,
|
||||
*,
|
||||
scale: float = 1.5,
|
||||
target_size: Tuple[int, int] = (200, 200),
|
||||
) -> None:
|
||||
self.refuge_root = Path(refuge_root)
|
||||
self.scale = scale
|
||||
self.target_size = target_size
|
||||
self._index: dict[str, list[Path]] = {}
|
||||
self._build_index()
|
||||
|
||||
def _build_index(self) -> None:
|
||||
"""Walk refuge_root and index all mask files by stem (stem → [paths])."""
|
||||
for mask_dir in self.refuge_root.rglob("*"):
|
||||
if mask_dir.is_dir() and mask_dir.name in _MASK_DIR_NAMES:
|
||||
for f in mask_dir.rglob("*"):
|
||||
if f.is_file() and f.suffix.lower() in _MASK_SUFFIXES:
|
||||
self._index.setdefault(f.stem, []).append(f)
|
||||
if not self._index:
|
||||
raise FileNotFoundError(
|
||||
f"No mask files found under {self.refuge_root!r}. "
|
||||
f"Expected directories named: {_MASK_DIR_NAMES}"
|
||||
)
|
||||
n_masks = sum(len(v) for v in self._index.values())
|
||||
print(f"[RefugeMaskCropper] indexed {n_masks} masks ({len(self._index)} unique stems)", flush=True)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Callable interface — drop-in preprocessor for ImageLoader
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
img: Image.Image,
|
||||
img_path: Optional[str | Path] = None,
|
||||
) -> Image.Image:
|
||||
stem = Path(img_path).stem if img_path else None
|
||||
mask_path = self._lookup(stem, img_path)
|
||||
disc_mask = _load_disc_mask(mask_path, img.size)
|
||||
box = _mask_to_crop_box(disc_mask, scale=self.scale, img_size=img.size)
|
||||
cropped = img.crop(box)
|
||||
return cropped.resize(
|
||||
(self.target_size[1], self.target_size[0]), Image.Resampling.BILINEAR
|
||||
)
|
||||
|
||||
def _lookup(self, stem: Optional[str], img_path: Optional[str | Path] = None) -> Path:
|
||||
if stem is None:
|
||||
raise ValueError("img_path is required to match the mask.")
|
||||
candidates = self._index.get(stem)
|
||||
if not candidates:
|
||||
raise KeyError(
|
||||
f"No mask found for image stem {stem!r}. "
|
||||
f"Available stems (sample): {list(self._index)[:5]}"
|
||||
)
|
||||
if len(candidates) == 1:
|
||||
return candidates[0]
|
||||
# Pick the mask whose directory components best overlap with img_path
|
||||
# (ignores the filename itself to handle extension differences)
|
||||
img_parts = set(Path(img_path).parent.parts) if img_path else set()
|
||||
return max(candidates, key=lambda m: len(set(m.parent.parts) & img_parts))
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return (
|
||||
f"RefugeMaskCropper(refuge_root={str(self.refuge_root)!r}, "
|
||||
f"scale={self.scale}, target_size={self.target_size}, "
|
||||
f"masks_indexed={len(self._index)})"
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Mask parsing
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _load_disc_mask(mask_path: Path, img_size: Tuple[int, int]) -> np.ndarray:
|
||||
"""
|
||||
Return a binary disc mask (uint8, 1=disc) from a REFUGE mask file.
|
||||
|
||||
Handles:
|
||||
- Grayscale BMP (REFUGE1): background≈128, disc=255, cup=0
|
||||
- RGB PNG (REFUGE2): background detected from image borders
|
||||
"""
|
||||
mask_img = Image.open(mask_path)
|
||||
|
||||
if mask_img.mode == "L" or mask_img.mode == "P":
|
||||
arr = np.asarray(mask_img.convert("L"), dtype=np.uint8)
|
||||
bg = _border_mode(arr)
|
||||
disc_mask = (arr != bg).astype(np.uint8)
|
||||
else:
|
||||
arr = np.asarray(mask_img.convert("RGB"), dtype=np.uint8)
|
||||
bg = _border_mode_rgb(arr)
|
||||
# disc = any non-background pixel
|
||||
bg_mask = np.all(arr == bg, axis=2)
|
||||
disc_mask = (~bg_mask).astype(np.uint8)
|
||||
|
||||
# Ensure mask matches image spatial size
|
||||
mh, mw = disc_mask.shape
|
||||
iw, ih = img_size
|
||||
if (mw, mh) != (iw, ih):
|
||||
disc_img = Image.fromarray(disc_mask * 255).resize((iw, ih), Image.NEAREST)
|
||||
disc_mask = (np.asarray(disc_img) > 0).astype(np.uint8)
|
||||
|
||||
return disc_mask
|
||||
|
||||
|
||||
def _border_mode(arr: np.ndarray, border: int = 5) -> int:
|
||||
"""Most common pixel value along the image border (grayscale)."""
|
||||
h, w = arr.shape
|
||||
border_pixels = np.concatenate([
|
||||
arr[:border, :].ravel(),
|
||||
arr[-border:, :].ravel(),
|
||||
arr[:, :border].ravel(),
|
||||
arr[:, -border:].ravel(),
|
||||
])
|
||||
return int(Counter(border_pixels.tolist()).most_common(1)[0][0])
|
||||
|
||||
|
||||
def _border_mode_rgb(arr: np.ndarray, border: int = 5) -> np.ndarray:
|
||||
"""Most common RGB colour along the image border."""
|
||||
h, w, _ = arr.shape
|
||||
border_pixels = np.concatenate([
|
||||
arr[:border, :].reshape(-1, 3),
|
||||
arr[-border:, :].reshape(-1, 3),
|
||||
arr[:, :border].reshape(-1, 3),
|
||||
arr[:, -border:].reshape(-1, 3),
|
||||
], axis=0)
|
||||
tuples = [tuple(row) for row in border_pixels.tolist()]
|
||||
most_common = Counter(tuples).most_common(1)[0][0]
|
||||
return np.array(most_common, dtype=np.uint8)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Bounding box from mask
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _mask_to_crop_box(
|
||||
disc_mask: np.ndarray,
|
||||
scale: float,
|
||||
img_size: Tuple[int, int],
|
||||
) -> Tuple[int, int, int, int]:
|
||||
"""
|
||||
Compute a square crop box centred on the disc with padding = scale * radius.
|
||||
|
||||
Returns (left, upper, right, lower) — ready for PIL Image.crop().
|
||||
Falls back to the full image if no disc pixels are found.
|
||||
"""
|
||||
coords = np.argwhere(disc_mask > 0) # (N, 2) in (row, col) order
|
||||
if coords.size == 0:
|
||||
w, h = img_size
|
||||
return (0, 0, w, h)
|
||||
|
||||
ys, xs = coords[:, 0], coords[:, 1]
|
||||
centre_x = float(xs.mean())
|
||||
centre_y = float(ys.mean())
|
||||
radius = max(float(xs.max() - xs.min()), float(ys.max() - ys.min())) / 2.0
|
||||
crop_radius = radius * scale
|
||||
|
||||
iw, ih = img_size
|
||||
left = int(max(0, centre_x - crop_radius))
|
||||
upper = int(max(0, centre_y - crop_radius))
|
||||
right = int(min(iw, centre_x + crop_radius))
|
||||
lower = int(min(ih, centre_y + crop_radius))
|
||||
|
||||
# Make square by expanding the shorter side
|
||||
cw, ch = right - left, lower - upper
|
||||
if cw < ch:
|
||||
diff = ch - cw
|
||||
left = max(0, left - diff // 2)
|
||||
right = min(iw, right + diff // 2)
|
||||
elif ch < cw:
|
||||
diff = cw - ch
|
||||
upper = max(0, upper - diff // 2)
|
||||
lower = min(ih, lower + diff // 2)
|
||||
|
||||
return (left, upper, right, lower)
|
||||
Binary file not shown.
Binary file not shown.
|
Before Width: | Height: | Size: 73 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 109 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 72 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 111 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 2.8 MiB |
Binary file not shown.
|
Before Width: | Height: | Size: 105 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 86 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 145 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 91 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 148 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 58 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 63 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 80 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 85 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 58 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 76 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 78 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 1003 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 925 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 95 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 81 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 130 KiB |
@@ -1,61 +0,0 @@
|
||||
{
|
||||
"figure": "/home/rpotter/hypertower/results/poster/papila_binary_roc_comparison.png",
|
||||
"curves": [
|
||||
{
|
||||
"label": "Clinical Data Only",
|
||||
"source_kind": "legacy_npy",
|
||||
"mode_dir": "analysis_data/pipeline_mdonly_500ep/binary/single",
|
||||
"run": null,
|
||||
"tower_path": null,
|
||||
"score_col": null,
|
||||
"probs_stem": "probs_classic",
|
||||
"auc_mean": 0.7311688311688311,
|
||||
"auc_std": 0.08461127279353979,
|
||||
"auc_pooled": 0.6812987012987013,
|
||||
"fold_count": 5,
|
||||
"n_total": 400
|
||||
},
|
||||
{
|
||||
"label": "Image Only",
|
||||
"source_kind": "v3_csv",
|
||||
"mode_dir": null,
|
||||
"run": "phase2/imageonly_resnet50_proper",
|
||||
"tower_path": "binary/single",
|
||||
"score_col": "prob_img_c1",
|
||||
"probs_stem": null,
|
||||
"auc_mean": 0.819264705882353,
|
||||
"auc_std": 0.0740851608018641,
|
||||
"auc_pooled": 0.8099584558823529,
|
||||
"fold_count": 50,
|
||||
"n_total": 4200
|
||||
},
|
||||
{
|
||||
"label": "Single Fusion",
|
||||
"source_kind": "v3_csv",
|
||||
"mode_dir": null,
|
||||
"run": "phase5/single_fused",
|
||||
"tower_path": "binary/single",
|
||||
"score_col": "prob_fused_c1",
|
||||
"probs_stem": null,
|
||||
"auc_mean": 0.8506801470588234,
|
||||
"auc_std": 0.07392522372615995,
|
||||
"auc_pooled": 0.8402091911764706,
|
||||
"fold_count": 50,
|
||||
"n_total": 4200
|
||||
},
|
||||
{
|
||||
"label": "Ensemble Fusion",
|
||||
"source_kind": "v3_csv",
|
||||
"mode_dir": null,
|
||||
"run": "phase5/logit_mlp_head",
|
||||
"tower_path": "binary/ensemble",
|
||||
"score_col": "prob_fused_c1",
|
||||
"probs_stem": null,
|
||||
"auc_mean": 0.8988235294117648,
|
||||
"auc_std": 0.06435919826554772,
|
||||
"auc_pooled": 0.8855705882352942,
|
||||
"fold_count": 50,
|
||||
"n_total": 2100
|
||||
}
|
||||
]
|
||||
}
|
||||
Binary file not shown.
|
Before Width: | Height: | Size: 219 KiB |
@@ -1,757 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Basic analytics helpers for PAPILA clinical data."""
|
||||
import re
|
||||
from pathlib import Path
|
||||
from typing import Iterable, List, Tuple
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import matplotlib.pyplot as plt
|
||||
from sklearn.base import clone
|
||||
from sklearn.metrics import roc_curve, auc, roc_auc_score, accuracy_score
|
||||
from sklearn.ensemble import RandomForestClassifier
|
||||
from sklearn.model_selection import StratifiedKFold
|
||||
from sklearn.linear_model import LogisticRegression
|
||||
from sklearn.neighbors import KNeighborsClassifier
|
||||
from sklearn.pipeline import Pipeline
|
||||
from sklearn.preprocessing import StandardScaler
|
||||
from sklearn.svm import SVC
|
||||
|
||||
from classes import build_papila_clinical
|
||||
|
||||
|
||||
class basic_analytics:
|
||||
def __init__(
|
||||
self,
|
||||
image_dir: str = "Papila/FundusImages",
|
||||
clinical_dir: str = "Papila/ClinicalData",
|
||||
label_col: str = "Diagnosis",
|
||||
cat_cols: Iterable[str] | None = None,
|
||||
exclude_cols: Iterable[str] | None = None,
|
||||
positive_label: int = 1,
|
||||
negative_label: int = 0,
|
||||
drop_labels: Iterable[int] = (2,),
|
||||
output_dir: Path | str = Path("analysis_data/basic_analysis"),
|
||||
debug: bool = False,
|
||||
) -> None:
|
||||
self.image_dir = image_dir
|
||||
self.clinical_dir = clinical_dir
|
||||
self.label_col = label_col
|
||||
self.cat_cols = (
|
||||
list(cat_cols)
|
||||
if cat_cols is not None
|
||||
else ["Gender", "Phakic/Pseudophakic"]
|
||||
)
|
||||
base_excludes = {"Pneumatic", "Perkins"}
|
||||
self.exclude_cols = base_excludes | set(exclude_cols or [])
|
||||
self.positive_label = positive_label
|
||||
self.negative_label = negative_label
|
||||
self.drop_labels = list(drop_labels or [])
|
||||
self.output_dir = Path(output_dir)
|
||||
self.debug = debug
|
||||
|
||||
@staticmethod
|
||||
def _sanitize(name: str) -> str:
|
||||
safe = re.sub(r"[^A-Za-z0-9._-]+", "_", str(name)).strip("_")
|
||||
return safe or "var"
|
||||
|
||||
def _build_clinical(self):
|
||||
return build_papila_clinical(
|
||||
image_dir=self.image_dir,
|
||||
clinical_dir=self.clinical_dir,
|
||||
label_col=self.label_col,
|
||||
cat_cols=self.cat_cols,
|
||||
)
|
||||
|
||||
def _select_binary_labels(
|
||||
self,
|
||||
labels: pd.Series,
|
||||
) -> Tuple[np.ndarray, np.ndarray]:
|
||||
labels_num = pd.to_numeric(labels, errors="coerce")
|
||||
use_num = labels_num.notna().any()
|
||||
lab = labels_num if use_num else labels.astype(str)
|
||||
|
||||
drop_set = set(self.drop_labels or [])
|
||||
keep = lab.isin([self.positive_label, self.negative_label])
|
||||
if drop_set:
|
||||
keep &= ~lab.isin(drop_set)
|
||||
|
||||
y = (lab == self.positive_label).astype(int)
|
||||
return y.values, keep.values
|
||||
|
||||
def _base_exclude(self) -> set:
|
||||
base_exclude = (
|
||||
{self.label_col, "Patient ID"} | self.exclude_cols | set(self.cat_cols)
|
||||
)
|
||||
if "eyeID" not in self.cat_cols:
|
||||
base_exclude.add("eyeID")
|
||||
return base_exclude
|
||||
|
||||
def _numeric_columns(self, df: pd.DataFrame) -> List[str]:
|
||||
base_exclude = self._base_exclude()
|
||||
candidate_cols = [c for c in df.columns if c not in base_exclude]
|
||||
numeric_cols: List[str] = []
|
||||
for col in candidate_cols:
|
||||
s = pd.to_numeric(df[col], errors="coerce")
|
||||
if s.notna().any():
|
||||
numeric_cols.append(col)
|
||||
return numeric_cols
|
||||
|
||||
@staticmethod
|
||||
def _compute_roc(
|
||||
y: np.ndarray, scores: np.ndarray
|
||||
) -> Tuple[np.ndarray, np.ndarray, np.ndarray, float]:
|
||||
fpr, tpr, thresholds = roc_curve(y, scores, pos_label=1)
|
||||
auc_val = float(auc(fpr, tpr))
|
||||
return fpr, tpr, thresholds, auc_val
|
||||
|
||||
@staticmethod
|
||||
def _best_threshold(
|
||||
fpr: np.ndarray, tpr: np.ndarray, thresholds: np.ndarray
|
||||
) -> Tuple[float, float, float]:
|
||||
youden = tpr - fpr
|
||||
idx = int(np.nanargmax(youden))
|
||||
return float(thresholds[idx]), float(tpr[idx]), float(fpr[idx])
|
||||
|
||||
@staticmethod
|
||||
def _plot_overlay(curves, title: str, out_path: Path) -> None:
|
||||
fig, ax = plt.subplots(figsize=(7, 5.5))
|
||||
cmap = plt.get_cmap("tab20")
|
||||
for i, (name, fpr, tpr, auc_val) in enumerate(curves):
|
||||
color = cmap(i % cmap.N)
|
||||
ax.plot(fpr, tpr, lw=1.6, color=color, label=f"{name} (AUC={auc_val:.3f})")
|
||||
ax.plot([0, 1], [0, 1], "k--", lw=1)
|
||||
ax.set_xlabel("False Positive Rate")
|
||||
ax.set_ylabel("True Positive Rate")
|
||||
ax.set_title(title)
|
||||
ax.legend(loc="upper left", fontsize="small")
|
||||
ax.grid(True, alpha=0.3, linestyle="--")
|
||||
fig.tight_layout()
|
||||
fig.savefig(out_path, dpi=170)
|
||||
plt.close(fig)
|
||||
|
||||
@staticmethod
|
||||
def _plot_per_feature(
|
||||
fpr: np.ndarray, tpr: np.ndarray, auc_val: float, title: str, out_path: Path
|
||||
) -> None:
|
||||
fig, ax = plt.subplots(figsize=(5.5, 4.5))
|
||||
ax.plot(fpr, tpr, lw=1.8, label=f"AUC={auc_val:.3f}")
|
||||
ax.plot([0, 1], [0, 1], "k--", lw=1)
|
||||
ax.set_xlabel("False Positive Rate")
|
||||
ax.set_ylabel("True Positive Rate")
|
||||
ax.set_title(title)
|
||||
ax.legend(loc="lower right")
|
||||
ax.grid(True, alpha=0.3, linestyle="--")
|
||||
fig.tight_layout()
|
||||
fig.savefig(out_path, dpi=170)
|
||||
plt.close(fig)
|
||||
|
||||
@staticmethod
|
||||
def _plot_roc_line(
|
||||
fpr: np.ndarray, tpr: np.ndarray, auc_val: float, title: str, out_path: Path
|
||||
) -> None:
|
||||
fig, ax = plt.subplots(figsize=(5.5, 4.5))
|
||||
ax.plot(fpr, tpr, lw=1.8, label=f"AUC={auc_val:.3f}")
|
||||
ax.plot([0, 1], [0, 1], "k--", lw=1)
|
||||
ax.set_xlabel("False Positive Rate")
|
||||
ax.set_ylabel("True Positive Rate")
|
||||
ax.set_title(title)
|
||||
ax.legend(loc="lower right")
|
||||
ax.grid(True, alpha=0.3, linestyle="--")
|
||||
fig.tight_layout()
|
||||
fig.savefig(out_path, dpi=170)
|
||||
plt.close(fig)
|
||||
|
||||
def _oof_scores(
|
||||
self,
|
||||
model,
|
||||
X: np.ndarray,
|
||||
y: np.ndarray,
|
||||
n_splits: int,
|
||||
random_state: int,
|
||||
) -> Tuple[np.ndarray, np.ndarray]:
|
||||
skf = StratifiedKFold(
|
||||
n_splits=n_splits, shuffle=True, random_state=random_state
|
||||
)
|
||||
scores = np.zeros(len(y), dtype=float)
|
||||
for train_idx, test_idx in skf.split(X, y):
|
||||
X_train, X_test = X[train_idx], X[test_idx]
|
||||
y_train = y[train_idx]
|
||||
if np.unique(y_train).size < 2:
|
||||
continue
|
||||
fitted = clone(model)
|
||||
fitted.fit(X_train, y_train)
|
||||
if hasattr(fitted, "predict_proba"):
|
||||
fold_scores = fitted.predict_proba(X_test)[:, 1]
|
||||
elif hasattr(fitted, "decision_function"):
|
||||
fold_scores = fitted.decision_function(X_test)
|
||||
else:
|
||||
fold_scores = fitted.predict(X_test)
|
||||
scores[test_idx] = fold_scores
|
||||
return y.astype(int), scores
|
||||
|
||||
def _cv_roc_curves(
|
||||
self,
|
||||
model,
|
||||
X: np.ndarray,
|
||||
y: np.ndarray,
|
||||
n_splits: int,
|
||||
random_state: int,
|
||||
) -> List[Tuple[np.ndarray, np.ndarray, float]]:
|
||||
skf = StratifiedKFold(
|
||||
n_splits=n_splits, shuffle=True, random_state=random_state
|
||||
)
|
||||
curves = []
|
||||
for train_idx, test_idx in skf.split(X, y):
|
||||
X_train, X_test = X[train_idx], X[test_idx]
|
||||
y_train, y_test = y[train_idx], y[test_idx]
|
||||
if np.unique(y_train).size < 2 or np.unique(y_test).size < 2:
|
||||
continue
|
||||
fitted = clone(model)
|
||||
fitted.fit(X_train, y_train)
|
||||
if hasattr(fitted, "predict_proba"):
|
||||
scores = fitted.predict_proba(X_test)[:, 1]
|
||||
elif hasattr(fitted, "decision_function"):
|
||||
scores = fitted.decision_function(X_test)
|
||||
else:
|
||||
scores = fitted.predict(X_test)
|
||||
fpr, tpr, _ = roc_curve(y_test, scores, pos_label=1)
|
||||
auc_val = float(auc(fpr, tpr))
|
||||
curves.append((fpr, tpr, auc_val))
|
||||
return curves
|
||||
|
||||
@staticmethod
|
||||
def _plot_mean_roc(
|
||||
curves: List[Tuple[np.ndarray, np.ndarray, float]],
|
||||
title: str,
|
||||
out_path: Path,
|
||||
) -> None:
|
||||
if not curves:
|
||||
return
|
||||
mean_fpr = np.linspace(0.0, 1.0, 200)
|
||||
tprs = []
|
||||
aucs = []
|
||||
for fpr, tpr, auc_val in curves:
|
||||
tpr_interp = np.interp(mean_fpr, fpr, tpr)
|
||||
tpr_interp[0] = 0.0
|
||||
tprs.append(tpr_interp)
|
||||
aucs.append(auc_val)
|
||||
mean_tpr = np.mean(tprs, axis=0)
|
||||
mean_tpr[-1] = 1.0
|
||||
std_tpr = np.std(tprs, axis=0)
|
||||
mean_auc = float(np.mean(aucs))
|
||||
std_auc = float(np.std(aucs, ddof=0))
|
||||
|
||||
fig, ax = plt.subplots(figsize=(5.8, 4.6))
|
||||
ax.plot(mean_fpr, mean_tpr, lw=2, label=f"AUC={mean_auc:.3f}±{std_auc:.3f}")
|
||||
ax.fill_between(
|
||||
mean_fpr,
|
||||
np.maximum(mean_tpr - std_tpr, 0),
|
||||
np.minimum(mean_tpr + std_tpr, 1),
|
||||
color="grey",
|
||||
alpha=0.2,
|
||||
)
|
||||
ax.plot([0, 1], [0, 1], "k--", lw=1)
|
||||
ax.set_xlabel("False Positive Rate")
|
||||
ax.set_ylabel("True Positive Rate")
|
||||
ax.set_title(title)
|
||||
ax.legend(loc="lower right")
|
||||
ax.grid(True, alpha=0.3, linestyle="--")
|
||||
fig.tight_layout()
|
||||
fig.savefig(out_path, dpi=170)
|
||||
plt.close(fig)
|
||||
|
||||
def _iter_categorical(self, df: pd.DataFrame, cols: List[str]):
|
||||
for col in cols:
|
||||
if col not in df.columns:
|
||||
continue
|
||||
s = df[col]
|
||||
vals = s.dropna().unique().tolist()
|
||||
try:
|
||||
vals = sorted(vals)
|
||||
except Exception:
|
||||
pass
|
||||
for v in vals:
|
||||
name = f"{col}=={v}"
|
||||
ind = (s == v).astype(int)
|
||||
yield name, ind
|
||||
|
||||
def _feature_matrix(
|
||||
self, df: pd.DataFrame, include_categorical: bool
|
||||
) -> Tuple[np.ndarray, np.ndarray, List[str]]:
|
||||
numeric_cols = self._numeric_columns(df)
|
||||
X_num = df[numeric_cols].apply(pd.to_numeric, errors="coerce")
|
||||
for col in numeric_cols:
|
||||
med = pd.to_numeric(X_num[col], errors="coerce").median()
|
||||
X_num[col] = pd.to_numeric(X_num[col], errors="coerce").fillna(med)
|
||||
|
||||
parts = [X_num]
|
||||
feat_names = list(X_num.columns)
|
||||
|
||||
if include_categorical and self.cat_cols:
|
||||
cat_cols = [c for c in self.cat_cols if c in df.columns]
|
||||
if cat_cols:
|
||||
df_cats = pd.get_dummies(
|
||||
df[cat_cols].astype("category"), drop_first=False, prefix=cat_cols
|
||||
)
|
||||
parts.append(df_cats)
|
||||
feat_names.extend(list(df_cats.columns))
|
||||
|
||||
X = pd.concat(parts, axis=1).values.astype(np.float32)
|
||||
labels = df[self.label_col]
|
||||
y_all, keep_mask = self._select_binary_labels(labels)
|
||||
y = y_all[keep_mask]
|
||||
X = X[keep_mask]
|
||||
return X, y.astype(int), feat_names
|
||||
|
||||
def univariate_roc(
|
||||
self, merge: bool = False, include_categorical: bool = False
|
||||
) -> pd.DataFrame:
|
||||
clinical = self._build_clinical()
|
||||
df = clinical.df.copy()
|
||||
labels = df[self.label_col]
|
||||
y_all, keep_mask = self._select_binary_labels(labels)
|
||||
|
||||
if self.debug:
|
||||
for col in ("IOP_raw", "IOP_corr"):
|
||||
if col not in df.columns:
|
||||
print(f"[debug] {col} missing from df")
|
||||
continue
|
||||
s = pd.to_numeric(df[col], errors="coerce")
|
||||
print(
|
||||
f"[debug] {col}: non-null={int(s.notna().sum())}, unique={int(s.nunique(dropna=True))}"
|
||||
)
|
||||
|
||||
plot_dir = self.output_dir / "papila_univariate_roc" / "plots"
|
||||
plot_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
rows = []
|
||||
curves = []
|
||||
|
||||
numeric_cols = self._numeric_columns(df)
|
||||
for col in numeric_cols:
|
||||
series = pd.to_numeric(df[col], errors="coerce")
|
||||
mask = keep_mask & series.notna().values
|
||||
y = y_all[mask]
|
||||
scores = series.values[mask].astype(float)
|
||||
if y.size < 2 or np.unique(y).size < 2:
|
||||
continue
|
||||
if np.nanmin(scores) == np.nanmax(scores):
|
||||
continue
|
||||
fpr, tpr, thresholds, auc_val = self._compute_roc(y, scores)
|
||||
thr, best_tpr, best_fpr = self._best_threshold(fpr, tpr, thresholds)
|
||||
direction = "high" if auc_val >= 0.5 else "low"
|
||||
title = f"{col} (n={y.size}, direction={direction})"
|
||||
if merge:
|
||||
out_path = plot_dir / f"roc_{self._sanitize(col)}.png"
|
||||
self._plot_per_feature(fpr, tpr, auc_val, title, out_path)
|
||||
curves.append((col, fpr, tpr, auc_val))
|
||||
rows.append(
|
||||
{
|
||||
"feature": col,
|
||||
"kind": "numeric",
|
||||
"n": int(y.size),
|
||||
"auc": auc_val,
|
||||
"direction": direction,
|
||||
"best_threshold": thr,
|
||||
"best_tpr": best_tpr,
|
||||
"best_fpr": best_fpr,
|
||||
"best_specificity": 1.0 - best_fpr,
|
||||
}
|
||||
)
|
||||
|
||||
if include_categorical:
|
||||
cat_cols_use = [c for c in self.cat_cols if c not in self.exclude_cols]
|
||||
for name, ind in self._iter_categorical(df, cat_cols_use):
|
||||
mask = keep_mask & ind.notna().values
|
||||
y = y_all[mask]
|
||||
scores = ind.values[mask].astype(float)
|
||||
if y.size < 2 or np.unique(y).size < 2:
|
||||
continue
|
||||
if np.nanmin(scores) == np.nanmax(scores):
|
||||
continue
|
||||
fpr, tpr, thresholds, auc_val = self._compute_roc(y, scores)
|
||||
thr, best_tpr, best_fpr = self._best_threshold(fpr, tpr, thresholds)
|
||||
direction = "high" if auc_val >= 0.5 else "low"
|
||||
title = f"{name} (n={y.size}, direction={direction})"
|
||||
if merge:
|
||||
out_path = plot_dir / f"roc_{self._sanitize(name)}.png"
|
||||
self._plot_per_feature(fpr, tpr, auc_val, title, out_path)
|
||||
curves.append((name, fpr, tpr, auc_val))
|
||||
rows.append(
|
||||
{
|
||||
"feature": name,
|
||||
"kind": "categorical",
|
||||
"n": int(y.size),
|
||||
"auc": auc_val,
|
||||
"direction": direction,
|
||||
"best_threshold": thr,
|
||||
"best_tpr": best_tpr,
|
||||
"best_fpr": best_fpr,
|
||||
"best_specificity": 1.0 - best_fpr,
|
||||
}
|
||||
)
|
||||
|
||||
if not rows:
|
||||
raise SystemExit(
|
||||
"No valid features produced ROC curves. Check labels and feature columns."
|
||||
)
|
||||
|
||||
overlay_path = plot_dir / "roc_overlay.png"
|
||||
if not merge:
|
||||
self._plot_overlay(curves, "Univariate ROC curves", overlay_path)
|
||||
|
||||
out_df = pd.DataFrame(rows).sort_values(by="auc", ascending=False)
|
||||
out_csv = self.output_dir / "papila_univariate_roc" / "summary.csv"
|
||||
out_csv.parent.mkdir(parents=True, exist_ok=True)
|
||||
out_df.to_csv(out_csv, index=False)
|
||||
return out_df
|
||||
|
||||
def random_forest(
|
||||
self,
|
||||
include_categorical: bool = True,
|
||||
n_estimators: int = 500,
|
||||
max_depth: int | None = None,
|
||||
min_samples_leaf: int = 1,
|
||||
max_features: str | None = "sqrt",
|
||||
class_weight: str | None = "balanced",
|
||||
max_samples: float | None = None,
|
||||
random_state: int = 42,
|
||||
top_n: int = 25,
|
||||
n_splits: int = 5,
|
||||
drop_missing: bool = False,
|
||||
nerf: bool = False,
|
||||
drop_age: bool = False,
|
||||
) -> pd.DataFrame:
|
||||
clinical = self._build_clinical()
|
||||
df = clinical.df.copy()
|
||||
original_exclude = set(self.exclude_cols)
|
||||
if nerf:
|
||||
self.exclude_cols = set(self.exclude_cols)
|
||||
drop_missing = True
|
||||
if drop_age:
|
||||
self.exclude_cols = set(self.exclude_cols) | {"Age"}
|
||||
if drop_missing:
|
||||
numeric_cols = self._numeric_columns(df)
|
||||
df = df.dropna(subset=numeric_cols)
|
||||
X, y, feat_names = self._feature_matrix(
|
||||
df, include_categorical=include_categorical
|
||||
)
|
||||
|
||||
if X.size == 0 or np.unique(y).size < 2:
|
||||
raise SystemExit(
|
||||
"Not enough data after filtering labels for Random Forest."
|
||||
)
|
||||
|
||||
if nerf:
|
||||
n_estimators = 200
|
||||
max_depth = 5
|
||||
min_samples_leaf = 5
|
||||
max_features = "sqrt"
|
||||
class_weight = None
|
||||
max_samples = 0.7
|
||||
self.exclude_cols = original_exclude
|
||||
|
||||
clf = RandomForestClassifier(
|
||||
n_estimators=n_estimators,
|
||||
max_depth=max_depth,
|
||||
min_samples_leaf=min_samples_leaf,
|
||||
max_features=max_features,
|
||||
class_weight=class_weight,
|
||||
max_samples=max_samples,
|
||||
random_state=random_state,
|
||||
n_jobs=-1,
|
||||
)
|
||||
clf.fit(X, y)
|
||||
importances = clf.feature_importances_.astype(float)
|
||||
|
||||
rows = []
|
||||
for name, val in zip(feat_names, importances):
|
||||
rows.append({"feature": name, "importance": float(val)})
|
||||
|
||||
out_df = pd.DataFrame(rows).sort_values(by="importance", ascending=False)
|
||||
out_dir = self.output_dir / (
|
||||
"papila_random_forest_nerfed" if nerf else "papila_random_forest"
|
||||
)
|
||||
out_dir.mkdir(parents=True, exist_ok=True)
|
||||
out_df.to_csv(out_dir / "feature_importance.csv", index=False)
|
||||
|
||||
top_df = out_df.head(top_n)
|
||||
fig, ax = plt.subplots(figsize=(7, 6))
|
||||
ax.barh(top_df["feature"], top_df["importance"], color="steelblue")
|
||||
ax.invert_yaxis()
|
||||
ax.set_xlabel("Importance (Gini)")
|
||||
ax.set_title(f"Random Forest Feature Importance")
|
||||
fig.tight_layout()
|
||||
fig.savefig(out_dir / "feature_importance_top.png", dpi=170)
|
||||
plt.close(fig)
|
||||
if self.debug:
|
||||
age_rows = out_df[out_df["feature"] == "Age"]
|
||||
if not age_rows.empty:
|
||||
age_imp = float(age_rows["importance"].iloc[0])
|
||||
print(f"[debug] RF importance Age = {age_imp:.4f}")
|
||||
|
||||
y_oof, scores_oof = self._oof_scores(clf, X, y, n_splits, random_state)
|
||||
fpr, tpr, _, auc_val = self._compute_roc(y_oof, scores_oof)
|
||||
curves = self._cv_roc_curves(clf, X, y, n_splits, random_state)
|
||||
self._plot_mean_roc(
|
||||
curves,
|
||||
"Random Forest ROC (mean ± SD)",
|
||||
out_dir / "roc_mean.png",
|
||||
)
|
||||
|
||||
return out_df
|
||||
|
||||
def _cv_binary_metrics(
|
||||
self,
|
||||
model,
|
||||
X: np.ndarray,
|
||||
y: np.ndarray,
|
||||
n_splits: int,
|
||||
random_state: int,
|
||||
) -> pd.DataFrame:
|
||||
skf = StratifiedKFold(
|
||||
n_splits=n_splits, shuffle=True, random_state=random_state
|
||||
)
|
||||
rows = []
|
||||
for fold, (train_idx, test_idx) in enumerate(skf.split(X, y), start=1):
|
||||
X_train, X_test = X[train_idx], X[test_idx]
|
||||
y_train, y_test = y[train_idx], y[test_idx]
|
||||
if np.unique(y_train).size < 2 or np.unique(y_test).size < 2:
|
||||
continue
|
||||
model.fit(X_train, y_train)
|
||||
if hasattr(model, "predict_proba"):
|
||||
scores = model.predict_proba(X_test)[:, 1]
|
||||
elif hasattr(model, "decision_function"):
|
||||
scores = model.decision_function(X_test)
|
||||
else:
|
||||
scores = model.predict(X_test)
|
||||
preds = model.predict(X_test)
|
||||
auc_val = float(roc_auc_score(y_test, scores))
|
||||
acc_val = float(accuracy_score(y_test, preds))
|
||||
rows.append(
|
||||
{
|
||||
"fold": int(fold),
|
||||
"n": int(len(y_test)),
|
||||
"auc": auc_val,
|
||||
"acc": acc_val,
|
||||
}
|
||||
)
|
||||
return pd.DataFrame(rows)
|
||||
|
||||
def svm(
|
||||
self,
|
||||
include_categorical: bool = True,
|
||||
kernel: str = "rbf",
|
||||
C: float = 1.0,
|
||||
gamma: str = "scale",
|
||||
n_splits: int = 5,
|
||||
random_state: int = 42,
|
||||
) -> pd.DataFrame:
|
||||
clinical = self._build_clinical()
|
||||
df = clinical.df.copy()
|
||||
X, y, feat_names = self._feature_matrix(
|
||||
df, include_categorical=include_categorical
|
||||
)
|
||||
if X.size == 0 or np.unique(y).size < 2:
|
||||
raise SystemExit("Not enough data after filtering labels for SVM.")
|
||||
|
||||
model = Pipeline(
|
||||
[
|
||||
("scale", StandardScaler()),
|
||||
(
|
||||
"svm",
|
||||
SVC(
|
||||
kernel=kernel,
|
||||
C=C,
|
||||
gamma=gamma,
|
||||
probability=True,
|
||||
class_weight="balanced",
|
||||
random_state=random_state,
|
||||
),
|
||||
),
|
||||
]
|
||||
)
|
||||
fold_df = self._cv_binary_metrics(model, X, y, n_splits, random_state)
|
||||
if fold_df.empty:
|
||||
raise SystemExit("SVM produced no valid folds (check class balance).")
|
||||
|
||||
y_oof, scores_oof = self._oof_scores(model, X, y, n_splits, random_state)
|
||||
fpr, tpr, _, auc_val = self._compute_roc(y_oof, scores_oof)
|
||||
curves = self._cv_roc_curves(model, X, y, n_splits, random_state)
|
||||
|
||||
summary = pd.DataFrame(
|
||||
[
|
||||
{
|
||||
"metric": "auc",
|
||||
"mean": float(fold_df["auc"].mean()),
|
||||
"std": float(fold_df["auc"].std(ddof=0)),
|
||||
"oof_auc": float(auc_val),
|
||||
},
|
||||
{
|
||||
"metric": "acc",
|
||||
"mean": float(fold_df["acc"].mean()),
|
||||
"std": float(fold_df["acc"].std(ddof=0)),
|
||||
},
|
||||
]
|
||||
)
|
||||
out_dir = self.output_dir / "papila_svm"
|
||||
out_dir.mkdir(parents=True, exist_ok=True)
|
||||
fold_df.to_csv(out_dir / "fold_metrics.csv", index=False)
|
||||
summary.to_csv(out_dir / "summary.csv", index=False)
|
||||
self._plot_mean_roc(
|
||||
curves,
|
||||
"SVM ROC (mean ± SD)",
|
||||
out_dir / "roc_mean.png",
|
||||
)
|
||||
return fold_df
|
||||
|
||||
def knn(
|
||||
self,
|
||||
include_categorical: bool = True,
|
||||
n_neighbors: int = 5,
|
||||
weights: str = "distance",
|
||||
n_splits: int = 5,
|
||||
random_state: int = 42,
|
||||
) -> pd.DataFrame:
|
||||
clinical = self._build_clinical()
|
||||
df = clinical.df.copy()
|
||||
X, y, feat_names = self._feature_matrix(
|
||||
df, include_categorical=include_categorical
|
||||
)
|
||||
if X.size == 0 or np.unique(y).size < 2:
|
||||
raise SystemExit("Not enough data after filtering labels for KNN.")
|
||||
|
||||
model = Pipeline(
|
||||
[
|
||||
("scale", StandardScaler()),
|
||||
("knn", KNeighborsClassifier(n_neighbors=n_neighbors, weights=weights)),
|
||||
]
|
||||
)
|
||||
fold_df = self._cv_binary_metrics(model, X, y, n_splits, random_state)
|
||||
if fold_df.empty:
|
||||
raise SystemExit("KNN produced no valid folds (check class balance).")
|
||||
|
||||
y_oof, scores_oof = self._oof_scores(model, X, y, n_splits, random_state)
|
||||
fpr, tpr, _, auc_val = self._compute_roc(y_oof, scores_oof)
|
||||
curves = self._cv_roc_curves(model, X, y, n_splits, random_state)
|
||||
|
||||
summary = pd.DataFrame(
|
||||
[
|
||||
{
|
||||
"metric": "auc",
|
||||
"mean": float(fold_df["auc"].mean()),
|
||||
"std": float(fold_df["auc"].std(ddof=0)),
|
||||
"oof_auc": float(auc_val),
|
||||
},
|
||||
{
|
||||
"metric": "acc",
|
||||
"mean": float(fold_df["acc"].mean()),
|
||||
"std": float(fold_df["acc"].std(ddof=0)),
|
||||
},
|
||||
]
|
||||
)
|
||||
out_dir = self.output_dir / "papila_knn"
|
||||
out_dir.mkdir(parents=True, exist_ok=True)
|
||||
fold_df.to_csv(out_dir / "fold_metrics.csv", index=False)
|
||||
summary.to_csv(out_dir / "summary.csv", index=False)
|
||||
self._plot_mean_roc(
|
||||
curves,
|
||||
"KNN ROC (mean ± SD)",
|
||||
out_dir / "roc_mean.png",
|
||||
)
|
||||
return fold_df
|
||||
|
||||
def logistic_regression(
|
||||
self,
|
||||
include_categorical: bool = True,
|
||||
C: float = 1.0,
|
||||
max_iter: int = 1000,
|
||||
n_splits: int = 5,
|
||||
random_state: int = 42,
|
||||
drop_age: bool = False,
|
||||
nerf: bool = False,
|
||||
) -> pd.DataFrame:
|
||||
clinical = self._build_clinical()
|
||||
df = clinical.df.copy()
|
||||
original_exclude = set(self.exclude_cols)
|
||||
if drop_age:
|
||||
self.exclude_cols = set(self.exclude_cols) | {"Age"}
|
||||
X, y, feat_names = self._feature_matrix(
|
||||
df, include_categorical=include_categorical
|
||||
)
|
||||
self.exclude_cols = original_exclude
|
||||
if X.size == 0 or np.unique(y).size < 2:
|
||||
raise SystemExit(
|
||||
"Not enough data after filtering labels for Logistic Regression."
|
||||
)
|
||||
|
||||
class_weight = "balanced"
|
||||
penalty = "l2"
|
||||
solver = "lbfgs"
|
||||
if nerf:
|
||||
C = 0.05
|
||||
class_weight = None
|
||||
penalty = "l1"
|
||||
solver = "liblinear"
|
||||
|
||||
model = Pipeline(
|
||||
[
|
||||
("scale", StandardScaler()),
|
||||
(
|
||||
"logreg",
|
||||
LogisticRegression(
|
||||
C=C,
|
||||
max_iter=max_iter,
|
||||
class_weight=class_weight,
|
||||
penalty=penalty,
|
||||
solver=solver,
|
||||
),
|
||||
),
|
||||
]
|
||||
)
|
||||
fold_df = self._cv_binary_metrics(model, X, y, n_splits, random_state)
|
||||
if fold_df.empty:
|
||||
raise SystemExit(
|
||||
"Logistic Regression produced no valid folds (check class balance)."
|
||||
)
|
||||
|
||||
y_oof, scores_oof = self._oof_scores(model, X, y, n_splits, random_state)
|
||||
fpr, tpr, _, auc_val = self._compute_roc(y_oof, scores_oof)
|
||||
curves = self._cv_roc_curves(model, X, y, n_splits, random_state)
|
||||
|
||||
summary = pd.DataFrame(
|
||||
[
|
||||
{
|
||||
"metric": "auc",
|
||||
"mean": float(fold_df["auc"].mean()),
|
||||
"std": float(fold_df["auc"].std(ddof=0)),
|
||||
"oof_auc": float(auc_val),
|
||||
},
|
||||
{
|
||||
"metric": "acc",
|
||||
"mean": float(fold_df["acc"].mean()),
|
||||
"std": float(fold_df["acc"].std(ddof=0)),
|
||||
},
|
||||
]
|
||||
)
|
||||
out_dir = self.output_dir / "papila_logistic_regression"
|
||||
out_dir.mkdir(parents=True, exist_ok=True)
|
||||
fold_df.to_csv(out_dir / "fold_metrics.csv", index=False)
|
||||
summary.to_csv(out_dir / "summary.csv", index=False)
|
||||
self._plot_mean_roc(
|
||||
curves,
|
||||
"Logistic Regression ROC (mean ± SD)",
|
||||
out_dir / "roc_mean.png",
|
||||
)
|
||||
return fold_df
|
||||
|
||||
|
||||
ba = basic_analytics()
|
||||
roc_df = ba.univariate_roc(merge=False, include_categorical=False)
|
||||
rf_df = ba.random_forest(include_categorical=True, nerf=False)
|
||||
svm_df = ba.svm(include_categorical=True)
|
||||
knn_df = ba.knn(include_categorical=True)
|
||||
lr_df = ba.logistic_regression(include_categorical=True, nerf=True)
|
||||
clinical = ba._build_clinical()
|
||||
clinical.df["Diagnosis"].value_counts()
|
||||
@@ -1,327 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Train a CNN (resnet50 backbone), extract logits, and train RF on logits+metadata with 5-fold CV."""
|
||||
from __future__ import annotations
|
||||
|
||||
import random
|
||||
from pathlib import Path
|
||||
import sys
|
||||
from typing import Dict, List, Tuple
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
from PIL import Image
|
||||
import torch
|
||||
from torch import nn
|
||||
from torch.utils.data import DataLoader, Dataset
|
||||
from torchvision import transforms
|
||||
from sklearn.ensemble import RandomForestClassifier
|
||||
from sklearn.metrics import accuracy_score, roc_auc_score
|
||||
|
||||
REPO_ROOT = Path(__file__).resolve().parents[2]
|
||||
if str(REPO_ROOT) not in sys.path:
|
||||
sys.path.insert(0, str(REPO_ROOT))
|
||||
|
||||
from classes import build_papila_clinical
|
||||
from classes.backbones import BACKBONES, load_backbone_weights
|
||||
|
||||
|
||||
# ---------------------------
|
||||
# Config (edit in IDE)
|
||||
# ---------------------------
|
||||
IMAGE_DIR = "Papila/FundusImages"
|
||||
CLINICAL_DIR = "Papila/ClinicalData"
|
||||
LABEL_COL = "Diagnosis"
|
||||
CAT_COLS = ["Gender", "Phakic/Pseudophakic"]
|
||||
EVAL_MODE = "binary" # "binary" or "multiclass"
|
||||
N_SPLITS = 5
|
||||
FOLD_SEED = 42
|
||||
HOLDOUT_SEED = 123
|
||||
HOLDOUT_PATIENTS_PER_CLASS = 6
|
||||
|
||||
BACKBONE_NAME = "resnet50"
|
||||
BATCH_SIZE = 8
|
||||
EPOCHS = 40
|
||||
LR = 1e-4
|
||||
WEIGHT_DECAY = 1e-5
|
||||
|
||||
RF_TREES = 500
|
||||
RF_MAX_DEPTH = None
|
||||
RF_MIN_SAMPLES_LEAF = 1
|
||||
|
||||
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
OUTPUT_DIR = Path("analysis_data/basic_analysis/cnn_logits_rf_cv")
|
||||
PRINT_EPOCH_REPORT = True
|
||||
EPOCH_REPORT_EVERY = 1
|
||||
|
||||
|
||||
class PapilaImageDataset(Dataset):
|
||||
def __init__(
|
||||
self, clinical, df: pd.DataFrame, label_col: str, img_transform
|
||||
) -> None:
|
||||
self.clinical = clinical
|
||||
self.df = df.reset_index(drop=True)
|
||||
self.label_col = label_col
|
||||
self.img_transform = img_transform
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self.df)
|
||||
|
||||
def __getitem__(self, idx: int):
|
||||
row = self.df.iloc[idx]
|
||||
img_path = self.clinical.get_image_path(row)
|
||||
image = Image.open(img_path).convert("RGB")
|
||||
x_img = self.img_transform(image)
|
||||
y = int(row[self.label_col])
|
||||
x_md = self.clinical.vectorize_row(row).astype(np.float32)
|
||||
return x_img, y, x_md
|
||||
|
||||
|
||||
class CNNHead(nn.Module):
|
||||
def __init__(self, backbone_name: str, num_classes: int) -> None:
|
||||
super().__init__()
|
||||
spec = BACKBONES[backbone_name]
|
||||
backbone = spec.ctor(weights=spec.weights_default)
|
||||
if backbone_name.startswith("refuge"):
|
||||
load_backbone_weights(backbone_name, backbone)
|
||||
out_dim, backbone = spec.strip(backbone)
|
||||
self.backbone = backbone
|
||||
self.head = nn.Linear(out_dim, num_classes)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
feats = self.backbone(x)
|
||||
return self.head(feats)
|
||||
|
||||
|
||||
def _set_seed(seed: int) -> None:
|
||||
random.seed(seed)
|
||||
np.random.seed(seed)
|
||||
torch.manual_seed(seed)
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.manual_seed_all(seed)
|
||||
|
||||
|
||||
def _auc_score(y_true: np.ndarray, probs: np.ndarray, num_classes: int) -> float:
|
||||
try:
|
||||
if num_classes == 2:
|
||||
return float(roc_auc_score(y_true, probs[:, 1]))
|
||||
return float(roc_auc_score(y_true, probs, multi_class="ovr", average="macro"))
|
||||
except Exception:
|
||||
return float("nan")
|
||||
|
||||
|
||||
def _prepare_clinical() -> Tuple[object, pd.DataFrame]:
|
||||
clinical = build_papila_clinical(
|
||||
image_dir=IMAGE_DIR,
|
||||
clinical_dir=CLINICAL_DIR,
|
||||
label_col=LABEL_COL,
|
||||
cat_cols=CAT_COLS,
|
||||
n_splits=N_SPLITS,
|
||||
random_seed=FOLD_SEED,
|
||||
)
|
||||
df = clinical.df.copy()
|
||||
if EVAL_MODE == "binary":
|
||||
df = df[df[LABEL_COL].isin([0, 1])].reset_index(drop=True)
|
||||
return clinical, df
|
||||
|
||||
|
||||
def _split_holdout_by_patient(df: pd.DataFrame) -> Tuple[pd.DataFrame, pd.DataFrame]:
|
||||
rng = np.random.default_rng(HOLDOUT_SEED)
|
||||
patient_label = (
|
||||
df.groupby("Patient ID")[LABEL_COL]
|
||||
.agg(lambda s: int(s.mode().iloc[0]))
|
||||
.reset_index()
|
||||
)
|
||||
holdout_patients = []
|
||||
for lbl, grp in patient_label.groupby(LABEL_COL):
|
||||
candidates = grp["Patient ID"].to_numpy()
|
||||
n = min(HOLDOUT_PATIENTS_PER_CLASS, len(candidates))
|
||||
if n <= 0:
|
||||
continue
|
||||
selected = rng.choice(candidates, size=n, replace=False)
|
||||
holdout_patients.extend(selected.tolist())
|
||||
holdout_patients = sorted(set(holdout_patients))
|
||||
holdout_df = df[df["Patient ID"].isin(holdout_patients)].reset_index(drop=True)
|
||||
train_df = df[~df["Patient ID"].isin(holdout_patients)].reset_index(drop=True)
|
||||
return train_df, holdout_df
|
||||
|
||||
|
||||
def _rebuild_clinical_from_df(clinical, df: pd.DataFrame) -> object:
|
||||
clinical.frames = [df.copy()]
|
||||
clinical.df = df.copy()
|
||||
clinical._infer_or_validate_feature_types()
|
||||
clinical._compute_numeric_stats()
|
||||
clinical._build_cat_maps()
|
||||
clinical._compute_feature_dim()
|
||||
clinical._build_kfold_indices()
|
||||
return clinical
|
||||
|
||||
|
||||
def _train_cnn(
|
||||
model: nn.Module, loader: DataLoader, num_classes: int, fold: int
|
||||
) -> None:
|
||||
model.train()
|
||||
optimizer = torch.optim.Adam(model.parameters(), lr=LR, weight_decay=WEIGHT_DECAY)
|
||||
criterion = nn.CrossEntropyLoss()
|
||||
for epoch in range(EPOCHS):
|
||||
running_loss = 0.0
|
||||
correct = 0
|
||||
total = 0
|
||||
for x_img, y, _x_md in loader:
|
||||
x_img = x_img.to(DEVICE)
|
||||
y = y.to(DEVICE)
|
||||
optimizer.zero_grad()
|
||||
logits = model(x_img)
|
||||
loss = criterion(logits, y)
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
running_loss += float(loss.item()) * int(y.size(0))
|
||||
pred = torch.argmax(logits, dim=1)
|
||||
correct += int((pred == y).sum().item())
|
||||
total += int(y.size(0))
|
||||
|
||||
if PRINT_EPOCH_REPORT and ((epoch + 1) % EPOCH_REPORT_EVERY == 0):
|
||||
avg_loss = running_loss / max(total, 1)
|
||||
train_acc = correct / max(total, 1)
|
||||
print(
|
||||
f"[fold {fold + 1}/{N_SPLITS}] epoch {epoch + 1}/{EPOCHS} "
|
||||
f"train_loss={avg_loss:.4f} train_acc={train_acc:.4f}"
|
||||
)
|
||||
|
||||
|
||||
def _infer_logits(
|
||||
model: nn.Module, loader: DataLoader
|
||||
) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
|
||||
model.eval()
|
||||
logits_all, probs_all, y_all, md_all = [], [], [], []
|
||||
with torch.no_grad():
|
||||
for x_img, y, x_md in loader:
|
||||
x_img = x_img.to(DEVICE)
|
||||
logits = model(x_img).cpu().numpy()
|
||||
probs = torch.softmax(torch.from_numpy(logits), dim=1).numpy()
|
||||
logits_all.append(logits)
|
||||
probs_all.append(probs)
|
||||
y_all.append(y.numpy())
|
||||
md_all.append(x_md.numpy())
|
||||
return (
|
||||
np.concatenate(y_all, axis=0),
|
||||
np.concatenate(logits_all, axis=0),
|
||||
np.concatenate(md_all, axis=0),
|
||||
)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
_set_seed(FOLD_SEED)
|
||||
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
num_classes = 2 if EVAL_MODE == "binary" else 3
|
||||
train_tf = transforms.Compose(
|
||||
[
|
||||
transforms.Resize((224, 224)),
|
||||
transforms.RandomHorizontalFlip(),
|
||||
transforms.ToTensor(),
|
||||
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
|
||||
]
|
||||
)
|
||||
eval_tf = transforms.Compose(
|
||||
[
|
||||
transforms.Resize((224, 224)),
|
||||
transforms.ToTensor(),
|
||||
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
|
||||
]
|
||||
)
|
||||
|
||||
clinical, df = _prepare_clinical()
|
||||
train_df, holdout_df = _split_holdout_by_patient(df)
|
||||
clinical = _rebuild_clinical_from_df(clinical, train_df)
|
||||
holdout_df.to_csv(OUTPUT_DIR / "holdout_patients.csv", index=False)
|
||||
|
||||
rows: List[Dict[str, object]] = []
|
||||
holdout_rows: List[Dict[str, object]] = []
|
||||
|
||||
for fold in range(N_SPLITS):
|
||||
print(f"\n[info] Starting fold {fold + 1}/{N_SPLITS}")
|
||||
fold_train_df, fold_val_df = clinical.get_split_dfs(fold)
|
||||
ds_train = PapilaImageDataset(clinical, fold_train_df, LABEL_COL, train_tf)
|
||||
ds_val = PapilaImageDataset(clinical, fold_val_df, LABEL_COL, eval_tf)
|
||||
ds_holdout = PapilaImageDataset(clinical, holdout_df, LABEL_COL, eval_tf)
|
||||
|
||||
dl_train = DataLoader(
|
||||
ds_train, batch_size=BATCH_SIZE, shuffle=True, num_workers=0
|
||||
)
|
||||
dl_val = DataLoader(ds_val, batch_size=BATCH_SIZE, shuffle=False, num_workers=0)
|
||||
dl_holdout = DataLoader(
|
||||
ds_holdout, batch_size=BATCH_SIZE, shuffle=False, num_workers=0
|
||||
)
|
||||
|
||||
model = CNNHead(BACKBONE_NAME, num_classes=num_classes).to(DEVICE)
|
||||
_train_cnn(model, dl_train, num_classes=num_classes, fold=fold)
|
||||
|
||||
y_tr, log_tr, md_tr = _infer_logits(
|
||||
model,
|
||||
DataLoader(ds_train, batch_size=BATCH_SIZE, shuffle=False, num_workers=0),
|
||||
)
|
||||
y_va, log_va, md_va = _infer_logits(model, dl_val)
|
||||
y_ho, log_ho, md_ho = _infer_logits(model, dl_holdout)
|
||||
|
||||
np.save(OUTPUT_DIR / f"fold{fold}_train_logits.npy", log_tr)
|
||||
np.save(OUTPUT_DIR / f"fold{fold}_val_logits.npy", log_va)
|
||||
np.save(OUTPUT_DIR / f"fold{fold}_holdout_logits.npy", log_ho)
|
||||
|
||||
X_tr = np.concatenate([log_tr, md_tr], axis=1)
|
||||
X_va = np.concatenate([log_va, md_va], axis=1)
|
||||
X_ho = np.concatenate([log_ho, md_ho], axis=1)
|
||||
|
||||
rf = RandomForestClassifier(
|
||||
n_estimators=RF_TREES,
|
||||
max_depth=RF_MAX_DEPTH,
|
||||
min_samples_leaf=RF_MIN_SAMPLES_LEAF,
|
||||
class_weight="balanced",
|
||||
random_state=FOLD_SEED + fold,
|
||||
n_jobs=-1,
|
||||
)
|
||||
rf.fit(X_tr, y_tr)
|
||||
|
||||
p_va = rf.predict_proba(X_va)
|
||||
p_ho = rf.predict_proba(X_ho)
|
||||
pred_va = np.argmax(p_va, axis=1)
|
||||
pred_ho = np.argmax(p_ho, axis=1)
|
||||
|
||||
rows.append(
|
||||
{
|
||||
"fold": fold,
|
||||
"val_acc": float(accuracy_score(y_va, pred_va)),
|
||||
"val_auc": _auc_score(y_va, p_va, num_classes),
|
||||
"n_val": int(len(y_va)),
|
||||
}
|
||||
)
|
||||
holdout_rows.append(
|
||||
{
|
||||
"fold": fold,
|
||||
"holdout_acc": float(accuracy_score(y_ho, pred_ho)),
|
||||
"holdout_auc": _auc_score(y_ho, p_ho, num_classes),
|
||||
"n_holdout": int(len(y_ho)),
|
||||
}
|
||||
)
|
||||
print(
|
||||
f"[info] Fold {fold + 1} RF: val_acc={rows[-1]['val_acc']:.4f} val_auc={rows[-1]['val_auc']:.4f} "
|
||||
f"| holdout_acc={holdout_rows[-1]['holdout_acc']:.4f} holdout_auc={holdout_rows[-1]['holdout_auc']:.4f}"
|
||||
)
|
||||
|
||||
fold_df = pd.DataFrame(rows)
|
||||
holdout_df = pd.DataFrame(holdout_rows)
|
||||
fold_df.to_csv(OUTPUT_DIR / "rf_val_metrics.csv", index=False)
|
||||
holdout_df.to_csv(OUTPUT_DIR / "rf_holdout_metrics.csv", index=False)
|
||||
|
||||
print("\nRF validation metrics:")
|
||||
print(fold_df.to_string(index=False, float_format=lambda x: f"{x:.4f}"))
|
||||
print("\nRF holdout metrics:")
|
||||
print(holdout_df.to_string(index=False, float_format=lambda x: f"{x:.4f}"))
|
||||
print(
|
||||
f"\nMeans: val_acc={fold_df['val_acc'].mean():.4f}, val_auc={fold_df['val_auc'].mean():.4f}, "
|
||||
f"holdout_acc={holdout_df['holdout_acc'].mean():.4f}, holdout_auc={holdout_df['holdout_auc'].mean():.4f}"
|
||||
)
|
||||
print(f"\nSaved outputs to: {OUTPUT_DIR}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,399 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
"""v2 port of cnn_logits_rf_cv: CNN logit extraction + RF CV using the v2 PAPILA stack.
|
||||
|
||||
Replaces the hardcoded-config v1 version. Data loading, fold splitting, and
|
||||
feature preparation all go through the v2 stack so that --exclude-cols,
|
||||
--iop-corr-method, etc. are first-class options.
|
||||
|
||||
Usage:
|
||||
python scripts/basic_analysis/cnn_logits_rf_cv_v2.py \
|
||||
--eval-mode binary --backbone resnet50 --epochs 40 \
|
||||
--exclude-cols Phakic/Pseudophakic Axial_Length \
|
||||
--run-name cnn_rf_nocrop_binary
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import random
|
||||
import time
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from pathlib import Path
|
||||
import sys
|
||||
from types import SimpleNamespace
|
||||
from typing import Dict, List, Tuple
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
from PIL import Image
|
||||
import torch
|
||||
from torch import nn
|
||||
from torch.utils.data import DataLoader
|
||||
from sklearn.ensemble import RandomForestClassifier
|
||||
from sklearn.metrics import accuracy_score, roc_auc_score
|
||||
|
||||
REPO_ROOT = Path(__file__).resolve().parents[2]
|
||||
if str(REPO_ROOT) not in sys.path:
|
||||
sys.path.insert(0, str(REPO_ROOT))
|
||||
|
||||
from classes.v2.papila_builders import build_papila_data
|
||||
from classes.v2.backbones import BACKBONES, load_backbone_weights
|
||||
from classes.v2.dataset import ClinicalDataset, _ClinicalView
|
||||
from classes.v2.split_manager import PatientFirstSplitManager
|
||||
from classes.v2.transforms import build_backbone_transform, build_eval_transform
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Model
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class CNNHead(nn.Module):
|
||||
def __init__(self, backbone_name: str, num_classes: int) -> None:
|
||||
super().__init__()
|
||||
spec = BACKBONES[backbone_name]
|
||||
backbone = spec.ctor(weights=spec.weights_default)
|
||||
if backbone_name.startswith("refuge"):
|
||||
load_backbone_weights(backbone_name, backbone)
|
||||
out_dim, backbone = spec.strip(backbone)
|
||||
self.backbone = backbone
|
||||
self.head = nn.Linear(out_dim, num_classes)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
return self.head(self.backbone(x))
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _auc_score(y_true: np.ndarray, probs: np.ndarray, num_classes: int) -> float:
|
||||
try:
|
||||
if num_classes == 2:
|
||||
return float(roc_auc_score(y_true, probs[:, 1]))
|
||||
return float(roc_auc_score(y_true, probs, multi_class="ovr", average="macro"))
|
||||
except Exception:
|
||||
return float("nan")
|
||||
|
||||
|
||||
def _train_cnn(
|
||||
model: nn.Module,
|
||||
loader: DataLoader,
|
||||
args,
|
||||
fold: int,
|
||||
device: torch.device,
|
||||
) -> None:
|
||||
model.train()
|
||||
optimizer = torch.optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay)
|
||||
criterion = nn.CrossEntropyLoss()
|
||||
for epoch in range(args.epochs):
|
||||
running_loss = total = correct = 0
|
||||
for x_img, _meta, y in loader:
|
||||
x_img = x_img.to(device)
|
||||
y = y.to(device=device, dtype=torch.long)
|
||||
optimizer.zero_grad()
|
||||
logits = model(x_img)
|
||||
loss = criterion(logits, y)
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
running_loss += float(loss.item()) * int(y.size(0))
|
||||
correct += int((logits.argmax(1) == y).sum().item())
|
||||
total += int(y.size(0))
|
||||
if (epoch + 1) % args.log_every == 0:
|
||||
print(
|
||||
f" [fold {fold+1}] epoch {epoch+1}/{args.epochs} "
|
||||
f"loss={running_loss/max(total,1):.4f} acc={correct/max(total,1):.4f}",
|
||||
flush=True,
|
||||
)
|
||||
|
||||
|
||||
def _infer_logits(
|
||||
model: nn.Module, loader: DataLoader, device: torch.device
|
||||
) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
|
||||
"""Returns (y_true, logits, metadata_vectors)."""
|
||||
model.eval()
|
||||
y_all, logits_all, md_all = [], [], []
|
||||
with torch.no_grad():
|
||||
for x_img, x_md, y in loader:
|
||||
logits = model(x_img.to(device)).cpu().numpy()
|
||||
y_np = y.numpy() if torch.is_tensor(y) else np.asarray(y)
|
||||
md_np = x_md.numpy() if torch.is_tensor(x_md) else np.asarray(x_md)
|
||||
y_all.append(y_np)
|
||||
logits_all.append(logits)
|
||||
md_all.append(md_np)
|
||||
return (
|
||||
np.concatenate(y_all),
|
||||
np.concatenate(logits_all),
|
||||
np.concatenate(md_all),
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# CLI
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def build_parser() -> argparse.ArgumentParser:
|
||||
ap = argparse.ArgumentParser(
|
||||
description="CNN logit extraction + RF CV (v2 PAPILA stack)"
|
||||
)
|
||||
ap.add_argument("--image-dir", default="Papila/FundusImages")
|
||||
ap.add_argument("--clinical-dir", default="Papila/ClinicalData")
|
||||
ap.add_argument("--label-col", default="Diagnosis")
|
||||
ap.add_argument("--cat-cols", nargs="*", default=["Gender", "Phakic/Pseudophakic"])
|
||||
ap.add_argument("--exclude-cols", nargs="*", default=[],
|
||||
help="Feature columns to drop from the clinical feature matrix.")
|
||||
ap.add_argument("--eval-mode", choices=["binary", "multiclass"], default="binary")
|
||||
ap.add_argument("--n-splits", type=int, default=5)
|
||||
ap.add_argument("--fold-seed", type=int, default=42)
|
||||
ap.add_argument("--holdout-per-class", type=int, default=6,
|
||||
help="Patients per class reserved for holdout (0 disables).")
|
||||
ap.add_argument("--holdout-seed", type=int, default=123)
|
||||
ap.add_argument("--backbone", default="resnet50")
|
||||
ap.add_argument("--batch-size", type=int, default=8)
|
||||
ap.add_argument("--epochs", type=int, default=40)
|
||||
ap.add_argument("--lr", type=float, default=1e-4)
|
||||
ap.add_argument("--weight-decay", type=float, default=1e-5)
|
||||
ap.add_argument("--rf-trees", type=int, default=500)
|
||||
ap.add_argument("--rf-max-depth", type=int, default=None)
|
||||
ap.add_argument("--rf-min-samples-leaf", type=int, default=1)
|
||||
ap.add_argument("--num-workers", type=int, default=0)
|
||||
ap.add_argument("--device", choices=["auto", "cpu", "cuda"], default="auto")
|
||||
ap.add_argument("--seed", type=int, default=1234)
|
||||
ap.add_argument("--log-every", type=int, default=5)
|
||||
ap.add_argument("--run-name", default=None)
|
||||
ap.add_argument("--output-root", default="analysis_data/basic_analysis/cnn_logits_rf_cv_v2")
|
||||
ap.add_argument("--iop-corr-method", choices=["ratio", "ols", "lad", "multi"], default="ratio")
|
||||
ap.add_argument("--iop-drop-raw", action="store_true", default=False)
|
||||
ap.add_argument("--in-memory-cache", action="store_true", default=True,
|
||||
help="Cache all images in RAM before training (default: on).")
|
||||
ap.add_argument("--no-in-memory-cache", action="store_false", dest="in_memory_cache",
|
||||
help="Disable in-memory image cache.")
|
||||
ap.add_argument("--cache-workers", type=int, default=4,
|
||||
help="Threads for prebuilding image cache (default: 4).")
|
||||
return ap
|
||||
|
||||
|
||||
def _prebuild_image_cache(df: "pd.DataFrame", data, n_workers: int,
|
||||
resize: int = 256) -> dict:
|
||||
"""Load, convert to RGB, resize to `resize`px short edge, and cache as uint8 arrays.
|
||||
|
||||
Storing pre-resized images means the per-batch transform only has to do
|
||||
CenterCrop + augmentation + ToTensor + Normalize on a small image rather
|
||||
than resizing a full-resolution fundus image every step.
|
||||
"""
|
||||
paths = list({str(data.get_image_path(row)) for _, row in df.iterrows()})
|
||||
cache: dict = {}
|
||||
print(f"[cache] Preloading {len(paths)} images (resize={resize}px) "
|
||||
f"with {n_workers} threads...", flush=True)
|
||||
|
||||
def _load(p: str):
|
||||
img = Image.open(p)
|
||||
if img.mode != "RGB":
|
||||
img = img.convert("RGB")
|
||||
w, h = img.size
|
||||
scale = resize / min(w, h)
|
||||
img = img.resize((round(w * scale), round(h * scale)), Image.BILINEAR)
|
||||
return p, np.asarray(img, dtype=np.uint8)
|
||||
|
||||
with ThreadPoolExecutor(max_workers=max(1, n_workers)) as ex:
|
||||
for path, arr in ex.map(_load, paths):
|
||||
cache[path] = arr
|
||||
|
||||
ex_shape = cache[paths[0]].shape
|
||||
print(f"[cache] Done — {len(cache)} images in RAM "
|
||||
f"({ex_shape[1]}×{ex_shape[0]} each).", flush=True)
|
||||
return cache
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Main
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def main() -> None:
|
||||
args = build_parser().parse_args()
|
||||
|
||||
random.seed(args.seed)
|
||||
np.random.seed(args.seed)
|
||||
torch.manual_seed(args.seed)
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.manual_seed_all(args.seed)
|
||||
|
||||
if args.device == "auto":
|
||||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
else:
|
||||
device = torch.device(args.device)
|
||||
print(f"Device: {device}", flush=True)
|
||||
|
||||
run_name = args.run_name or time.strftime("%Y%m%d_%H%M%S")
|
||||
out_dir = Path(args.output_root) / run_name
|
||||
out_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# ---------- data ----------
|
||||
print("Loading PAPILA data...", flush=True)
|
||||
data = build_papila_data(
|
||||
image_dir=args.image_dir,
|
||||
clinical_dir=args.clinical_dir,
|
||||
label_col=args.label_col,
|
||||
cat_cols=list(args.cat_cols),
|
||||
n_splits=args.n_splits,
|
||||
random_seed=args.fold_seed,
|
||||
iop_corr_method=args.iop_corr_method,
|
||||
iop_drop_raw=args.iop_drop_raw,
|
||||
exclude_cols=list(args.exclude_cols or []),
|
||||
)
|
||||
print(f"Loaded: {len(data.df)} rows feature_dim={data.feature_dim}", flush=True)
|
||||
|
||||
df_mode = data.df.copy()
|
||||
if args.eval_mode == "binary":
|
||||
df_mode = df_mode[df_mode[args.label_col].isin([0, 1])].reset_index(drop=True)
|
||||
num_classes = 2 if args.eval_mode == "binary" else int(df_mode[args.label_col].nunique())
|
||||
print(
|
||||
f"eval_mode={args.eval_mode} num_classes={num_classes} "
|
||||
f"rows={len(df_mode)} patients={df_mode['Patient ID'].nunique()}",
|
||||
flush=True,
|
||||
)
|
||||
|
||||
# ---------- splits ----------
|
||||
split_manager = PatientFirstSplitManager(
|
||||
patient_col="Patient ID", label_col=args.label_col
|
||||
)
|
||||
split_args = SimpleNamespace(
|
||||
eval_mode=args.eval_mode,
|
||||
holdout_per_class=args.holdout_per_class,
|
||||
holdout_seed=args.holdout_seed,
|
||||
n_splits=args.n_splits,
|
||||
fold_seed=args.fold_seed,
|
||||
)
|
||||
clinical_ns = SimpleNamespace(df=df_mode, label_col=args.label_col)
|
||||
plans = split_manager.build_plans(clinical=clinical_ns, args=split_args, profile=None)
|
||||
n_folds = min(args.n_splits, len(plans))
|
||||
|
||||
if plans and plans[0].holdout is not None:
|
||||
plans[0].holdout.to_csv(out_dir / "holdout_patients.csv", index=False)
|
||||
|
||||
# ---------- image cache ----------
|
||||
image_cache = None
|
||||
if args.in_memory_cache:
|
||||
image_cache = _prebuild_image_cache(df_mode, data, args.cache_workers)
|
||||
|
||||
# ---------- transforms ----------
|
||||
train_tf = build_backbone_transform(args.backbone, augment=True)
|
||||
eval_tf = build_eval_transform(args.backbone)
|
||||
|
||||
rows: List[Dict] = []
|
||||
holdout_rows: List[Dict] = []
|
||||
|
||||
for fold, split in enumerate(plans[:n_folds]):
|
||||
print(f"\n[info] Fold {fold+1}/{n_folds}", flush=True)
|
||||
random.seed(args.seed + fold * 100)
|
||||
np.random.seed(args.seed + fold * 100)
|
||||
torch.manual_seed(args.seed + fold * 100)
|
||||
|
||||
view_train = _ClinicalView(data, split.train)
|
||||
view_val = _ClinicalView(data, split.val)
|
||||
|
||||
dl_train = DataLoader(
|
||||
ClinicalDataset(view_train, train_tf, image_cache=image_cache),
|
||||
batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers,
|
||||
)
|
||||
dl_val = DataLoader(
|
||||
ClinicalDataset(view_val, eval_tf, image_cache=image_cache),
|
||||
batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers,
|
||||
)
|
||||
dl_holdout = None
|
||||
if split.holdout is not None and not split.holdout.empty:
|
||||
dl_holdout = DataLoader(
|
||||
ClinicalDataset(_ClinicalView(data, split.holdout), eval_tf,
|
||||
image_cache=image_cache),
|
||||
batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers,
|
||||
)
|
||||
|
||||
# Train CNN
|
||||
model = CNNHead(args.backbone, num_classes=num_classes).to(device)
|
||||
_train_cnn(model, dl_train, args, fold=fold, device=device)
|
||||
|
||||
# Extract logits (re-run train without augmentation for RF features)
|
||||
dl_train_eval = DataLoader(
|
||||
ClinicalDataset(view_train, eval_tf, image_cache=image_cache),
|
||||
batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers,
|
||||
)
|
||||
y_tr, log_tr, md_tr = _infer_logits(model, dl_train_eval, device)
|
||||
y_va, log_va, md_va = _infer_logits(model, dl_val, device)
|
||||
|
||||
np.save(out_dir / f"fold{fold}_train_logits.npy", log_tr)
|
||||
np.save(out_dir / f"fold{fold}_val_logits.npy", log_va)
|
||||
|
||||
X_tr = np.concatenate([log_tr, md_tr], axis=1)
|
||||
X_va = np.concatenate([log_va, md_va], axis=1)
|
||||
|
||||
rf = RandomForestClassifier(
|
||||
n_estimators=args.rf_trees,
|
||||
max_depth=args.rf_max_depth,
|
||||
min_samples_leaf=args.rf_min_samples_leaf,
|
||||
class_weight="balanced",
|
||||
random_state=args.fold_seed + fold,
|
||||
n_jobs=-1,
|
||||
)
|
||||
rf.fit(X_tr, y_tr)
|
||||
|
||||
p_va = rf.predict_proba(X_va)
|
||||
pred_va = np.argmax(p_va, axis=1)
|
||||
rows.append({
|
||||
"fold": fold,
|
||||
"val_acc": float(accuracy_score(y_va, pred_va)),
|
||||
"val_auc": _auc_score(y_va, p_va, num_classes),
|
||||
"n_val": int(len(y_va)),
|
||||
})
|
||||
|
||||
if dl_holdout is not None:
|
||||
y_ho, log_ho, md_ho = _infer_logits(model, dl_holdout, device)
|
||||
np.save(out_dir / f"fold{fold}_holdout_logits.npy", log_ho)
|
||||
X_ho = np.concatenate([log_ho, md_ho], axis=1)
|
||||
p_ho = rf.predict_proba(X_ho)
|
||||
pred_ho = np.argmax(p_ho, axis=1)
|
||||
holdout_rows.append({
|
||||
"fold": fold,
|
||||
"holdout_acc": float(accuracy_score(y_ho, pred_ho)),
|
||||
"holdout_auc": _auc_score(y_ho, p_ho, num_classes),
|
||||
"n_holdout": int(len(y_ho)),
|
||||
})
|
||||
print(
|
||||
f"[info] Fold {fold+1} RF: val_acc={rows[-1]['val_acc']:.4f}"
|
||||
f" val_auc={rows[-1]['val_auc']:.4f}"
|
||||
f" | holdout_acc={holdout_rows[-1]['holdout_acc']:.4f}"
|
||||
f" holdout_auc={holdout_rows[-1]['holdout_auc']:.4f}",
|
||||
flush=True,
|
||||
)
|
||||
else:
|
||||
print(
|
||||
f"[info] Fold {fold+1} RF: val_acc={rows[-1]['val_acc']:.4f}"
|
||||
f" val_auc={rows[-1]['val_auc']:.4f}",
|
||||
flush=True,
|
||||
)
|
||||
|
||||
# ---------- save + summarise ----------
|
||||
fold_df = pd.DataFrame(rows)
|
||||
fold_df.to_csv(out_dir / "rf_val_metrics.csv", index=False)
|
||||
print("\nRF validation metrics:")
|
||||
print(fold_df.to_string(index=False, float_format=lambda x: f"{x:.4f}"))
|
||||
|
||||
if holdout_rows:
|
||||
ho_df = pd.DataFrame(holdout_rows)
|
||||
ho_df.to_csv(out_dir / "rf_holdout_metrics.csv", index=False)
|
||||
print("\nRF holdout metrics:")
|
||||
print(ho_df.to_string(index=False, float_format=lambda x: f"{x:.4f}"))
|
||||
print(
|
||||
f"\nMeans: val_acc={fold_df['val_acc'].mean():.4f}"
|
||||
f" val_auc={fold_df['val_auc'].mean():.4f}"
|
||||
f" holdout_acc={ho_df['holdout_acc'].mean():.4f}"
|
||||
f" holdout_auc={ho_df['holdout_auc'].mean():.4f}"
|
||||
)
|
||||
else:
|
||||
print(
|
||||
f"\nMeans: val_acc={fold_df['val_acc'].mean():.4f}"
|
||||
f" val_auc={fold_df['val_auc'].mean():.4f}"
|
||||
)
|
||||
|
||||
print(f"\nSaved outputs to: {out_dir}", flush=True)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,264 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Univariate ROC curves for PAPILA clinical variables."""
|
||||
import re
|
||||
from pathlib import Path
|
||||
from typing import Iterable, List, Tuple
|
||||
import sys
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import matplotlib.pyplot as plt
|
||||
from sklearn.metrics import roc_curve, auc
|
||||
|
||||
REPO_ROOT = Path(__file__).resolve().parents[2]
|
||||
if str(REPO_ROOT) not in sys.path:
|
||||
sys.path.insert(0, str(REPO_ROOT))
|
||||
|
||||
from classes import build_papila_clinical
|
||||
|
||||
|
||||
# ---------------------------
|
||||
# Config (edit in IDE)
|
||||
# ---------------------------
|
||||
IMAGE_DIR = "Papila/FundusImages"
|
||||
CLINICAL_DIR = "Papila/ClinicalData"
|
||||
LABEL_COL = "Diagnosis"
|
||||
CAT_COLS = ["Gender", "Phakic/Pseudophakic"]
|
||||
EXCLUDE_COLS = {"Pneumatic", "Perkins"}
|
||||
INCLUDE_CATEGORICAL = False
|
||||
POSITIVE_LABEL = 1
|
||||
NEGATIVE_LABEL = 0
|
||||
DROP_LABELS = [2]
|
||||
OUTPUT_DIR = Path("analysis_data/basic_analysis/papila_univariate_roc")
|
||||
DEBUG_PRINTS = False
|
||||
PLOT_PER_FEATURE = False
|
||||
DI_OPTRE_COL_PREFIXES = ("dioptre",)
|
||||
ADD_DIOPTRE_ABS = True
|
||||
ADD_DIOPTRE_SQUARED = True
|
||||
|
||||
|
||||
def _sanitize(name: str) -> str:
|
||||
safe = re.sub(r"[^A-Za-z0-9._-]+", "_", str(name)).strip("_")
|
||||
return safe or "var"
|
||||
|
||||
|
||||
def _select_binary_labels(labels: pd.Series,
|
||||
positive_label: int,
|
||||
negative_label: int,
|
||||
drop_labels: Iterable[int]) -> Tuple[np.ndarray, np.ndarray]:
|
||||
labels_num = pd.to_numeric(labels, errors="coerce")
|
||||
use_num = labels_num.notna().any()
|
||||
lab = labels_num if use_num else labels.astype(str)
|
||||
|
||||
drop_set = set(drop_labels or [])
|
||||
keep = lab.isin([positive_label, negative_label])
|
||||
if drop_set:
|
||||
keep &= ~lab.isin(drop_set)
|
||||
|
||||
y = (lab == positive_label).astype(int)
|
||||
return y.values, keep.values
|
||||
|
||||
|
||||
def _compute_roc(y: np.ndarray, scores: np.ndarray) -> Tuple[np.ndarray, np.ndarray, np.ndarray, float]:
|
||||
fpr, tpr, thresholds = roc_curve(y, scores, pos_label=1)
|
||||
auc_val = float(auc(fpr, tpr))
|
||||
return fpr, tpr, thresholds, auc_val
|
||||
|
||||
|
||||
def _best_threshold(fpr: np.ndarray, tpr: np.ndarray, thresholds: np.ndarray) -> Tuple[float, float, float]:
|
||||
youden = tpr - fpr
|
||||
idx = int(np.nanargmax(youden))
|
||||
return float(thresholds[idx]), float(tpr[idx]), float(fpr[idx])
|
||||
|
||||
|
||||
def _plot_roc(fpr: np.ndarray, tpr: np.ndarray, auc_val: float, title: str, out_path: Path) -> None:
|
||||
fig, ax = plt.subplots(figsize=(5.5, 4.5))
|
||||
ax.plot(fpr, tpr, lw=1.8, label=f"AUC={auc_val:.3f}")
|
||||
ax.plot([0, 1], [0, 1], "k--", lw=1)
|
||||
ax.set_xlabel("False Positive Rate")
|
||||
ax.set_ylabel("True Positive Rate")
|
||||
ax.set_title(title)
|
||||
ax.legend(loc="lower right")
|
||||
ax.grid(True, alpha=0.3, linestyle="--")
|
||||
fig.tight_layout()
|
||||
fig.savefig(out_path, dpi=170)
|
||||
plt.close(fig)
|
||||
|
||||
|
||||
def _plot_overlay(curves, title: str, out_path: Path) -> None:
|
||||
fig, ax = plt.subplots(figsize=(7, 5.5))
|
||||
cmap = plt.get_cmap("tab20")
|
||||
for i, (name, fpr, tpr, auc_val) in enumerate(curves):
|
||||
color = cmap(i % cmap.N)
|
||||
ax.plot(fpr, tpr, lw=1.6, color=color, label=f"{name} (AUC={auc_val:.3f})")
|
||||
ax.plot([0, 1], [0, 1], "k--", lw=1)
|
||||
ax.set_xlabel("False Positive Rate")
|
||||
ax.set_ylabel("True Positive Rate")
|
||||
ax.set_title(title)
|
||||
ax.legend(loc="upper left", fontsize="small")
|
||||
ax.grid(True, alpha=0.3, linestyle="--")
|
||||
fig.tight_layout()
|
||||
fig.savefig(out_path, dpi=170)
|
||||
plt.close(fig)
|
||||
|
||||
|
||||
def _iter_numeric(df: pd.DataFrame, cols: List[str]):
|
||||
for col in cols:
|
||||
if col not in df.columns:
|
||||
continue
|
||||
s = pd.to_numeric(df[col], errors="coerce")
|
||||
yield col, s
|
||||
|
||||
|
||||
def _is_dioptre_col(col: str) -> bool:
|
||||
name = str(col).strip().lower()
|
||||
return any(name.startswith(prefix) for prefix in DI_OPTRE_COL_PREFIXES)
|
||||
|
||||
|
||||
def _iter_numeric_with_transforms(df: pd.DataFrame, cols: List[str]):
|
||||
for col, s in _iter_numeric(df, cols):
|
||||
yield col, s
|
||||
if _is_dioptre_col(col):
|
||||
if ADD_DIOPTRE_ABS:
|
||||
yield f"{col}_abs", s.abs()
|
||||
if ADD_DIOPTRE_SQUARED:
|
||||
yield f"{col}_sq", s.pow(2)
|
||||
|
||||
|
||||
def _include_in_overlay(feature_name: str) -> bool:
|
||||
name = str(feature_name).strip().lower()
|
||||
if _is_dioptre_col(name) and not name.endswith("_abs"):
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def _iter_categorical(df: pd.DataFrame, cols: List[str]):
|
||||
for col in cols:
|
||||
if col not in df.columns:
|
||||
continue
|
||||
s = df[col]
|
||||
vals = s.dropna().unique().tolist()
|
||||
try:
|
||||
vals = sorted(vals)
|
||||
except Exception:
|
||||
pass
|
||||
for v in vals:
|
||||
name = f"{col}=={v}"
|
||||
ind = (s == v).astype(int)
|
||||
yield name, ind
|
||||
|
||||
|
||||
def main() -> None:
|
||||
clinical = build_papila_clinical(
|
||||
image_dir=IMAGE_DIR,
|
||||
clinical_dir=CLINICAL_DIR,
|
||||
label_col=LABEL_COL,
|
||||
cat_cols=CAT_COLS,
|
||||
)
|
||||
|
||||
df = clinical.df.copy()
|
||||
labels = df[LABEL_COL]
|
||||
y_all, keep_mask = _select_binary_labels(labels, POSITIVE_LABEL, NEGATIVE_LABEL, DROP_LABELS)
|
||||
|
||||
out_dir = OUTPUT_DIR
|
||||
plot_dir = out_dir / "plots"
|
||||
plot_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
rows = []
|
||||
curves = []
|
||||
|
||||
base_exclude = {LABEL_COL, "Patient ID"} | EXCLUDE_COLS | set(CAT_COLS)
|
||||
if "eyeID" not in CAT_COLS:
|
||||
base_exclude.add("eyeID")
|
||||
candidate_cols = [c for c in df.columns if c not in base_exclude]
|
||||
numeric_cols = []
|
||||
for col in candidate_cols:
|
||||
s = pd.to_numeric(df[col], errors="coerce")
|
||||
if s.notna().any():
|
||||
numeric_cols.append(col)
|
||||
|
||||
if DEBUG_PRINTS:
|
||||
for col in ("IOP_raw", "IOP_corr"):
|
||||
if col not in df.columns:
|
||||
print(f"[debug] {col} missing from df")
|
||||
continue
|
||||
s = pd.to_numeric(df[col], errors="coerce")
|
||||
print(f"[debug] {col}: non-null={int(s.notna().sum())}, unique={int(s.nunique(dropna=True))}")
|
||||
for col, series in _iter_numeric_with_transforms(df, numeric_cols):
|
||||
mask = keep_mask & series.notna().values
|
||||
y = y_all[mask]
|
||||
scores = series.values[mask].astype(float)
|
||||
if y.size < 2 or np.unique(y).size < 2:
|
||||
continue
|
||||
if np.nanmin(scores) == np.nanmax(scores):
|
||||
continue
|
||||
fpr, tpr, thresholds, auc_val = _compute_roc(y, scores)
|
||||
thr, best_tpr, best_fpr = _best_threshold(fpr, tpr, thresholds)
|
||||
direction = "high" if auc_val >= 0.5 else "low"
|
||||
title = f"{col} (n={y.size}, direction={direction})"
|
||||
if PLOT_PER_FEATURE:
|
||||
out_path = plot_dir / f"roc_{_sanitize(col)}.png"
|
||||
_plot_roc(fpr, tpr, auc_val, title, out_path)
|
||||
if _include_in_overlay(col):
|
||||
curves.append((col, fpr, tpr, auc_val))
|
||||
rows.append({
|
||||
"feature": col,
|
||||
"kind": "numeric",
|
||||
"n": int(y.size),
|
||||
"auc": auc_val,
|
||||
"direction": direction,
|
||||
"best_threshold": thr,
|
||||
"best_tpr": best_tpr,
|
||||
"best_fpr": best_fpr,
|
||||
"best_specificity": 1.0 - best_fpr,
|
||||
})
|
||||
|
||||
if INCLUDE_CATEGORICAL:
|
||||
cat_cols_use = [c for c in clinical.cat_cols if c not in EXCLUDE_COLS]
|
||||
for name, ind in _iter_categorical(df, cat_cols_use):
|
||||
mask = keep_mask & ind.notna().values
|
||||
y = y_all[mask]
|
||||
scores = ind.values[mask].astype(float)
|
||||
if y.size < 2 or np.unique(y).size < 2:
|
||||
continue
|
||||
if np.nanmin(scores) == np.nanmax(scores):
|
||||
continue
|
||||
fpr, tpr, thresholds, auc_val = _compute_roc(y, scores)
|
||||
thr, best_tpr, best_fpr = _best_threshold(fpr, tpr, thresholds)
|
||||
direction = "high" if auc_val >= 0.5 else "low"
|
||||
title = f"{name} (n={y.size}, direction={direction})"
|
||||
if PLOT_PER_FEATURE:
|
||||
out_path = plot_dir / f"roc_{_sanitize(name)}.png"
|
||||
_plot_roc(fpr, tpr, auc_val, title, out_path)
|
||||
if _include_in_overlay(name):
|
||||
curves.append((name, fpr, tpr, auc_val))
|
||||
rows.append({
|
||||
"feature": name,
|
||||
"kind": "categorical",
|
||||
"n": int(y.size),
|
||||
"auc": auc_val,
|
||||
"direction": direction,
|
||||
"best_threshold": thr,
|
||||
"best_tpr": best_tpr,
|
||||
"best_fpr": best_fpr,
|
||||
"best_specificity": 1.0 - best_fpr,
|
||||
})
|
||||
|
||||
if not rows:
|
||||
raise SystemExit("No valid features produced ROC curves. Check labels and feature columns.")
|
||||
|
||||
overlay_path = plot_dir / "roc_overlay.png"
|
||||
_plot_overlay(curves, "Univariate ROC curves", overlay_path)
|
||||
|
||||
out_df = pd.DataFrame(rows).sort_values(by="auc", ascending=False)
|
||||
out_dir.mkdir(parents=True, exist_ok=True)
|
||||
out_df.to_csv(out_dir / "summary.csv", index=False)
|
||||
print(out_df.to_string(index=False, float_format=lambda x: f"{x:.4f}"))
|
||||
print(f"\nSaved overlay plot to: {overlay_path}")
|
||||
if PLOT_PER_FEATURE:
|
||||
print(f"Saved per-feature plots to: {plot_dir}")
|
||||
print(f"Saved summary to: {out_dir / 'summary.csv'}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,142 +0,0 @@
|
||||
"""Generate ground-truth mask overlays for REFUGE and Papila samples."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
from pathlib import Path
|
||||
import sys
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[1]
|
||||
sys.path.append(str(ROOT))
|
||||
|
||||
import numpy as np
|
||||
from collections import Counter
|
||||
from PIL import Image
|
||||
from PIL.Image import Resampling
|
||||
|
||||
from classes.unet_segmenter import UNetSegmenter
|
||||
|
||||
REFUGE_ROOT = Path("REFUGE")
|
||||
DEFAULT_MANIFEST = Path("manifest.csv")
|
||||
OUTPUT_DIR = Path("temp/gt_test")
|
||||
|
||||
|
||||
def to_mask_colors(disc: np.ndarray, cup: np.ndarray) -> Image.Image:
|
||||
h, w = disc.shape
|
||||
canvas = np.ones((h, w, 3), dtype=np.uint8) * 255
|
||||
disc_mask = disc.astype(bool)
|
||||
cup_mask = cup.astype(bool)
|
||||
canvas[disc_mask] = [128, 128, 128]
|
||||
canvas[cup_mask] = [0, 0, 0]
|
||||
return Image.fromarray(canvas)
|
||||
|
||||
|
||||
def overlay(
|
||||
original: Image.Image, mask_rgb: Image.Image, alpha: float = 0.6
|
||||
) -> Image.Image:
|
||||
mask_rgba = mask_rgb.convert("RGBA")
|
||||
updates = np.array(mask_rgba, dtype=np.float32)
|
||||
updates[..., 3] = alpha * 255 * (updates[..., :3] != 255).any(axis=-1)
|
||||
base = original.convert("RGBA")
|
||||
return Image.alpha_composite(
|
||||
base, Image.fromarray(updates.astype(np.uint8))
|
||||
).convert("RGB")
|
||||
|
||||
|
||||
def original_mask_to_rgb(mask_path: Path) -> Image.Image:
|
||||
mask_img = Image.open(mask_path)
|
||||
arr = np.asarray(mask_img)
|
||||
h, w = arr.shape[:2]
|
||||
canvas = np.ones((h, w, 3), dtype=np.uint8) * 255
|
||||
|
||||
if arr.ndim == 2:
|
||||
border = np.concatenate([arr[0, :], arr[-1, :], arr[:, 0], arr[:, -1]])
|
||||
bg_value = Counter(border.tolist()).most_common(1)[0][0]
|
||||
disc_mask = arr != bg_value
|
||||
fg_counts = Counter(arr[arr != bg_value].flatten())
|
||||
if fg_counts:
|
||||
# For REFUGE-style masks: cup should be the darkest (minimum value)
|
||||
cup_value = min(fg_counts.keys())
|
||||
cup_mask = arr == cup_value
|
||||
else:
|
||||
cup_mask = np.zeros_like(arr, dtype=bool)
|
||||
else:
|
||||
edges = np.concatenate(
|
||||
[arr[0, :, :], arr[-1, :, :], arr[:, 0, :], arr[:, -1, :]], axis=0
|
||||
)
|
||||
bg_color = Counter(map(tuple, edges)).most_common(1)[0][0]
|
||||
disc_mask = ~np.all(arr == bg_color, axis=-1)
|
||||
color_counts = Counter(map(tuple, arr.reshape(-1, arr.shape[2])))
|
||||
cup_mask = np.zeros((h, w), dtype=bool)
|
||||
candidates = {}
|
||||
for color, count in color_counts.items():
|
||||
if color == bg_color:
|
||||
continue
|
||||
mask = np.all(arr == color, axis=-1)
|
||||
candidates[color] = mask
|
||||
if candidates:
|
||||
# For REFUGE-style masks: cup should be the darkest color (closest to black)
|
||||
cup_color = min(candidates.keys(), key=lambda color: sum(color))
|
||||
cup_mask = candidates[cup_color]
|
||||
disc_mask = disc_mask.astype(bool)
|
||||
cup_mask = cup_mask & disc_mask
|
||||
|
||||
canvas[disc_mask] = [128, 128, 128]
|
||||
canvas[cup_mask] = [0, 0, 0]
|
||||
return Image.fromarray(canvas)
|
||||
|
||||
|
||||
def process_entries(
|
||||
segmenter: UNetSegmenter, entries, prefix: str, count: int, dest: Path
|
||||
) -> None:
|
||||
for entry in entries[:count]:
|
||||
img_path = Path(entry.image_path)
|
||||
if not img_path.exists():
|
||||
continue
|
||||
orig = Image.open(img_path).convert("RGB")
|
||||
image = segmenter.preprocess_image(orig)
|
||||
disc, cup = segmenter.load_masks(entry)
|
||||
disc_coords = segmenter._mask_to_coords(disc)
|
||||
cup_coords = segmenter._mask_to_coords(cup)
|
||||
print(
|
||||
f"{entry.sample_id}: disc coords {disc_coords.shape[0] if disc_coords is not None else 0}, "
|
||||
f"cup coords {cup_coords.shape[0] if cup_coords is not None else 0}"
|
||||
)
|
||||
mask_rgb = to_mask_colors(disc, cup)
|
||||
overlay_img = overlay(image, mask_rgb)
|
||||
mask_rgb.save(dest / f"{prefix}_{entry.sample_id}_mask.png")
|
||||
overlay_img.save(dest / f"{prefix}_{entry.sample_id}_overlay.png")
|
||||
|
||||
if entry.annotation_type_disc == "mask":
|
||||
gt_mask_rgb = original_mask_to_rgb(entry.annotation_disc)
|
||||
gt_overlay = overlay(
|
||||
orig.resize(gt_mask_rgb.size, Resampling.BILINEAR), gt_mask_rgb
|
||||
)
|
||||
gt_mask_rgb.save(dest / f"{prefix}_{entry.sample_id}_gt_mask.png")
|
||||
gt_overlay.save(dest / f"{prefix}_{entry.sample_id}_gt_overlay.png")
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Inspect ground-truth masks for REFUGE and Papila"
|
||||
)
|
||||
parser.add_argument("--manifest", type=Path, default=DEFAULT_MANIFEST)
|
||||
parser.add_argument("--output", type=Path, default=OUTPUT_DIR)
|
||||
parser.add_argument(
|
||||
"--count", type=int, default=10, help="Number of samples per dataset"
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
args.output.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
segmenter = UNetSegmenter(args.manifest)
|
||||
refuge_entries = [e for e in segmenter._manifest if e.dataset == "refuge"]
|
||||
papila_entries = [e for e in segmenter._manifest if e.dataset == "papila"]
|
||||
|
||||
process_entries(segmenter, refuge_entries, "refuge", args.count, args.output)
|
||||
process_entries(segmenter, papila_entries, "papila", args.count, args.output)
|
||||
print(f"Saved overlays to {args.output}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,725 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Compare single-eye OD baseline vs SiameseImageTower bilateral model.
|
||||
|
||||
Key differences from compare_dual_eye_towers.py:
|
||||
- Uses SiameseImageTower (shared backbone, f_mean + f_delta output).
|
||||
- Reports BEST-epoch val metrics per fold (not final-epoch), with the
|
||||
corresponding holdout metrics snapped at the same checkpoint.
|
||||
- Both models are always evaluated on patient-level samples (matched n).
|
||||
- Optionally includes two-single-merge as a second reference point.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import copy
|
||||
import csv
|
||||
import json
|
||||
import random
|
||||
import sys
|
||||
import time
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
from typing import Optional
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from sklearn.metrics import roc_auc_score
|
||||
from torch import nn
|
||||
from torch.utils.data import DataLoader
|
||||
|
||||
REPO_ROOT = Path(__file__).resolve().parents[2]
|
||||
if str(REPO_ROOT) not in sys.path:
|
||||
sys.path.insert(0, str(REPO_ROOT))
|
||||
|
||||
from classes.v2 import (
|
||||
ImageTower,
|
||||
PatientFirstSplitManager,
|
||||
SiameseImageTower,
|
||||
SlotDataset,
|
||||
build_papila_data,
|
||||
build_papila_profile,
|
||||
slot_collate,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Reproducibility
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def seed_everything(seed: int) -> None:
|
||||
random.seed(seed)
|
||||
np.random.seed(seed)
|
||||
torch.manual_seed(seed)
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.manual_seed_all(seed)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Models
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class BaselineCNN(nn.Module):
|
||||
"""Single-eye (OD) image tower with a linear head."""
|
||||
|
||||
def __init__(self, *, backbone: str, freeze_ratio: float, num_classes: int, augment: bool):
|
||||
super().__init__()
|
||||
self.tower = ImageTower(
|
||||
backbone=backbone,
|
||||
freeze_ratio=freeze_ratio,
|
||||
augment=augment,
|
||||
use_se=False,
|
||||
)
|
||||
self.head = nn.Linear(self.tower.out_dim, num_classes)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
return self.head(self.tower(x))
|
||||
|
||||
|
||||
class SiameseCNN(nn.Module):
|
||||
"""
|
||||
Bilateral image model using SiameseImageTower.
|
||||
forward(x_od, x_os) -> logits
|
||||
"""
|
||||
|
||||
def __init__(self, *, backbone: str, freeze_ratio: float, num_classes: int, augment: bool):
|
||||
super().__init__()
|
||||
self.tower = SiameseImageTower(
|
||||
backbone=backbone,
|
||||
freeze_ratio=freeze_ratio,
|
||||
augment=augment,
|
||||
use_se=False,
|
||||
)
|
||||
self.head = nn.Linear(self.tower.out_dim, num_classes)
|
||||
|
||||
def forward(self, x_od: torch.Tensor, x_os: torch.Tensor) -> torch.Tensor:
|
||||
return self.head(self.tower(x_od, x_os))
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Data helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def filter_od_samples(patient_samples: list[dict]) -> list[dict]:
|
||||
"""Extract patient-level OD-only samples (image_1 = OD)."""
|
||||
out = []
|
||||
for s in patient_samples:
|
||||
if s.get("image_1") is not None and s.get("label_1") is not None:
|
||||
out.append({"id_1": s.get("id_1"), "image_1": s["image_1"], "label_1": s["label_1"]})
|
||||
return out
|
||||
|
||||
|
||||
def filter_bilateral_samples(patient_samples: list[dict]) -> list[dict]:
|
||||
return [
|
||||
s for s in patient_samples
|
||||
if s.get("image_1") is not None
|
||||
and s.get("image_2") is not None
|
||||
and s.get("label_1") is not None
|
||||
]
|
||||
|
||||
|
||||
def make_loader(samples, slots, *, image_transform, batch_size, shuffle, num_workers) -> DataLoader:
|
||||
ds = SlotDataset(samples, slots, image_transform=image_transform)
|
||||
return DataLoader(
|
||||
ds,
|
||||
batch_size=batch_size,
|
||||
shuffle=shuffle,
|
||||
num_workers=num_workers,
|
||||
collate_fn=slot_collate,
|
||||
)
|
||||
|
||||
|
||||
def to_label_tensor(labels, device: torch.device) -> torch.Tensor:
|
||||
if torch.is_tensor(labels):
|
||||
return labels.to(device=device, dtype=torch.long)
|
||||
return torch.as_tensor(labels, dtype=torch.long, device=device)
|
||||
|
||||
|
||||
def _drop_mixed_label_patients(df, *, patient_col: str, label_col: str):
|
||||
import pandas as pd
|
||||
per_patient = (
|
||||
df.groupby(patient_col)[label_col]
|
||||
.agg(lambda s: set(pd.to_numeric(s, errors="coerce").dropna().astype(int).tolist()))
|
||||
)
|
||||
mixed = [pid for pid, labels in per_patient.items() if len(labels) > 1]
|
||||
if not mixed:
|
||||
return df, []
|
||||
return df[~df[patient_col].isin(mixed)].reset_index(drop=True), mixed
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Score helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _score(y_true_chunks, y_prob_chunks, num_classes: int):
|
||||
if not y_true_chunks:
|
||||
return float("nan"), float("nan"), 0
|
||||
y = np.concatenate(y_true_chunks)
|
||||
p = np.concatenate(y_prob_chunks)
|
||||
acc = float((p.argmax(1) == y).mean())
|
||||
try:
|
||||
auc = (
|
||||
float(roc_auc_score(y, p[:, 1]))
|
||||
if num_classes == 2
|
||||
else float(roc_auc_score(y, p, multi_class="ovr", average="macro"))
|
||||
)
|
||||
except Exception:
|
||||
auc = float("nan")
|
||||
return acc, auc, int(len(y))
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Train / evaluate
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def train_baseline_epoch(model, loader, opt, device):
|
||||
model.train()
|
||||
total_loss = total_correct = total_n = 0
|
||||
for batch in loader:
|
||||
x = batch.get("image_1")
|
||||
y = batch.get("label_1")
|
||||
if not torch.is_tensor(x):
|
||||
continue
|
||||
y = to_label_tensor(y, device)
|
||||
x = x.to(device)
|
||||
logits = model(x)
|
||||
loss = F.cross_entropy(logits, y)
|
||||
opt.zero_grad()
|
||||
loss.backward()
|
||||
opt.step()
|
||||
bs = y.shape[0]
|
||||
total_loss += float(loss.item()) * bs
|
||||
total_correct += int((logits.argmax(1) == y).sum())
|
||||
total_n += bs
|
||||
return (total_loss / total_n if total_n else float("nan"),
|
||||
total_correct / total_n if total_n else float("nan"))
|
||||
|
||||
|
||||
def train_siamese_epoch(model, loader, opt, device):
|
||||
model.train()
|
||||
total_loss = total_correct = total_n = 0
|
||||
for batch in loader:
|
||||
x1 = batch.get("image_1")
|
||||
x2 = batch.get("image_2")
|
||||
y = batch.get("label_1")
|
||||
if not torch.is_tensor(x1) or not torch.is_tensor(x2):
|
||||
continue
|
||||
y = to_label_tensor(y, device)
|
||||
logits = model(x1.to(device), x2.to(device))
|
||||
loss = F.cross_entropy(logits, y)
|
||||
opt.zero_grad()
|
||||
loss.backward()
|
||||
opt.step()
|
||||
bs = y.shape[0]
|
||||
total_loss += float(loss.item()) * bs
|
||||
total_correct += int((logits.argmax(1) == y).sum())
|
||||
total_n += bs
|
||||
return (total_loss / total_n if total_n else float("nan"),
|
||||
total_correct / total_n if total_n else float("nan"))
|
||||
|
||||
|
||||
def evaluate_baseline(model, loader, device, num_classes):
|
||||
model.eval()
|
||||
y_true, y_prob = [], []
|
||||
with torch.no_grad():
|
||||
for batch in loader:
|
||||
x = batch.get("image_1")
|
||||
y = batch.get("label_1")
|
||||
if not torch.is_tensor(x):
|
||||
continue
|
||||
y_t = to_label_tensor(y, device)
|
||||
p = F.softmax(model(x.to(device)), dim=1).cpu().numpy()
|
||||
y_true.append(y_t.cpu().numpy())
|
||||
y_prob.append(p)
|
||||
return _score(y_true, y_prob, num_classes)
|
||||
|
||||
|
||||
def evaluate_siamese(model, loader, device, num_classes):
|
||||
model.eval()
|
||||
y_true, y_prob = [], []
|
||||
with torch.no_grad():
|
||||
for batch in loader:
|
||||
x1 = batch.get("image_1")
|
||||
x2 = batch.get("image_2")
|
||||
y = batch.get("label_1")
|
||||
if not torch.is_tensor(x1) or not torch.is_tensor(x2):
|
||||
continue
|
||||
y_t = to_label_tensor(y, device)
|
||||
p = F.softmax(model(x1.to(device), x2.to(device)), dim=1).cpu().numpy()
|
||||
y_true.append(y_t.cpu().numpy())
|
||||
y_prob.append(p)
|
||||
return _score(y_true, y_prob, num_classes)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Result dataclass
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@dataclass
|
||||
class FoldResult:
|
||||
mode: str
|
||||
fold: int
|
||||
# Best-epoch validation metrics
|
||||
best_epoch: int
|
||||
baseline_best_val_auc: float
|
||||
baseline_best_val_acc: float
|
||||
siamese_best_val_auc: float
|
||||
siamese_best_val_acc: float
|
||||
# Holdout metrics at the respective best-epoch checkpoint
|
||||
baseline_holdout_auc: float
|
||||
baseline_holdout_acc: float
|
||||
baseline_holdout_n: int
|
||||
siamese_holdout_auc: float
|
||||
siamese_holdout_acc: float
|
||||
siamese_holdout_n: int
|
||||
# Sample sizes
|
||||
baseline_n: int
|
||||
siamese_n: int
|
||||
|
||||
|
||||
def _nan() -> float:
|
||||
return float("nan")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Main fold runner
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def run_fold(
|
||||
fold: int,
|
||||
split,
|
||||
mode: str,
|
||||
args,
|
||||
device: torch.device,
|
||||
data,
|
||||
num_classes: int,
|
||||
profile_od,
|
||||
profile_patient,
|
||||
fold_dir: Path,
|
||||
) -> FoldResult:
|
||||
holdout_df = split.holdout
|
||||
|
||||
# ---- build samples ----
|
||||
bilat_train = filter_bilateral_samples(profile_patient.build_samples(df=split.train, clinical=data))
|
||||
bilat_val = filter_bilateral_samples(profile_patient.build_samples(df=split.val, clinical=data))
|
||||
od_train = filter_od_samples(bilat_train)
|
||||
od_val = filter_od_samples(bilat_val)
|
||||
|
||||
bilat_holdout = []
|
||||
od_holdout = []
|
||||
if holdout_df is not None and not holdout_df.empty:
|
||||
bilat_holdout = filter_bilateral_samples(profile_patient.build_samples(df=holdout_df, clinical=data))
|
||||
od_holdout = filter_od_samples(bilat_holdout)
|
||||
|
||||
# ---- models ----
|
||||
baseline = BaselineCNN(
|
||||
backbone=args.backbone, freeze_ratio=args.freeze_ratio,
|
||||
num_classes=num_classes, augment=args.augment,
|
||||
).to(device)
|
||||
siamese = SiameseCNN(
|
||||
backbone=args.backbone, freeze_ratio=args.freeze_ratio,
|
||||
num_classes=num_classes, augment=args.augment,
|
||||
).to(device)
|
||||
|
||||
slots_od = profile_od.slot_descriptors()
|
||||
slots_patient = profile_patient.slot_descriptors()
|
||||
|
||||
# ---- loaders ----
|
||||
loader_kw = dict(batch_size=args.batch_size, num_workers=args.num_workers)
|
||||
train_base = make_loader(od_train, slots_od, image_transform=baseline.tower.transform, shuffle=True, **loader_kw)
|
||||
val_base = make_loader(od_val, slots_od, image_transform=baseline.tower.transform, shuffle=False, **loader_kw)
|
||||
train_siam = make_loader(bilat_train, slots_patient, image_transform=siamese.tower.transform, shuffle=True, **loader_kw)
|
||||
val_siam = make_loader(bilat_val, slots_patient, image_transform=siamese.tower.transform, shuffle=False, **loader_kw)
|
||||
|
||||
ho_base = (
|
||||
make_loader(od_holdout, slots_od, image_transform=baseline.tower.transform, shuffle=False, **loader_kw)
|
||||
if od_holdout else None
|
||||
)
|
||||
ho_siam = (
|
||||
make_loader(bilat_holdout, slots_patient, image_transform=siamese.tower.transform, shuffle=False, **loader_kw)
|
||||
if bilat_holdout else None
|
||||
)
|
||||
|
||||
opt_base = torch.optim.Adam(baseline.parameters(), lr=args.lr)
|
||||
opt_siam = torch.optim.Adam(siamese.parameters(), lr=args.lr)
|
||||
|
||||
# ---- epoch log ----
|
||||
epoch_log_path = fold_dir / "epoch_log.csv"
|
||||
epoch_fields = [
|
||||
"fold", "epoch",
|
||||
"base_train_loss", "base_train_acc",
|
||||
"base_val_auc", "base_val_acc", "base_val_n",
|
||||
"siam_train_loss", "siam_train_acc",
|
||||
"siam_val_auc", "siam_val_acc", "siam_val_n",
|
||||
"ho_base_auc", "ho_base_acc", "ho_base_n",
|
||||
"ho_siam_auc", "ho_siam_acc", "ho_siam_n",
|
||||
]
|
||||
epoch_fp = epoch_log_path.open("w", newline="", encoding="utf-8")
|
||||
epoch_writer = csv.DictWriter(epoch_fp, fieldnames=epoch_fields)
|
||||
epoch_writer.writeheader()
|
||||
|
||||
def _f(v):
|
||||
return None if (v is None or (isinstance(v, float) and np.isnan(v))) else round(float(v), 6)
|
||||
|
||||
# ---- best-epoch tracking ----
|
||||
best_base_auc = -1.0
|
||||
best_siam_auc = -1.0
|
||||
best_base_state: Optional[dict] = None
|
||||
best_siam_state: Optional[dict] = None
|
||||
best_base_val_acc = _nan()
|
||||
best_siam_val_acc = _nan()
|
||||
# Holdout metrics snapped at best-val checkpoint
|
||||
snap_ho_base_auc = _nan()
|
||||
snap_ho_base_acc = _nan()
|
||||
snap_ho_base_n = 0
|
||||
snap_ho_siam_auc = _nan()
|
||||
snap_ho_siam_acc = _nan()
|
||||
snap_ho_siam_n = 0
|
||||
best_epoch = 0
|
||||
|
||||
print(
|
||||
f" [fold {fold+1}] training {args.epochs} epochs | "
|
||||
f"baseline n_train={len(od_train)} n_val={len(od_val)} | "
|
||||
f"siamese n_train={len(bilat_train)} n_val={len(bilat_val)}",
|
||||
flush=True,
|
||||
)
|
||||
|
||||
for epoch in range(args.epochs):
|
||||
bl_loss, bl_acc = train_baseline_epoch(baseline, train_base, opt_base, device)
|
||||
si_loss, si_acc = train_siamese_epoch(siamese, train_siam, opt_siam, device)
|
||||
|
||||
b_val_acc, b_val_auc, b_val_n = evaluate_baseline(baseline, val_base, device, num_classes)
|
||||
s_val_acc, s_val_auc, s_val_n = evaluate_siamese( siamese, val_siam, device, num_classes)
|
||||
|
||||
# Holdout at this epoch (always evaluated for logging, cheaply)
|
||||
hb_auc, hb_acc, hb_n = (_nan(), _nan(), 0)
|
||||
hs_auc, hs_acc, hs_n = (_nan(), _nan(), 0)
|
||||
if ho_base is not None:
|
||||
hb_acc, hb_auc, hb_n = evaluate_baseline(baseline, ho_base, device, num_classes)
|
||||
if ho_siam is not None:
|
||||
hs_acc, hs_auc, hs_n = evaluate_siamese(siamese, ho_siam, device, num_classes)
|
||||
|
||||
# Best-epoch tracking: snapshot state independently per model
|
||||
if not np.isnan(b_val_auc) and b_val_auc > best_base_auc:
|
||||
best_base_auc = b_val_auc
|
||||
best_base_val_acc = b_val_acc
|
||||
best_base_state = copy.deepcopy(baseline.state_dict())
|
||||
snap_ho_base_auc = hb_auc
|
||||
snap_ho_base_acc = hb_acc
|
||||
snap_ho_base_n = hb_n
|
||||
|
||||
if not np.isnan(s_val_auc) and s_val_auc > best_siam_auc:
|
||||
best_siam_auc = s_val_auc
|
||||
best_siam_val_acc = s_val_acc
|
||||
best_siam_state = copy.deepcopy(siamese.state_dict())
|
||||
snap_ho_siam_auc = hs_auc
|
||||
snap_ho_siam_acc = hs_acc
|
||||
snap_ho_siam_n = hs_n
|
||||
best_epoch = epoch + 1
|
||||
|
||||
row = {
|
||||
"fold": fold, "epoch": epoch + 1,
|
||||
"base_train_loss": _f(bl_loss), "base_train_acc": _f(bl_acc),
|
||||
"base_val_auc": _f(b_val_auc), "base_val_acc": _f(b_val_acc), "base_val_n": b_val_n,
|
||||
"siam_train_loss": _f(si_loss), "siam_train_acc": _f(si_acc),
|
||||
"siam_val_auc": _f(s_val_auc), "siam_val_acc": _f(s_val_acc), "siam_val_n": s_val_n,
|
||||
"ho_base_auc": _f(hb_auc), "ho_base_acc": _f(hb_acc), "ho_base_n": hb_n,
|
||||
"ho_siam_auc": _f(hs_auc), "ho_siam_acc": _f(hs_acc), "ho_siam_n": hs_n,
|
||||
}
|
||||
epoch_writer.writerow(row)
|
||||
epoch_fp.flush()
|
||||
|
||||
if args.log_every > 0 and (epoch + 1) % args.log_every == 0:
|
||||
print(
|
||||
f" ep {epoch+1:>3}/{args.epochs} "
|
||||
f"base val AUC={b_val_auc:.4f} siam val AUC={s_val_auc:.4f} "
|
||||
f"(best base={best_base_auc:.4f} best siam={best_siam_auc:.4f})",
|
||||
flush=True,
|
||||
)
|
||||
|
||||
epoch_fp.close()
|
||||
|
||||
# Save best checkpoints
|
||||
if best_base_state is not None:
|
||||
torch.save(best_base_state, fold_dir / "best_baseline.pt")
|
||||
if best_siam_state is not None:
|
||||
torch.save(best_siam_state, fold_dir / "best_siamese.pt")
|
||||
|
||||
result = FoldResult(
|
||||
mode=mode, fold=fold,
|
||||
best_epoch=best_epoch,
|
||||
baseline_best_val_auc=best_base_auc,
|
||||
baseline_best_val_acc=best_base_val_acc,
|
||||
siamese_best_val_auc=best_siam_auc,
|
||||
siamese_best_val_acc=best_siam_val_acc,
|
||||
baseline_holdout_auc=snap_ho_base_auc,
|
||||
baseline_holdout_acc=snap_ho_base_acc,
|
||||
baseline_holdout_n=snap_ho_base_n,
|
||||
siamese_holdout_auc=snap_ho_siam_auc,
|
||||
siamese_holdout_acc=snap_ho_siam_acc,
|
||||
siamese_holdout_n=snap_ho_siam_n,
|
||||
baseline_n=len(od_val),
|
||||
siamese_n=len(bilat_val),
|
||||
)
|
||||
|
||||
print(
|
||||
f" [fold {fold+1}] BEST "
|
||||
f"base val AUC={best_base_auc:.4f} acc={best_base_val_acc:.4f} "
|
||||
f"siam val AUC={best_siam_auc:.4f} acc={best_siam_val_acc:.4f} "
|
||||
f"(siam best epoch={best_epoch})",
|
||||
flush=True,
|
||||
)
|
||||
if snap_ho_base_n > 0 or snap_ho_siam_n > 0:
|
||||
print(
|
||||
f" [fold {fold+1}] HOUT "
|
||||
f"base AUC={snap_ho_base_auc:.4f} acc={snap_ho_base_acc:.4f} (n={snap_ho_base_n}) "
|
||||
f"siam AUC={snap_ho_siam_auc:.4f} acc={snap_ho_siam_acc:.4f} (n={snap_ho_siam_n})",
|
||||
flush=True,
|
||||
)
|
||||
|
||||
return result
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Summary helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _summary(results: list[FoldResult]) -> dict:
|
||||
def _means(vals):
|
||||
v = np.array([x for x in vals if not np.isnan(x)], dtype=float)
|
||||
return (float(np.mean(v)) if len(v) else None,
|
||||
float(np.std(v)) if len(v) else None)
|
||||
|
||||
b_val_aucs = [r.baseline_best_val_auc for r in results]
|
||||
s_val_aucs = [r.siamese_best_val_auc for r in results]
|
||||
b_ho_aucs = [r.baseline_holdout_auc for r in results]
|
||||
s_ho_aucs = [r.siamese_holdout_auc for r in results]
|
||||
b_val_accs = [r.baseline_best_val_acc for r in results]
|
||||
s_val_accs = [r.siamese_best_val_acc for r in results]
|
||||
b_ho_accs = [r.baseline_holdout_acc for r in results]
|
||||
s_ho_accs = [r.siamese_holdout_acc for r in results]
|
||||
|
||||
deltas_val_auc = [s - b for b, s in zip(b_val_aucs, s_val_aucs)
|
||||
if not np.isnan(b) and not np.isnan(s)]
|
||||
deltas_ho_auc = [s - b for b, s in zip(b_ho_aucs, s_ho_aucs)
|
||||
if not np.isnan(b) and not np.isnan(s)]
|
||||
|
||||
bva_m, bva_s = _means(b_val_aucs)
|
||||
sva_m, sva_s = _means(s_val_aucs)
|
||||
bha_m, bha_s = _means(b_ho_aucs)
|
||||
sha_m, sha_s = _means(s_ho_aucs)
|
||||
|
||||
return {
|
||||
"baseline_best_val": {"auc_mean": bva_m, "auc_std": bva_s, "acc_mean": _means(b_val_accs)[0]},
|
||||
"siamese_best_val": {"auc_mean": sva_m, "auc_std": sva_s, "acc_mean": _means(s_val_accs)[0]},
|
||||
"delta_val_auc": {"mean": float(np.mean(deltas_val_auc)) if deltas_val_auc else None,
|
||||
"std": float(np.std(deltas_val_auc)) if deltas_val_auc else None},
|
||||
"baseline_holdout": {"auc_mean": bha_m, "auc_std": bha_s, "acc_mean": _means(b_ho_accs)[0]},
|
||||
"siamese_holdout": {"auc_mean": sha_m, "auc_std": sha_s, "acc_mean": _means(s_ho_accs)[0]},
|
||||
"delta_holdout_auc": {"mean": float(np.mean(deltas_ho_auc)) if deltas_ho_auc else None,
|
||||
"std": float(np.std(deltas_ho_auc)) if deltas_ho_auc else None},
|
||||
}
|
||||
|
||||
|
||||
def _print_summary(mode: str, s: dict) -> None:
|
||||
def f(v):
|
||||
return "nan" if v is None else f"{v:.4f}"
|
||||
|
||||
bv = s["baseline_best_val"]
|
||||
sv = s["siamese_best_val"]
|
||||
dv = s["delta_val_auc"]
|
||||
bh = s["baseline_holdout"]
|
||||
sh = s["siamese_holdout"]
|
||||
dh = s["delta_holdout_auc"]
|
||||
|
||||
print(f"\n=== Summary [{mode}] (best-epoch metrics) ===")
|
||||
print(f" val baseline AUC={f(bv['auc_mean'])}±{f(bv['auc_std'])} acc={f(bv['acc_mean'])}")
|
||||
print(f" val siamese AUC={f(sv['auc_mean'])}±{f(sv['auc_std'])} acc={f(sv['acc_mean'])}")
|
||||
print(f" val delta AUC={f(dv['mean'])}±{f(dv['std'])}")
|
||||
print(f" hout baseline AUC={f(bh['auc_mean'])}±{f(bh['auc_std'])} acc={f(bh['acc_mean'])}")
|
||||
print(f" hout siamese AUC={f(sh['auc_mean'])}±{f(sh['auc_std'])} acc={f(sh['acc_mean'])}")
|
||||
print(f" hout delta AUC={f(dh['mean'])}±{f(dh['std'])}")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# CLI
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def parse_args():
|
||||
ap = argparse.ArgumentParser(
|
||||
description="Baseline single-eye vs SiameseImageTower bilateral comparison."
|
||||
)
|
||||
ap.add_argument("--image-dir", default="Papila/FundusImages")
|
||||
ap.add_argument("--clinical-dir", default="Papila/ClinicalData")
|
||||
ap.add_argument("--label-col", default="Diagnosis")
|
||||
ap.add_argument("--cat-cols", nargs="*", default=["Gender", "Phakic/Pseudophakic"])
|
||||
ap.add_argument("--eval-mode", choices=["binary", "multiclass"], default="binary")
|
||||
ap.add_argument(
|
||||
"--eval-modes", nargs="+", choices=["binary", "multiclass"], default=None,
|
||||
help="Run multiple modes in one pass, e.g. --eval-modes binary multiclass",
|
||||
)
|
||||
ap.add_argument("--n-splits", type=int, default=5)
|
||||
ap.add_argument("--fold-seed", type=int, default=42)
|
||||
ap.add_argument("--holdout-per-class", type=int, default=0)
|
||||
ap.add_argument("--holdout-seed", type=int, default=123)
|
||||
ap.add_argument("--folds", type=int, default=5)
|
||||
ap.add_argument("--epochs", type=int, default=40)
|
||||
ap.add_argument("--batch-size", type=int, default=8)
|
||||
ap.add_argument("--lr", type=float, default=1e-4)
|
||||
ap.add_argument("--backbone", default="refugelike")
|
||||
ap.add_argument("--freeze-ratio", type=float, default=0.0)
|
||||
ap.add_argument("--augment", action="store_true")
|
||||
ap.add_argument("--num-workers", type=int, default=0)
|
||||
ap.add_argument("--device", choices=["auto", "cpu", "cuda"], default="auto")
|
||||
ap.add_argument("--seed", type=int, default=1234)
|
||||
ap.add_argument("--run-name", default=None)
|
||||
ap.add_argument("--output-root", default="analysis_data/basic_analysis")
|
||||
ap.add_argument(
|
||||
"--exclude-mixed-patients", action="store_true",
|
||||
help="Drop patients whose two eyes have different labels before splitting.",
|
||||
)
|
||||
ap.add_argument(
|
||||
"--log-every", type=int, default=5,
|
||||
help="Print epoch progress every N epochs (0 to disable).",
|
||||
)
|
||||
return ap.parse_args()
|
||||
|
||||
|
||||
def choose_device(name: str) -> torch.device:
|
||||
if name == "cuda":
|
||||
if not torch.cuda.is_available():
|
||||
raise RuntimeError("--device cuda requested but CUDA is not available.")
|
||||
return torch.device("cuda")
|
||||
if name == "cpu":
|
||||
return torch.device("cpu")
|
||||
return torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Entry point
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
device = choose_device(args.device)
|
||||
seed_everything(args.seed)
|
||||
|
||||
print(f"Device: {device}", flush=True)
|
||||
print("Loading PAPILA data...", flush=True)
|
||||
data = build_papila_data(
|
||||
image_dir=args.image_dir,
|
||||
clinical_dir=args.clinical_dir,
|
||||
label_col=args.label_col,
|
||||
cat_cols=list(args.cat_cols),
|
||||
n_splits=args.n_splits,
|
||||
random_seed=args.fold_seed,
|
||||
)
|
||||
print(f"Loaded: {len(data.df)} rows", flush=True)
|
||||
|
||||
ts = time.strftime("%Y%m%d_%H%M%S")
|
||||
run_name = args.run_name or f"siamese_compare_{ts}"
|
||||
out_dir = Path(args.output_root) / run_name
|
||||
out_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
eval_modes = args.eval_modes if args.eval_modes else [args.eval_mode]
|
||||
|
||||
all_results: dict[str, list[FoldResult]] = {}
|
||||
summaries: dict[str, dict] = {}
|
||||
|
||||
for mode in eval_modes:
|
||||
import pandas as pd
|
||||
df_mode = data.df.copy()
|
||||
|
||||
if args.exclude_mixed_patients:
|
||||
before = df_mode["Patient ID"].nunique()
|
||||
df_mode, mixed = _drop_mixed_label_patients(
|
||||
df_mode, patient_col="Patient ID", label_col=args.label_col
|
||||
)
|
||||
print(f"[{mode}] dropped {len(mixed)} mixed-label patients "
|
||||
f"({before} -> {df_mode['Patient ID'].nunique()})", flush=True)
|
||||
|
||||
if mode == "binary":
|
||||
df_mode = df_mode[df_mode[args.label_col].isin([0, 1])].reset_index(drop=True)
|
||||
|
||||
num_classes = 2 if mode == "binary" else int(df_mode[args.label_col].nunique())
|
||||
print(
|
||||
f"\n[{mode}] num_classes={num_classes} rows={len(df_mode)} "
|
||||
f"patients={df_mode['Patient ID'].nunique()}",
|
||||
flush=True,
|
||||
)
|
||||
|
||||
split_manager = PatientFirstSplitManager(
|
||||
patient_col="Patient ID", label_col=args.label_col
|
||||
)
|
||||
split_args = SimpleNamespace(
|
||||
eval_mode=mode,
|
||||
holdout_per_class=args.holdout_per_class,
|
||||
holdout_seed=args.holdout_seed,
|
||||
n_splits=args.n_splits,
|
||||
fold_seed=args.fold_seed,
|
||||
)
|
||||
clinical_ns = SimpleNamespace(df=df_mode, label_col=args.label_col)
|
||||
plans = split_manager.build_plans(clinical=clinical_ns, args=split_args, profile=None)
|
||||
n_folds = min(args.folds, len(plans))
|
||||
|
||||
profile_od = build_papila_profile(
|
||||
patient_col="Patient ID", label_col=args.label_col, sample_mode="eye"
|
||||
)
|
||||
profile_patient = build_papila_profile(
|
||||
patient_col="Patient ID", label_col=args.label_col, sample_mode="patient"
|
||||
)
|
||||
|
||||
mode_dir = out_dir / mode
|
||||
mode_dir.mkdir(exist_ok=True)
|
||||
|
||||
fold_results: list[FoldResult] = []
|
||||
for fold in range(n_folds):
|
||||
fold_seed = args.seed + fold * 100
|
||||
seed_everything(fold_seed)
|
||||
fold_dir = mode_dir / f"fold{fold}"
|
||||
fold_dir.mkdir(exist_ok=True)
|
||||
|
||||
print(f"\n[{mode}] fold {fold+1}/{n_folds}", flush=True)
|
||||
result = run_fold(
|
||||
fold=fold,
|
||||
split=plans[fold],
|
||||
mode=mode,
|
||||
args=args,
|
||||
device=device,
|
||||
data=data,
|
||||
num_classes=num_classes,
|
||||
profile_od=profile_od,
|
||||
profile_patient=profile_patient,
|
||||
fold_dir=fold_dir,
|
||||
)
|
||||
fold_results.append(result)
|
||||
|
||||
# Write per-mode CSV
|
||||
fold_csv = out_dir / f"{mode}_fold_results.csv"
|
||||
csv_fields = list(FoldResult.__dataclass_fields__.keys())
|
||||
with fold_csv.open("w", newline="", encoding="utf-8") as fh:
|
||||
w = csv.DictWriter(fh, fieldnames=csv_fields)
|
||||
w.writeheader()
|
||||
for r in fold_results:
|
||||
w.writerow({k: getattr(r, k) for k in csv_fields})
|
||||
|
||||
summary = _summary(fold_results)
|
||||
_print_summary(mode, summary)
|
||||
|
||||
all_results[mode] = fold_results
|
||||
summaries[mode] = summary
|
||||
|
||||
payload = {
|
||||
"run_name": run_name,
|
||||
"timestamp": ts,
|
||||
"config": vars(args),
|
||||
"summaries": summaries,
|
||||
}
|
||||
(out_dir / "summary.json").write_text(json.dumps(payload, indent=2), encoding="utf-8")
|
||||
print(f"\nOutputs written to: {out_dir}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,117 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Rank multiclass runs by mean holdout AUC (fused) across folds."""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, Iterable, List, Optional
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
|
||||
|
||||
# ---------------------------
|
||||
# Config (edit in IDE)
|
||||
# ---------------------------
|
||||
ANALYSIS_DIR = Path("analysis_data/grid_search")
|
||||
TOP_N = 20
|
||||
HEAD = "fused" # fused | image | metadata
|
||||
OUTPUT_CSV = Path("analysis_data/grid_search/plots/best_holdout_multiclass.csv")
|
||||
|
||||
|
||||
def _read_json(path: Path) -> Optional[Dict[str, Any]]:
|
||||
if not path.exists():
|
||||
return None
|
||||
try:
|
||||
data = json.loads(path.read_text())
|
||||
except Exception:
|
||||
return None
|
||||
return data if isinstance(data, dict) else None
|
||||
|
||||
|
||||
def _infer_mode(summary: Optional[Dict[str, Any]]) -> Optional[str]:
|
||||
if not summary:
|
||||
return None
|
||||
eval_mode = summary.get("eval_mode")
|
||||
if isinstance(eval_mode, str):
|
||||
mode = eval_mode.strip().lower()
|
||||
if mode == "binary":
|
||||
return "binary"
|
||||
if mode in {"multiclass", "multi", "multi-class"}:
|
||||
return "multiclass"
|
||||
num_classes = summary.get("num_classes")
|
||||
if isinstance(num_classes, (int, float)):
|
||||
return "binary" if int(num_classes) <= 2 else "multiclass"
|
||||
class_names = summary.get("class_names")
|
||||
if isinstance(class_names, list) and class_names:
|
||||
return "binary" if len(class_names) <= 2 else "multiclass"
|
||||
return None
|
||||
|
||||
|
||||
def _simple_fields(summary: Dict[str, Any]) -> Dict[str, Any]:
|
||||
keep: Dict[str, Any] = {}
|
||||
for key, val in summary.items():
|
||||
if key == "fold_metrics":
|
||||
continue
|
||||
if isinstance(val, (str, int, float, bool)) or val is None:
|
||||
keep[key] = val
|
||||
return keep
|
||||
|
||||
|
||||
def _collect_fold_values(summary: Dict[str, Any], metric_key: str) -> List[float]:
|
||||
values: List[float] = []
|
||||
for entry in summary.get("fold_metrics") or []:
|
||||
if not isinstance(entry, dict):
|
||||
continue
|
||||
stats = entry.get("stats") if isinstance(entry.get("stats"), dict) else {}
|
||||
val = stats.get(metric_key)
|
||||
if isinstance(val, (int, float)):
|
||||
values.append(float(val))
|
||||
return values
|
||||
|
||||
|
||||
def main() -> None:
|
||||
metric_key = f"holdout_auc_{HEAD}"
|
||||
rows: List[Dict[str, Any]] = []
|
||||
|
||||
for run_dir in sorted(ANALYSIS_DIR.iterdir()):
|
||||
if not run_dir.is_dir():
|
||||
continue
|
||||
summary = _read_json(run_dir / "summary.json")
|
||||
mode = _infer_mode(summary)
|
||||
if mode != "multiclass":
|
||||
continue
|
||||
|
||||
values = _collect_fold_values(summary, metric_key)
|
||||
if not values:
|
||||
continue
|
||||
|
||||
mean_val = float(np.mean(values))
|
||||
std_val = float(np.std(values, ddof=1)) if len(values) > 1 else float("nan")
|
||||
|
||||
row = {
|
||||
"run_id": summary.get("run_id", run_dir.name),
|
||||
"run_dir": str(run_dir),
|
||||
"metric": metric_key,
|
||||
"mean": mean_val,
|
||||
"std": std_val,
|
||||
"n_folds": len(values),
|
||||
**_simple_fields(summary),
|
||||
}
|
||||
rows.append(row)
|
||||
|
||||
if not rows:
|
||||
raise SystemExit("No multiclass runs with holdout AUC found.")
|
||||
|
||||
df = pd.DataFrame(rows).sort_values(by="mean", ascending=False)
|
||||
top_df = df.head(TOP_N) if TOP_N else df
|
||||
|
||||
OUTPUT_CSV.parent.mkdir(parents=True, exist_ok=True)
|
||||
df.to_csv(OUTPUT_CSV, index=False)
|
||||
|
||||
print(top_df.to_string(index=False, float_format=lambda x: f"{x:.4f}"))
|
||||
print(f"\nSaved full ranking to: {OUTPUT_CSV}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,376 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Compare cached crop bounds/features vs GT-derived crops from the manifest."""
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
from typing import Dict, List, Optional, Tuple
|
||||
|
||||
import sys
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
from PIL import Image
|
||||
import torch
|
||||
from torchvision import transforms
|
||||
|
||||
REPO_ROOT = Path(__file__).resolve().parents[2]
|
||||
if str(REPO_ROOT) not in sys.path:
|
||||
sys.path.insert(0, str(REPO_ROOT))
|
||||
|
||||
from classes.hypertower import ManifestImageCropper, UNetImageCropper
|
||||
from classes.refuge_segmentation import UNet as RefugeUNet
|
||||
|
||||
|
||||
# ---------------------------
|
||||
# Config (edit in IDE)
|
||||
# ---------------------------
|
||||
CACHE_DIR = Path("analysis_data/hypertower_crops")
|
||||
MANIFEST_PATH = Path("manifest.csv")
|
||||
IMAGE_DIR = Path("Papila/FundusImages")
|
||||
SCALE = 2.5
|
||||
MAX_SAMPLES = 200 # set None to scan all
|
||||
TOL_BOUNDS = 1.0 # pixels
|
||||
TOL_FEATURES = 1e-3
|
||||
UNET_VARIANTS = [
|
||||
("norm_imagenet", Path("models/unet_segmenter/norm_imagenet/best.pt"), "imagenet"),
|
||||
("normalize_none", Path("models/unet_segmenter/normalize_none/best.pt"), "none"),
|
||||
("norm_per_image", Path("models/unet_segmenter/norm_per_image/best.pt"), "per_image"),
|
||||
]
|
||||
REFUGE_SEG_WEIGHTS = Path("models/refuge/segmentation/refuge_segmentation_best.pt")
|
||||
|
||||
|
||||
def _load_cache(path: Path) -> Optional[Dict[str, np.ndarray]]:
|
||||
try:
|
||||
data = np.load(path, allow_pickle=False)
|
||||
except Exception:
|
||||
return None
|
||||
return {k: data[k] for k in data.files}
|
||||
|
||||
|
||||
def _parse_stem(path: Path) -> str:
|
||||
# expects RET###OS_s250.npz -> RET###OS
|
||||
stem = path.stem
|
||||
if "_s" in stem:
|
||||
stem = stem.split("_s")[0]
|
||||
return stem
|
||||
|
||||
|
||||
def _image_path_from_stem(stem: str) -> Optional[Path]:
|
||||
cand = IMAGE_DIR / f"{stem}.jpg"
|
||||
if cand.exists():
|
||||
return cand
|
||||
cand = IMAGE_DIR / f"{stem}.png"
|
||||
if cand.exists():
|
||||
return cand
|
||||
return None
|
||||
|
||||
|
||||
def _gt_info(
|
||||
cropper: ManifestImageCropper, image_path: Path
|
||||
) -> Optional[Dict[str, float]]:
|
||||
try:
|
||||
image = Image.open(image_path).convert("RGB")
|
||||
except Exception:
|
||||
return None
|
||||
info = cropper._compute_crop_info(image, image_path)
|
||||
return info
|
||||
|
||||
|
||||
def _unet_info(
|
||||
cropper: UNetImageCropper, image_path: Path
|
||||
) -> Optional[Dict[str, float]]:
|
||||
try:
|
||||
image = Image.open(image_path).convert("RGB")
|
||||
except Exception:
|
||||
return None
|
||||
info = cropper._compute_crop_info(image, image_path)
|
||||
return info
|
||||
|
||||
|
||||
def _load_refuge_model(device: str) -> Optional[RefugeUNet]:
|
||||
if not REFUGE_SEG_WEIGHTS.exists():
|
||||
return None
|
||||
model = RefugeUNet()
|
||||
try:
|
||||
state = torch.load(REFUGE_SEG_WEIGHTS, map_location=device)
|
||||
except Exception:
|
||||
return None
|
||||
state_dict = state.get("model", state) if isinstance(state, dict) else state
|
||||
try:
|
||||
model.load_state_dict(state_dict)
|
||||
except Exception:
|
||||
return None
|
||||
model.to(device)
|
||||
model.eval()
|
||||
return model
|
||||
|
||||
|
||||
def _refuge_seg_info(
|
||||
model: RefugeUNet, device: str, image_path: Path
|
||||
) -> Optional[Dict[str, float]]:
|
||||
try:
|
||||
image = Image.open(image_path).convert("RGB")
|
||||
except Exception:
|
||||
return None
|
||||
original_size = image.size
|
||||
image_resized = image.resize((512, 512), Image.BILINEAR)
|
||||
tensor = transforms.ToTensor()(image_resized).unsqueeze(0).to(device)
|
||||
with torch.no_grad():
|
||||
logits = model(tensor)
|
||||
mask = torch.sigmoid(logits)[0, 0]
|
||||
mask_np = (mask.cpu().numpy() > 0.5).astype(np.float32)
|
||||
mask_img = Image.fromarray(mask_np)
|
||||
mask_img = mask_img.resize(original_size, Image.NEAREST)
|
||||
mask_np = np.array(mask_img, dtype=np.float32)
|
||||
coords = np.argwhere(mask_np > 0.5)
|
||||
if coords.size == 0:
|
||||
return None
|
||||
ys, xs = coords[:, 0], coords[:, 1]
|
||||
centre_x = float(xs.mean())
|
||||
centre_y = float(ys.mean())
|
||||
width = float(xs.max() - xs.min())
|
||||
height = float(ys.max() - ys.min())
|
||||
diameter = max(width, height)
|
||||
radius = diameter / 2.0
|
||||
crop_radius = radius * SCALE
|
||||
left = max(0.0, centre_x - crop_radius)
|
||||
upper = max(0.0, centre_y - crop_radius)
|
||||
right = min(float(image.width), centre_x + crop_radius)
|
||||
lower = min(float(image.height), centre_y + crop_radius)
|
||||
return {
|
||||
"left": left,
|
||||
"upper": upper,
|
||||
"right": right,
|
||||
"lower": lower,
|
||||
}
|
||||
|
||||
|
||||
def _diff_bounds(cache: Dict[str, np.ndarray], gt: Dict[str, float]) -> Optional[float]:
|
||||
keys = ("left", "upper", "right", "lower")
|
||||
if not all(k in cache for k in keys):
|
||||
return None
|
||||
diffs = [abs(float(cache[k]) - float(gt[k])) for k in keys]
|
||||
return float(max(diffs))
|
||||
|
||||
|
||||
def _diff_features(
|
||||
cache: Dict[str, np.ndarray], gt: Dict[str, float]
|
||||
) -> Optional[float]:
|
||||
if "features" not in cache or "features" not in gt:
|
||||
return None
|
||||
cf = np.asarray(cache["features"], dtype=float).ravel()
|
||||
gf = np.asarray(gt["features"], dtype=float).ravel()
|
||||
if cf.shape != gf.shape:
|
||||
return None
|
||||
return float(np.max(np.abs(cf - gf)))
|
||||
|
||||
|
||||
def main() -> None:
|
||||
if not CACHE_DIR.exists():
|
||||
raise SystemExit(f"Cache dir not found: {CACHE_DIR}")
|
||||
if not MANIFEST_PATH.exists():
|
||||
raise SystemExit(f"Manifest not found: {MANIFEST_PATH}")
|
||||
|
||||
cache_files = sorted(CACHE_DIR.glob(f"*_s{int(SCALE * 100)}.npz"))
|
||||
if MAX_SAMPLES is not None:
|
||||
cache_files = cache_files[:MAX_SAMPLES]
|
||||
print(f"[debug] cache files found: {len(cache_files)}")
|
||||
|
||||
try:
|
||||
manifest_df = pd.read_csv(MANIFEST_PATH)
|
||||
except Exception as exc:
|
||||
raise SystemExit(f"Failed to read manifest: {exc}")
|
||||
manifest_images = manifest_df.get("image_path")
|
||||
if manifest_images is None:
|
||||
raise SystemExit("Manifest is missing image_path column.")
|
||||
manifest_images = manifest_images.dropna().astype(str)
|
||||
manifest_stems = {Path(p).stem for p in manifest_images}
|
||||
print(f"[debug] manifest image_path count: {len(manifest_images)}")
|
||||
print(f"[debug] manifest unique stems: {len(manifest_stems)}")
|
||||
|
||||
cache_stems = {_parse_stem(p) for p in cache_files}
|
||||
overlap = cache_stems & manifest_stems
|
||||
print(
|
||||
f"[debug] cache stems: {len(cache_stems)} overlap with manifest stems: {len(overlap)}"
|
||||
)
|
||||
if cache_files:
|
||||
print(f"[debug] example cache stems: {sorted(list(cache_stems))[:5]}")
|
||||
if manifest_stems:
|
||||
print(f"[debug] example manifest stems: {sorted(list(manifest_stems))[:5]}")
|
||||
|
||||
cropper = ManifestImageCropper(
|
||||
manifest_path=MANIFEST_PATH,
|
||||
scale=SCALE,
|
||||
target_size=224,
|
||||
cache_dir=None,
|
||||
)
|
||||
|
||||
rows: List[Dict[str, object]] = []
|
||||
for cache_path in cache_files:
|
||||
cache = _load_cache(cache_path)
|
||||
if cache is None:
|
||||
continue
|
||||
stem = _parse_stem(cache_path)
|
||||
image_path = _image_path_from_stem(stem)
|
||||
if image_path is None:
|
||||
continue
|
||||
|
||||
gt = _gt_info(cropper, image_path)
|
||||
if gt is None:
|
||||
continue
|
||||
|
||||
bounds_diff = _diff_bounds(cache, gt)
|
||||
feat_diff = _diff_features(cache, gt)
|
||||
|
||||
rows.append(
|
||||
{
|
||||
"file": cache_path.name,
|
||||
"bounds_diff": bounds_diff,
|
||||
"features_diff": feat_diff,
|
||||
"bounds_match": bounds_diff is not None and bounds_diff <= TOL_BOUNDS,
|
||||
"features_match": feat_diff is not None and feat_diff <= TOL_FEATURES,
|
||||
}
|
||||
)
|
||||
|
||||
if not rows:
|
||||
print("[warn] No cache entries matched GT manifest entries.")
|
||||
else:
|
||||
df = pd.DataFrame(rows)
|
||||
print(df.head(10).to_string(index=False))
|
||||
print("\nSummary:")
|
||||
print(df[["bounds_diff", "features_diff"]].describe().to_string())
|
||||
if df["bounds_match"].notna().any():
|
||||
match_rate = df["bounds_match"].mean()
|
||||
print(f"\nBounds match rate (<= {TOL_BOUNDS}px): {match_rate:.3f}")
|
||||
if df["features_match"].notna().any():
|
||||
match_rate = df["features_match"].mean()
|
||||
print(f"Features match rate (<= {TOL_FEATURES}): {match_rate:.3f}")
|
||||
|
||||
print("\nUNet variant comparisons (no cache writes):")
|
||||
for name, weights, normalize in UNET_VARIANTS:
|
||||
if not weights.exists():
|
||||
print(f"[warn] {name}: weights not found at {weights}")
|
||||
continue
|
||||
|
||||
unet = UNetImageCropper(
|
||||
manifest_path=MANIFEST_PATH,
|
||||
weights_path=weights,
|
||||
normalize=normalize,
|
||||
threshold=0.5,
|
||||
tta=False,
|
||||
scale=SCALE,
|
||||
target_size=224,
|
||||
cache_dir=None, # ensure no cache writes
|
||||
)
|
||||
|
||||
u_rows: List[Dict[str, object]] = []
|
||||
missing_images = 0
|
||||
unet_none = 0
|
||||
cache_missing = 0
|
||||
exceptions = 0
|
||||
for cache_path in cache_files:
|
||||
cache = _load_cache(cache_path)
|
||||
if cache is None:
|
||||
cache_missing += 1
|
||||
continue
|
||||
stem = _parse_stem(cache_path)
|
||||
image_path = _image_path_from_stem(stem)
|
||||
if image_path is None:
|
||||
missing_images += 1
|
||||
continue
|
||||
try:
|
||||
info = _unet_info(unet, image_path)
|
||||
except Exception:
|
||||
exceptions += 1
|
||||
continue
|
||||
if info is None:
|
||||
unet_none += 1
|
||||
continue
|
||||
bounds_diff = _diff_bounds(cache, info)
|
||||
feat_diff = _diff_features(cache, info)
|
||||
u_rows.append(
|
||||
{
|
||||
"bounds_diff": bounds_diff,
|
||||
"features_diff": feat_diff,
|
||||
"bounds_match": bounds_diff is not None
|
||||
and bounds_diff <= TOL_BOUNDS,
|
||||
"features_match": feat_diff is not None
|
||||
and feat_diff <= TOL_FEATURES,
|
||||
}
|
||||
)
|
||||
|
||||
if not u_rows:
|
||||
print(
|
||||
f"[warn] {name}: no comparisons computed "
|
||||
f"(cache_missing={cache_missing}, missing_images={missing_images}, "
|
||||
f"unet_none={unet_none}, exceptions={exceptions})"
|
||||
)
|
||||
continue
|
||||
u_df = pd.DataFrame(u_rows)
|
||||
b_mean = float(u_df["bounds_diff"].mean())
|
||||
f_mean = float(u_df["features_diff"].mean())
|
||||
b_match = float(u_df["bounds_match"].mean())
|
||||
f_match = float(u_df["features_match"].mean())
|
||||
|
||||
print(
|
||||
f"{name}: mean bounds diff={b_mean:.3f}, mean feat diff={f_mean:.6f}, "
|
||||
f"bounds match rate={b_match:.3f}, features match rate={f_match:.3f}"
|
||||
)
|
||||
|
||||
print("\nRefuge segmentation model comparison (bounds only, no cache writes):")
|
||||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
refuge_model = _load_refuge_model(device)
|
||||
if refuge_model is None:
|
||||
print(f"[warn] refuge_segmentation_best.pt not found or failed to load at {REFUGE_SEG_WEIGHTS}")
|
||||
return
|
||||
|
||||
r_rows: List[Dict[str, object]] = []
|
||||
missing_images = 0
|
||||
cache_missing = 0
|
||||
model_none = 0
|
||||
exceptions = 0
|
||||
for cache_path in cache_files:
|
||||
cache = _load_cache(cache_path)
|
||||
if cache is None:
|
||||
cache_missing += 1
|
||||
continue
|
||||
stem = _parse_stem(cache_path)
|
||||
image_path = _image_path_from_stem(stem)
|
||||
if image_path is None:
|
||||
missing_images += 1
|
||||
continue
|
||||
try:
|
||||
info = _refuge_seg_info(refuge_model, device, image_path)
|
||||
except Exception:
|
||||
exceptions += 1
|
||||
continue
|
||||
if info is None:
|
||||
model_none += 1
|
||||
continue
|
||||
bounds_diff = _diff_bounds(cache, info)
|
||||
r_rows.append(
|
||||
{
|
||||
"bounds_diff": bounds_diff,
|
||||
"bounds_match": bounds_diff is not None and bounds_diff <= TOL_BOUNDS,
|
||||
}
|
||||
)
|
||||
|
||||
if not r_rows:
|
||||
print(
|
||||
"[warn] refuge_segmentation_best: no comparisons computed "
|
||||
f"(cache_missing={cache_missing}, missing_images={missing_images}, "
|
||||
f"model_none={model_none}, exceptions={exceptions})"
|
||||
)
|
||||
return
|
||||
|
||||
r_df = pd.DataFrame(r_rows)
|
||||
b_mean = float(r_df["bounds_diff"].mean())
|
||||
b_match = float(r_df["bounds_match"].mean())
|
||||
print(
|
||||
f"refuge_segmentation_best: mean bounds diff={b_mean:.3f}, "
|
||||
f"bounds match rate={b_match:.3f}"
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,249 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Compare suspect AUC from image tower vs crop-derived geometry (CDR)."""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import Dict, List, Optional, Tuple
|
||||
|
||||
import sys
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
from PIL import Image
|
||||
from sklearn.metrics import roc_auc_score
|
||||
|
||||
REPO_ROOT = Path(__file__).resolve().parents[2]
|
||||
if str(REPO_ROOT) not in sys.path:
|
||||
sys.path.insert(0, str(REPO_ROOT))
|
||||
|
||||
from classes import build_papila_clinical
|
||||
from classes.hypertower import UNetImageCropper, ManifestImageCropper
|
||||
|
||||
|
||||
# ---------------------------
|
||||
# Config (edit in IDE)
|
||||
# ---------------------------
|
||||
RUN_DIRS = [
|
||||
Path("analysis_data/1030_Balanced_Unet_Perimg_Resnet_SE16NormB_SE16NormT_multi_fused/1030_Balanced_Unet_Perimg_Resnet_SE16NormB_SE16NormT_multi_fused_20251030_091842"),
|
||||
Path("analysis_data/1030_Balanced_GT_Perimg_Resnet_SE16NormB_SE16NormT_multi_fused/1030_Balanced_GT_Perimg_Resnet_SE16NormB_SE16NormT_multi_fused_20251030_113730"),
|
||||
]
|
||||
GEOM_CACHE_ROOT = Path("analysis_data/geometry_cache")
|
||||
SUSPECT_LABEL = 2
|
||||
|
||||
|
||||
def _load_json(path: Path) -> Dict:
|
||||
if not path.exists():
|
||||
return {}
|
||||
try:
|
||||
return json.loads(path.read_text())
|
||||
except Exception:
|
||||
return {}
|
||||
|
||||
|
||||
def _drop_holdout_rows(clinical, holdout_path: Path) -> None:
|
||||
if not holdout_path.exists():
|
||||
return
|
||||
holdout = pd.read_csv(holdout_path)
|
||||
if holdout.empty:
|
||||
return
|
||||
if "Patient ID" not in holdout.columns or "eyeID" not in holdout.columns:
|
||||
return
|
||||
holdout_keys = set(zip(holdout["Patient ID"].astype(int), holdout["eyeID"].astype(str)))
|
||||
df = clinical.df.copy()
|
||||
df["_key"] = list(zip(df["Patient ID"].astype(int), df["eyeID"].astype(str)))
|
||||
df = df[~df["_key"].isin(holdout_keys)].drop(columns=["_key"]).reset_index(drop=True)
|
||||
|
||||
clinical.frames = [df.copy()]
|
||||
clinical.df = df.copy()
|
||||
clinical._infer_or_validate_feature_types()
|
||||
clinical._compute_numeric_stats()
|
||||
clinical._build_cat_maps()
|
||||
clinical._compute_feature_dim()
|
||||
clinical._build_kfold_indices()
|
||||
|
||||
|
||||
def _make_cropper(args: Dict, cache_dir: Path):
|
||||
manifest = args.get("img_crop_manifest")
|
||||
if not manifest:
|
||||
raise RuntimeError("img_crop_manifest missing; cannot compute geometry features.")
|
||||
|
||||
scale = float(args.get("img_crop_scale", 2.5))
|
||||
target_size = int(args.get("img_crop_size", 224))
|
||||
use_gt = bool(args.get("img_crop_gt", False))
|
||||
|
||||
if use_gt:
|
||||
return ManifestImageCropper(
|
||||
manifest_path=Path(manifest),
|
||||
scale=scale,
|
||||
target_size=target_size,
|
||||
cache_dir=cache_dir,
|
||||
)
|
||||
|
||||
weights = args.get("img_crop_weights")
|
||||
if not weights:
|
||||
raise RuntimeError("img_crop_weights missing for UNet cropper.")
|
||||
|
||||
normalize = args.get("img_crop_normalize", "per_image")
|
||||
threshold = float(args.get("img_crop_threshold", 0.5))
|
||||
tta = bool(args.get("img_crop_tta", False))
|
||||
return UNetImageCropper(
|
||||
manifest_path=Path(manifest),
|
||||
weights_path=Path(weights),
|
||||
normalize=normalize,
|
||||
threshold=threshold,
|
||||
tta=tta,
|
||||
scale=scale,
|
||||
target_size=target_size,
|
||||
cache_dir=cache_dir,
|
||||
)
|
||||
|
||||
|
||||
def _geometry_scores(
|
||||
clinical,
|
||||
cropper,
|
||||
test_df: pd.DataFrame,
|
||||
) -> Tuple[np.ndarray, np.ndarray]:
|
||||
scores: List[float] = []
|
||||
keep_mask: List[bool] = []
|
||||
for _, row in test_df.iterrows():
|
||||
img_path = clinical.get_image_path(row)
|
||||
try:
|
||||
image = Image.open(img_path).convert("RGB")
|
||||
except Exception:
|
||||
scores.append(float("nan"))
|
||||
keep_mask.append(False)
|
||||
continue
|
||||
feats = cropper.geometry_features(image, img_path)
|
||||
if feats is None or len(feats) == 0:
|
||||
scores.append(float("nan"))
|
||||
keep_mask.append(False)
|
||||
else:
|
||||
scores.append(float(feats[0])) # area_ratio (CDR)
|
||||
keep_mask.append(True)
|
||||
return np.asarray(scores, dtype=float), np.asarray(keep_mask, dtype=bool)
|
||||
|
||||
|
||||
def _suspect_auc(y_true: np.ndarray, scores: np.ndarray) -> float:
|
||||
y = (y_true == SUSPECT_LABEL).astype(int)
|
||||
if y.sum() == 0 or y.sum() == len(y):
|
||||
return float("nan")
|
||||
return float(roc_auc_score(y, scores))
|
||||
|
||||
|
||||
def main() -> None:
|
||||
rows: List[Dict[str, object]] = []
|
||||
|
||||
for run_dir in RUN_DIRS:
|
||||
cli_path = run_dir / "cli_args.json"
|
||||
cli_args = _load_json(cli_path)
|
||||
if not cli_args:
|
||||
print(f"[warn] Missing cli_args.json in {run_dir}")
|
||||
continue
|
||||
|
||||
label_col = cli_args.get("label_col", "Diagnosis")
|
||||
cat_cols = cli_args.get("cat_cols", ["Gender", "Phakic/Pseudophakic"])
|
||||
n_splits = int(cli_args.get("n_splits", 5))
|
||||
fold_seed = int(cli_args.get("fold_seed", 42))
|
||||
eval_mode = str(cli_args.get("eval_mode", "multiclass")).lower()
|
||||
|
||||
clinical = build_papila_clinical(
|
||||
image_dir=cli_args.get("image_dir", "Papila/FundusImages"),
|
||||
clinical_dir=cli_args.get("clinical_dir", "Papila/ClinicalData"),
|
||||
label_col=label_col,
|
||||
cat_cols=cat_cols,
|
||||
n_splits=n_splits,
|
||||
random_seed=fold_seed,
|
||||
)
|
||||
|
||||
if eval_mode == "binary":
|
||||
clinical.df = clinical.df[clinical.df[label_col].isin([0, 1])].reset_index(drop=True)
|
||||
clinical.frames = [clinical.df.copy()]
|
||||
clinical._infer_or_validate_feature_types()
|
||||
clinical._compute_numeric_stats()
|
||||
clinical._build_cat_maps()
|
||||
clinical._compute_feature_dim()
|
||||
clinical._build_kfold_indices()
|
||||
|
||||
_drop_holdout_rows(clinical, run_dir / "holdout.csv")
|
||||
|
||||
cache_dir = GEOM_CACHE_ROOT / run_dir.name
|
||||
cache_dir.mkdir(parents=True, exist_ok=True)
|
||||
cropper = _make_cropper(cli_args, cache_dir=cache_dir)
|
||||
|
||||
all_geom_scores: List[float] = []
|
||||
all_img_scores: List[float] = []
|
||||
all_y: List[int] = []
|
||||
|
||||
for fold in range(n_splits):
|
||||
y_path = run_dir / f"fold{fold}_y_true.npy"
|
||||
p_img_path = run_dir / f"fold{fold}_probs_img.npy"
|
||||
if not y_path.exists() or not p_img_path.exists():
|
||||
continue
|
||||
|
||||
y_true = np.load(y_path)
|
||||
probs_img = np.load(p_img_path)
|
||||
if probs_img.ndim != 2 or probs_img.shape[1] <= SUSPECT_LABEL:
|
||||
continue
|
||||
|
||||
_, test_df = clinical.get_split_dfs(fold)
|
||||
if len(test_df) != len(y_true):
|
||||
print(
|
||||
f"[warn] {run_dir.name} fold{fold}: test_df len {len(test_df)} != y_true len {len(y_true)}"
|
||||
)
|
||||
|
||||
geom_scores, keep_mask = _geometry_scores(clinical, cropper, test_df)
|
||||
if keep_mask.sum() == 0:
|
||||
print(f"[warn] {run_dir.name} fold{fold}: no valid geometry features")
|
||||
continue
|
||||
|
||||
y_fold = y_true[: len(geom_scores)][keep_mask]
|
||||
geom_fold = geom_scores[keep_mask]
|
||||
img_fold = probs_img[: len(geom_scores), SUSPECT_LABEL][keep_mask]
|
||||
|
||||
geom_auc = _suspect_auc(y_fold, geom_fold)
|
||||
img_auc = _suspect_auc(y_fold, img_fold)
|
||||
|
||||
rows.append(
|
||||
{
|
||||
"run": run_dir.name,
|
||||
"fold": fold,
|
||||
"metric": "suspect_auc",
|
||||
"image_auc": img_auc,
|
||||
"geometry_auc": geom_auc,
|
||||
"n": int(len(y_fold)),
|
||||
}
|
||||
)
|
||||
|
||||
all_geom_scores.append(geom_fold)
|
||||
all_img_scores.append(img_fold)
|
||||
all_y.append(y_fold)
|
||||
|
||||
if all_y:
|
||||
y_all = np.concatenate(all_y)
|
||||
geom_all = np.concatenate(all_geom_scores)
|
||||
img_all = np.concatenate(all_img_scores)
|
||||
rows.append(
|
||||
{
|
||||
"run": run_dir.name,
|
||||
"fold": "all",
|
||||
"metric": "suspect_auc",
|
||||
"image_auc": _suspect_auc(y_all, img_all),
|
||||
"geometry_auc": _suspect_auc(y_all, geom_all),
|
||||
"n": int(len(y_all)),
|
||||
}
|
||||
)
|
||||
|
||||
if not rows:
|
||||
raise SystemExit("No results produced; check run paths and files.")
|
||||
|
||||
df = pd.DataFrame(rows)
|
||||
out_path = GEOM_CACHE_ROOT / "suspect_auc_geometry_vs_image.csv"
|
||||
out_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
df.to_csv(out_path, index=False)
|
||||
print(df.to_string(index=False, float_format=lambda x: f"{x:.4f}"))
|
||||
print(f"\nSaved: {out_path}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,314 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Shared helpers for building grid search analytics.
|
||||
|
||||
The class below will gradually accumulate reusable utilities for working with
|
||||
grid search outputs (summary.json, cli_args.json, etc).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import csv
|
||||
import re
|
||||
from pathlib import Path
|
||||
from typing import Dict, Iterable, Iterator, List, Optional
|
||||
|
||||
import numpy as np
|
||||
|
||||
DEFAULT_EXCLUDE_KEYS = {
|
||||
"run_id",
|
||||
"fold_metrics",
|
||||
"best_metric",
|
||||
"best_metric_mode",
|
||||
"best_metric_mean",
|
||||
"best_metric_std",
|
||||
"eval_mode",
|
||||
"n_splits",
|
||||
"num_classes",
|
||||
}
|
||||
|
||||
|
||||
class GridSearchAnalytics:
|
||||
"""Utility wrapper for inspecting grid search result directories."""
|
||||
|
||||
def __init__(self,
|
||||
analysis_dir: Path | str,
|
||||
exclude_keys: Optional[Iterable[str]] = None) -> None:
|
||||
self.analysis_dir = Path(analysis_dir)
|
||||
if not self.analysis_dir.exists():
|
||||
raise FileNotFoundError(f"analysis_dir does not exist: {self.analysis_dir}")
|
||||
self.exclude_keys = set(exclude_keys or DEFAULT_EXCLUDE_KEYS)
|
||||
|
||||
def iter_run_dirs(self, shallow: bool = True) -> Iterator[Path]:
|
||||
"""
|
||||
Yield run directories containing grid search artifacts.
|
||||
|
||||
Shallow iteration only walks direct children. Deep iteration scans the
|
||||
entire subtree.
|
||||
"""
|
||||
candidates: Iterable[Path]
|
||||
if shallow:
|
||||
candidates = (p for p in sorted(self.analysis_dir.iterdir()) if p.is_dir())
|
||||
else:
|
||||
candidates = (p for p in self.analysis_dir.rglob("*") if p.is_dir())
|
||||
for run_dir in candidates:
|
||||
summary = run_dir / "summary.json"
|
||||
cli = run_dir / "cli_args.json"
|
||||
if summary.exists() or cli.exists():
|
||||
yield run_dir
|
||||
|
||||
def read_summary(self, run_dir: Path) -> Optional[Dict[str, object]]:
|
||||
"""Load summary.json for a run directory."""
|
||||
return self._read_json(run_dir / "summary.json")
|
||||
|
||||
def read_cli_args(self, run_dir: Path) -> Optional[Dict[str, object]]:
|
||||
"""Load cli_args.json for a run directory."""
|
||||
return self._read_json(run_dir / "cli_args.json")
|
||||
|
||||
def read_run_id(self,
|
||||
run_dir: Path,
|
||||
summary: Optional[Dict[str, object]]) -> str:
|
||||
if summary:
|
||||
rid = summary.get("run_id")
|
||||
if isinstance(rid, str) and rid:
|
||||
return rid
|
||||
return run_dir.name
|
||||
|
||||
def task_from_summary(self, summary: Optional[Dict[str, object]]) -> Optional[str]:
|
||||
"""Infer task (binary vs multiclass) from a summary payload."""
|
||||
if not summary:
|
||||
return None
|
||||
eval_mode = summary.get("eval_mode")
|
||||
if isinstance(eval_mode, str):
|
||||
mode = eval_mode.strip().lower()
|
||||
if mode == "binary":
|
||||
return "binary"
|
||||
if mode in {"multiclass", "multi", "multi-class"}:
|
||||
return "multiclass"
|
||||
num_classes = summary.get("num_classes")
|
||||
if isinstance(num_classes, (int, float)):
|
||||
return "binary" if int(num_classes) <= 2 else "multiclass"
|
||||
return None
|
||||
|
||||
def flatten_config(self,
|
||||
data: Dict[str, object],
|
||||
prefix: str = "",
|
||||
exclude_keys: Optional[Iterable[str]] = None) -> Dict[str, object]:
|
||||
"""Flatten nested CLI args or config dictionaries for analysis."""
|
||||
out: Dict[str, object] = {}
|
||||
excludes = set(exclude_keys or self.exclude_keys)
|
||||
for key, value in data.items():
|
||||
if key in excludes or key.startswith("best_"):
|
||||
continue
|
||||
full_key = f"{prefix}{key}" if not prefix else f"{prefix}.{key}"
|
||||
if isinstance(value, dict):
|
||||
out.update(self.flatten_config(value, full_key, exclude_keys=excludes))
|
||||
continue
|
||||
if isinstance(value, list):
|
||||
continue
|
||||
out[full_key] = value
|
||||
return out
|
||||
|
||||
def fusion_correction_events(self,
|
||||
output_csv: Path | str | None = None,
|
||||
shallow: bool = True) -> Path:
|
||||
"""
|
||||
Build a table of cases where the fused head is correct while both towers
|
||||
are wrong. Rows are written to CSV for downstream analysis.
|
||||
"""
|
||||
output_path = Path(output_csv) if output_csv else Path("analysis_data/grid_search_analytics/fusion_corrections.csv")
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
rows: List[Dict[str, object]] = []
|
||||
for run_dir in self.iter_run_dirs(shallow=shallow):
|
||||
summary = self.read_summary(run_dir)
|
||||
run_id = self.read_run_id(run_dir, summary)
|
||||
folds = self._available_folds(run_dir, summary)
|
||||
for fold in folds:
|
||||
y_true = self._load_y_true(run_dir, fold)
|
||||
if y_true is None:
|
||||
continue
|
||||
epoch_prob_paths = self._collect_epoch_prob_paths(run_dir, fold)
|
||||
if not epoch_prob_paths:
|
||||
# Per-epoch dumps were not found; fall back to the saved fold-level probabilities.
|
||||
base_paths = self._collect_base_prob_paths(run_dir, fold)
|
||||
if base_paths:
|
||||
epoch_hint = self._fold_epoch_hint(summary, fold)
|
||||
epoch_prob_paths = {epoch_hint if epoch_hint is not None else 0: base_paths}
|
||||
for epoch, paths in epoch_prob_paths.items():
|
||||
arrays = {head: self._load_probs_array(path) for head, path in paths.items()}
|
||||
if not self._has_all_heads(arrays):
|
||||
continue
|
||||
events = self._fusion_corrections_for_probs(y_true, arrays, run_id, fold, epoch)
|
||||
rows.extend(events)
|
||||
|
||||
if rows:
|
||||
fieldnames = [
|
||||
"run_id",
|
||||
"fold",
|
||||
"epoch",
|
||||
"index",
|
||||
"y_true",
|
||||
"pred_fused",
|
||||
"pred_img",
|
||||
"pred_md",
|
||||
"conf_fused",
|
||||
"conf_img",
|
||||
"conf_md",
|
||||
]
|
||||
with output_path.open("w", newline="") as f:
|
||||
writer = csv.DictWriter(f, fieldnames=fieldnames)
|
||||
writer.writeheader()
|
||||
writer.writerows(rows)
|
||||
else:
|
||||
output_path.write_text("")
|
||||
return output_path
|
||||
|
||||
@staticmethod
|
||||
def _read_json(path: Path) -> Optional[Dict[str, object]]:
|
||||
if not path.exists():
|
||||
return None
|
||||
try:
|
||||
data = json.loads(path.read_text())
|
||||
except Exception:
|
||||
return None
|
||||
if not isinstance(data, dict):
|
||||
return None
|
||||
return data
|
||||
|
||||
@staticmethod
|
||||
def _fold_epoch_hint(summary: Optional[Dict[str, object]], fold: int) -> Optional[int]:
|
||||
if not summary:
|
||||
return None
|
||||
fold_metrics = summary.get("fold_metrics") or []
|
||||
for entry in fold_metrics:
|
||||
if not isinstance(entry, dict):
|
||||
continue
|
||||
if entry.get("fold") == fold:
|
||||
stats = entry.get("stats") if isinstance(entry.get("stats"), dict) else {}
|
||||
epoch = stats.get("epoch") or entry.get("best_epoch")
|
||||
if isinstance(epoch, (int, float)):
|
||||
return int(epoch)
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _available_folds(run_dir: Path, summary: Optional[Dict[str, object]]) -> List[int]:
|
||||
folds: List[int] = []
|
||||
if summary:
|
||||
for entry in summary.get("fold_metrics") or []:
|
||||
if not isinstance(entry, dict):
|
||||
continue
|
||||
fold_idx = entry.get("fold")
|
||||
if isinstance(fold_idx, int):
|
||||
folds.append(fold_idx)
|
||||
if not folds:
|
||||
pattern = re.compile(r"fold(\d+)_y_true\.npy$")
|
||||
for path in run_dir.glob("fold*_y_true.npy"):
|
||||
match = pattern.match(path.name)
|
||||
if match:
|
||||
folds.append(int(match.group(1)))
|
||||
return sorted(set(folds))
|
||||
|
||||
@staticmethod
|
||||
def _load_y_true(run_dir: Path, fold: int) -> Optional[np.ndarray]:
|
||||
path = run_dir / f"fold{fold}_y_true.npy"
|
||||
if not path.exists():
|
||||
return None
|
||||
try:
|
||||
return np.load(path)
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _collect_epoch_prob_paths(run_dir: Path, fold: int) -> Dict[int, Dict[str, Path]]:
|
||||
pattern = re.compile(rf"fold{fold}_epoch(\d+)_probs_(\w+)\.npy$")
|
||||
epoch_paths: Dict[int, Dict[str, Path]] = {}
|
||||
for path in run_dir.glob(f"fold{fold}_epoch*_probs_*.npy"):
|
||||
match = pattern.match(path.name)
|
||||
if not match:
|
||||
continue
|
||||
epoch = int(match.group(1))
|
||||
head = match.group(2)
|
||||
epoch_paths.setdefault(epoch, {})[head] = path
|
||||
return epoch_paths
|
||||
|
||||
@staticmethod
|
||||
def _collect_base_prob_paths(run_dir: Path, fold: int) -> Dict[str, Path]:
|
||||
paths: Dict[str, Path] = {}
|
||||
for head in ("fused", "img", "md"):
|
||||
candidate = run_dir / f"fold{fold}_probs_{head}.npy"
|
||||
if candidate.exists():
|
||||
paths[head] = candidate
|
||||
return paths
|
||||
|
||||
@staticmethod
|
||||
def _load_probs_array(path: Path) -> Optional[np.ndarray]:
|
||||
try:
|
||||
return np.load(path)
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _prepare_probs(arr: np.ndarray) -> Optional[np.ndarray]:
|
||||
if arr is None:
|
||||
return None
|
||||
probs = np.asarray(arr, dtype=float)
|
||||
if probs.ndim == 1:
|
||||
probs = np.stack([1.0 - probs, probs], axis=1)
|
||||
if probs.ndim != 2:
|
||||
return None
|
||||
return probs
|
||||
|
||||
@staticmethod
|
||||
def _has_all_heads(arrays: Dict[str, Optional[np.ndarray]]) -> bool:
|
||||
needed = ("fused", "img", "md")
|
||||
return all(arrays.get(head) is not None for head in needed)
|
||||
|
||||
def _fusion_corrections_for_probs(self,
|
||||
y_true: np.ndarray,
|
||||
arrays: Dict[str, np.ndarray],
|
||||
run_id: str,
|
||||
fold: int,
|
||||
epoch: int) -> List[Dict[str, object]]:
|
||||
fused = self._prepare_probs(arrays.get("fused"))
|
||||
img = self._prepare_probs(arrays.get("img"))
|
||||
md = self._prepare_probs(arrays.get("md"))
|
||||
if fused is None or img is None or md is None:
|
||||
return []
|
||||
if not (len(fused) == len(img) == len(md) == len(y_true)):
|
||||
return []
|
||||
|
||||
fused_pred = fused.argmax(axis=1)
|
||||
img_pred = img.argmax(axis=1)
|
||||
md_pred = md.argmax(axis=1)
|
||||
|
||||
fused_conf = np.take_along_axis(fused, fused_pred[:, None], axis=1).squeeze(1)
|
||||
img_conf = np.take_along_axis(img, img_pred[:, None], axis=1).squeeze(1)
|
||||
md_conf = np.take_along_axis(md, md_pred[:, None], axis=1).squeeze(1)
|
||||
|
||||
mask = (fused_pred == y_true) & (img_pred != y_true) & (md_pred != y_true)
|
||||
indices = np.nonzero(mask)[0]
|
||||
|
||||
events: List[Dict[str, object]] = []
|
||||
for idx in indices:
|
||||
events.append({
|
||||
"run_id": run_id,
|
||||
"fold": fold,
|
||||
"epoch": epoch,
|
||||
"index": int(idx),
|
||||
"y_true": int(y_true[idx]),
|
||||
"pred_fused": int(fused_pred[idx]),
|
||||
"pred_img": int(img_pred[idx]),
|
||||
"pred_md": int(md_pred[idx]),
|
||||
"conf_fused": float(fused_conf[idx]),
|
||||
"conf_img": float(img_conf[idx]),
|
||||
"conf_md": float(md_conf[idx]),
|
||||
})
|
||||
return events
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
analytics = GridSearchAnalytics(Path("analysis_data/grid_search"))
|
||||
output = analytics.fusion_correction_events()
|
||||
print(f"Fusion correction events written to {output}")
|
||||
@@ -1,729 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Build an HTML heatmap-style grid for grid search runs.
|
||||
|
||||
Each column is a run. The header shows mean metrics (auc, acc, holdout_auc,
|
||||
holdout_acc). Rows encode hyperparameter options as red/green boxes.
|
||||
|
||||
Example:
|
||||
python scripts/grid_search_analytics/grid_search_heatmap.py \
|
||||
--analysis-dir analysis_data/grid_search \
|
||||
--task binary \
|
||||
--sort-by holdout_auc --desc \
|
||||
--top 40 \
|
||||
--format plot \
|
||||
--output analysis_data/grid_search_heatmap.png
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
from html import escape
|
||||
from pathlib import Path
|
||||
from typing import Dict, Iterable, List, Optional, Tuple
|
||||
|
||||
DEFAULT_EXCLUDE_KEYS = {
|
||||
"run_id",
|
||||
"fold_metrics",
|
||||
"best_metric",
|
||||
"best_metric_mode",
|
||||
"best_metric_mean",
|
||||
"best_metric_std",
|
||||
"eval_mode",
|
||||
"n_splits",
|
||||
"num_classes",
|
||||
}
|
||||
|
||||
|
||||
def to_float(value: Optional[object]) -> Optional[float]:
|
||||
if value is None:
|
||||
return None
|
||||
if isinstance(value, (int, float)):
|
||||
num = float(value)
|
||||
if math.isnan(num):
|
||||
return None
|
||||
return num
|
||||
if not isinstance(value, str):
|
||||
return None
|
||||
value = value.strip()
|
||||
if not value:
|
||||
return None
|
||||
try:
|
||||
num = float(value)
|
||||
except ValueError:
|
||||
return None
|
||||
if math.isnan(num):
|
||||
return None
|
||||
return num
|
||||
|
||||
|
||||
def mean(values: List[float]) -> Optional[float]:
|
||||
return (sum(values) / len(values)) if values else None
|
||||
|
||||
|
||||
def read_summary(run_dir: Path) -> Optional[Dict[str, object]]:
|
||||
summary_path = run_dir / "summary.json"
|
||||
if not summary_path.exists():
|
||||
return None
|
||||
try:
|
||||
data = json.loads(summary_path.read_text())
|
||||
except Exception:
|
||||
return None
|
||||
if not isinstance(data, dict):
|
||||
return None
|
||||
return data
|
||||
|
||||
|
||||
def read_cli_args(run_dir: Path) -> Optional[Dict[str, object]]:
|
||||
cli_path = run_dir / "cli_args.json"
|
||||
if not cli_path.exists():
|
||||
return None
|
||||
try:
|
||||
data = json.loads(cli_path.read_text())
|
||||
except Exception:
|
||||
return None
|
||||
if not isinstance(data, dict):
|
||||
return None
|
||||
return data
|
||||
|
||||
|
||||
def read_run_id(run_dir: Path, summary: Optional[Dict[str, object]]) -> str:
|
||||
if summary:
|
||||
rid = summary.get("run_id")
|
||||
if isinstance(rid, str) and rid:
|
||||
return rid
|
||||
return run_dir.name
|
||||
|
||||
|
||||
def task_from_summary(summary: Optional[Dict[str, object]]) -> Optional[str]:
|
||||
if not summary:
|
||||
return None
|
||||
eval_mode = summary.get("eval_mode")
|
||||
if isinstance(eval_mode, str):
|
||||
mode = eval_mode.strip().lower()
|
||||
if mode == "binary":
|
||||
return "binary"
|
||||
if mode in {"multiclass", "multi", "multi-class"}:
|
||||
return "multiclass"
|
||||
num_classes = summary.get("num_classes")
|
||||
if isinstance(num_classes, (int, float)):
|
||||
return "binary" if int(num_classes) <= 2 else "multiclass"
|
||||
return None
|
||||
|
||||
|
||||
def metric_from_stats(stats: Dict[str, object], metric: str) -> Optional[float]:
|
||||
if metric.startswith("holdout_") and stats.get("holdout_best_monitor") == metric:
|
||||
best_val = to_float(stats.get("holdout_best_so_far"))
|
||||
if best_val is not None:
|
||||
return best_val
|
||||
return to_float(stats.get(metric))
|
||||
|
||||
|
||||
def mean_metric(summary: Dict[str, object], metric: str) -> Optional[float]:
|
||||
folds = summary.get("fold_metrics") or []
|
||||
if not isinstance(folds, list) or not folds:
|
||||
return None
|
||||
values = []
|
||||
for fold in folds:
|
||||
stats = fold.get("stats") if isinstance(fold, dict) else None
|
||||
if not isinstance(stats, dict):
|
||||
return None
|
||||
val = metric_from_stats(stats, metric)
|
||||
if val is None:
|
||||
return None
|
||||
values.append(val)
|
||||
return mean(values)
|
||||
|
||||
|
||||
def flatten_config(data: Dict[str, object],
|
||||
prefix: str = "",
|
||||
exclude_keys: Optional[Iterable[str]] = None) -> Dict[str, object]:
|
||||
out: Dict[str, object] = {}
|
||||
excludes = set(exclude_keys or [])
|
||||
for key, value in data.items():
|
||||
if key in excludes or key.startswith("best_"):
|
||||
continue
|
||||
full_key = f"{prefix}{key}" if not prefix else f"{prefix}.{key}"
|
||||
if isinstance(value, dict):
|
||||
out.update(flatten_config(value, full_key, exclude_keys=excludes))
|
||||
continue
|
||||
if isinstance(value, list):
|
||||
continue
|
||||
out[full_key] = value
|
||||
return out
|
||||
|
||||
|
||||
def sort_value_key(value: object) -> Tuple[int, object]:
|
||||
if value is None:
|
||||
return (2, "")
|
||||
if isinstance(value, bool):
|
||||
return (0, int(value))
|
||||
if isinstance(value, (int, float)):
|
||||
return (0, value)
|
||||
return (1, str(value))
|
||||
|
||||
|
||||
def format_value(value: object) -> str:
|
||||
if value is None:
|
||||
return ""
|
||||
if isinstance(value, bool):
|
||||
return "true" if value else "false"
|
||||
if isinstance(value, int):
|
||||
return str(value)
|
||||
if isinstance(value, float):
|
||||
return f"{value:.6g}"
|
||||
return str(value)
|
||||
|
||||
|
||||
def format_metric(value: Optional[float]) -> str:
|
||||
if value is None:
|
||||
return ""
|
||||
return f"{value:.4f}"
|
||||
|
||||
|
||||
def render_progress(current: int, total: Optional[int], matched: int) -> str:
|
||||
if total:
|
||||
width = 30
|
||||
filled = int(width * current / total)
|
||||
bar = "#" * filled + "-" * (width - filled)
|
||||
return f"[{bar}] {current}/{total} matched {matched}"
|
||||
return f"Scanned {current} dirs, matched {matched}"
|
||||
|
||||
|
||||
def iter_run_dirs(root: Path, shallow: bool, show_progress: bool) -> Iterable[Path]:
|
||||
if shallow:
|
||||
entries = [entry for entry in root.iterdir() if entry.is_dir()]
|
||||
entries.sort(key=lambda p: p.name)
|
||||
total = len(entries)
|
||||
matched = 0
|
||||
last_update = 0.0
|
||||
for idx, entry in enumerate(entries, start=1):
|
||||
if show_progress:
|
||||
now = time.monotonic()
|
||||
if now - last_update >= 0.1 or idx == total:
|
||||
msg = render_progress(idx, total, matched)
|
||||
print(f"\rScanning {msg}", end="", file=sys.stderr, flush=True)
|
||||
last_update = now
|
||||
if (entry / "summary.json").is_file():
|
||||
matched += 1
|
||||
yield entry
|
||||
if show_progress:
|
||||
print(file=sys.stderr)
|
||||
return
|
||||
|
||||
matched = 0
|
||||
scanned = 0
|
||||
last_update = 0.0
|
||||
for dirpath, dirnames, filenames in os.walk(root):
|
||||
scanned += 1
|
||||
if show_progress:
|
||||
now = time.monotonic()
|
||||
if now - last_update >= 0.2:
|
||||
msg = render_progress(scanned, None, matched)
|
||||
print(f"\rScanning {msg}", end="", file=sys.stderr, flush=True)
|
||||
last_update = now
|
||||
if "summary.json" in filenames:
|
||||
matched += 1
|
||||
yield Path(dirpath)
|
||||
if show_progress:
|
||||
msg = render_progress(scanned, None, matched)
|
||||
print(f"\rScanning {msg}", end="", file=sys.stderr, flush=True)
|
||||
print(file=sys.stderr)
|
||||
|
||||
|
||||
def build_html(runs: List[Dict[str, object]],
|
||||
row_specs: List[Tuple[str, object]],
|
||||
title: str,
|
||||
filters: List[str]) -> str:
|
||||
lines: List[str] = []
|
||||
lines.append("<!doctype html>")
|
||||
lines.append("<html lang=\"en\">")
|
||||
lines.append("<head>")
|
||||
lines.append("<meta charset=\"utf-8\">")
|
||||
lines.append(f"<title>{escape(title)}</title>")
|
||||
lines.append("<style>")
|
||||
lines.append(":root { --green: #4caf50; --red: #d9534f; --grid: #d0d0d0; --header: #f0f0f0; }")
|
||||
lines.append("body { margin: 0; padding: 16px; font-family: \"Courier New\", monospace; }")
|
||||
lines.append(".wrap { overflow-x: auto; }")
|
||||
lines.append("table { border-collapse: collapse; font-size: 12px; }")
|
||||
lines.append("th, td { border: 1px solid var(--grid); padding: 4px; text-align: center; }")
|
||||
lines.append("th.row-label { text-align: left; background: var(--header); position: sticky; left: 0; }")
|
||||
lines.append("thead th { background: var(--header); position: sticky; top: 0; z-index: 1; }")
|
||||
lines.append("th.run-id { writing-mode: vertical-rl; transform: rotate(180deg); white-space: nowrap; }")
|
||||
lines.append("td.cell { width: 14px; height: 14px; padding: 0; }")
|
||||
lines.append("td.on { background: var(--green); }")
|
||||
lines.append("td.off { background: var(--red); }")
|
||||
lines.append(".meta { margin-bottom: 12px; }")
|
||||
lines.append("</style>")
|
||||
lines.append("</head>")
|
||||
lines.append("<body>")
|
||||
lines.append(f"<h2>{escape(title)}</h2>")
|
||||
if filters:
|
||||
lines.append("<div class=\"meta\">")
|
||||
for item in filters:
|
||||
lines.append(f"<div>{escape(item)}</div>")
|
||||
lines.append("</div>")
|
||||
lines.append("<div class=\"wrap\">")
|
||||
lines.append("<table>")
|
||||
lines.append("<thead>")
|
||||
lines.append("<tr>")
|
||||
lines.append("<th>run_id</th>")
|
||||
for run in runs:
|
||||
run_id = escape(str(run.get("run_id", "")))
|
||||
rel_path = escape(str(run.get("relative_path", "")))
|
||||
title_attr = f" title=\"{rel_path}\"" if rel_path else ""
|
||||
lines.append(f"<th class=\"run-id\"{title_attr}>{run_id}</th>")
|
||||
lines.append("</tr>")
|
||||
for metric_key, label in [
|
||||
("auc", "auc"),
|
||||
("acc", "acc"),
|
||||
("holdout_auc", "holdout_auc"),
|
||||
("holdout_acc", "holdout_acc"),
|
||||
]:
|
||||
lines.append("<tr>")
|
||||
lines.append(f"<th>{label}</th>")
|
||||
for run in runs:
|
||||
metrics = run.get("metrics", {})
|
||||
value = metrics.get(metric_key) if isinstance(metrics, dict) else None
|
||||
lines.append(f"<td>{escape(format_metric(value))}</td>")
|
||||
lines.append("</tr>")
|
||||
lines.append("</thead>")
|
||||
lines.append("<tbody>")
|
||||
for key, value in row_specs:
|
||||
label = f"{key}={format_value(value)}"
|
||||
lines.append("<tr>")
|
||||
lines.append(f"<th class=\"row-label\">{escape(label)}</th>")
|
||||
for run in runs:
|
||||
config = run.get("config", {})
|
||||
current = config.get(key) if isinstance(config, dict) else None
|
||||
cell_class = "on" if current == value else "off"
|
||||
lines.append(f"<td class=\"cell {cell_class}\"></td>")
|
||||
lines.append("</tr>")
|
||||
lines.append("</tbody>")
|
||||
lines.append("</table>")
|
||||
lines.append("</div>")
|
||||
lines.append("</body>")
|
||||
lines.append("</html>")
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
def truncate(text: str, width: int) -> str:
|
||||
if len(text) <= width:
|
||||
return text
|
||||
if width <= 3:
|
||||
return text[:width]
|
||||
return text[:width - 3] + "..."
|
||||
|
||||
|
||||
def build_text_grid(runs: List[Dict[str, object]],
|
||||
row_specs: List[Tuple[str, object]],
|
||||
filters: List[str],
|
||||
col_width: int,
|
||||
row_width: int,
|
||||
color: bool) -> str:
|
||||
sep = " "
|
||||
lines: List[str] = []
|
||||
if filters:
|
||||
lines.extend(filters)
|
||||
lines.append("")
|
||||
|
||||
def pad(text: str, width: int) -> str:
|
||||
return truncate(text, width).ljust(width)
|
||||
|
||||
def colorize(text: str, enabled: bool) -> str:
|
||||
if not color:
|
||||
return text
|
||||
color_code = "\x1b[32m" if enabled else "\x1b[31m"
|
||||
return f"{color_code}{text}\x1b[0m"
|
||||
|
||||
def row_line(label: str, values: List[str]) -> str:
|
||||
return pad(label, row_width) + sep + sep.join(pad(v, col_width) for v in values)
|
||||
|
||||
run_ids = [str(run.get("run_id", "")) for run in runs]
|
||||
lines.append(row_line("run_id", run_ids))
|
||||
for metric_key, label in [
|
||||
("auc", "auc"),
|
||||
("acc", "acc"),
|
||||
("holdout_auc", "holdout_auc"),
|
||||
("holdout_acc", "holdout_acc"),
|
||||
]:
|
||||
values = []
|
||||
for run in runs:
|
||||
metrics = run.get("metrics", {})
|
||||
value = metrics.get(metric_key) if isinstance(metrics, dict) else None
|
||||
values.append(format_metric(value))
|
||||
lines.append(row_line(label, values))
|
||||
|
||||
divider = "-" * row_width + sep + sep.join("-" * col_width for _ in runs)
|
||||
lines.append(divider)
|
||||
|
||||
for key, value in row_specs:
|
||||
label = f"{key}={format_value(value)}"
|
||||
cells: List[str] = []
|
||||
for run in runs:
|
||||
config = run.get("config", {})
|
||||
current = config.get(key) if isinstance(config, dict) else None
|
||||
enabled = current == value
|
||||
cell = colorize("##", enabled) if enabled else colorize("..", enabled)
|
||||
cells.append(cell)
|
||||
lines.append(row_line(label, cells))
|
||||
|
||||
lines.append("")
|
||||
lines.append("Legend: ##=on ..=off")
|
||||
if color:
|
||||
lines.append("Colors: green=on red=off")
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
def parse_figsize(value: Optional[str], n_cols: int, n_rows: int) -> Tuple[float, float]:
|
||||
if value:
|
||||
parts = [p.strip() for p in value.split(",") if p.strip()]
|
||||
if len(parts) == 2:
|
||||
try:
|
||||
return float(parts[0]), float(parts[1])
|
||||
except ValueError:
|
||||
pass
|
||||
width = min(40.0, max(8.0, n_cols * 0.3))
|
||||
height = min(40.0, max(6.0, (n_rows + 6) * 0.3))
|
||||
return width, height
|
||||
|
||||
|
||||
def plot_heatmap(runs: List[Dict[str, object]],
|
||||
row_specs: List[Tuple[str, object]],
|
||||
filters: List[str],
|
||||
output_path: Optional[Path],
|
||||
figsize: Tuple[float, float],
|
||||
dpi: int,
|
||||
show: bool) -> None:
|
||||
if not show:
|
||||
import matplotlib
|
||||
matplotlib.use("Agg")
|
||||
|
||||
import matplotlib.pyplot as plt
|
||||
from matplotlib.colors import ListedColormap
|
||||
|
||||
try:
|
||||
import seaborn as sns
|
||||
except ImportError:
|
||||
sns = None
|
||||
|
||||
metric_labels = ["auc", "acc", "holdout_auc", "holdout_acc"]
|
||||
metric_matrix: List[List[float]] = []
|
||||
for label in metric_labels:
|
||||
row: List[float] = []
|
||||
for run in runs:
|
||||
metrics = run.get("metrics", {})
|
||||
value = metrics.get(label) if isinstance(metrics, dict) else None
|
||||
row.append(float(value) if value is not None else float("nan"))
|
||||
metric_matrix.append(row)
|
||||
|
||||
param_labels = [f"{key}={format_value(value)}" for key, value in row_specs]
|
||||
param_matrix: List[List[int]] = []
|
||||
for key, value in row_specs:
|
||||
row = []
|
||||
for run in runs:
|
||||
config = run.get("config", {})
|
||||
current = config.get(key) if isinstance(config, dict) else None
|
||||
row.append(1 if current == value else 0)
|
||||
param_matrix.append(row)
|
||||
|
||||
fig = plt.figure(figsize=figsize, dpi=dpi)
|
||||
grid_rows = 2 if param_matrix else 1
|
||||
height_ratios = [2, max(2, len(param_matrix) * 0.5)] if param_matrix else [2]
|
||||
gs = fig.add_gridspec(grid_rows, 1, height_ratios=height_ratios, hspace=0.05)
|
||||
|
||||
ax_metrics = fig.add_subplot(gs[0, 0])
|
||||
if sns:
|
||||
sns.heatmap(
|
||||
metric_matrix,
|
||||
ax=ax_metrics,
|
||||
cmap="viridis",
|
||||
annot=True,
|
||||
fmt=".3f",
|
||||
cbar=True,
|
||||
yticklabels=metric_labels,
|
||||
xticklabels=False,
|
||||
)
|
||||
else:
|
||||
im = ax_metrics.imshow(metric_matrix, aspect="auto", cmap="viridis")
|
||||
ax_metrics.set_yticks(range(len(metric_labels)))
|
||||
ax_metrics.set_yticklabels(metric_labels)
|
||||
fig.colorbar(im, ax=ax_metrics, fraction=0.02, pad=0.01)
|
||||
for i, row in enumerate(metric_matrix):
|
||||
for j, value in enumerate(row):
|
||||
if math.isnan(value):
|
||||
continue
|
||||
ax_metrics.text(j, i, f"{value:.3f}", ha="center", va="center", fontsize=7, color="white")
|
||||
ax_metrics.set_ylabel("metrics")
|
||||
|
||||
if param_matrix:
|
||||
ax_params = fig.add_subplot(gs[1, 0], sharex=ax_metrics)
|
||||
cmap = ListedColormap(["#d9534f", "#4caf50"])
|
||||
if sns:
|
||||
sns.heatmap(
|
||||
param_matrix,
|
||||
ax=ax_params,
|
||||
cmap=cmap,
|
||||
cbar=False,
|
||||
yticklabels=param_labels,
|
||||
xticklabels=[run.get("run_id", "") for run in runs],
|
||||
vmin=0,
|
||||
vmax=1,
|
||||
)
|
||||
else:
|
||||
ax_params.imshow(param_matrix, aspect="auto", cmap=cmap, vmin=0, vmax=1)
|
||||
ax_params.set_yticks(range(len(param_labels)))
|
||||
ax_params.set_yticklabels(param_labels)
|
||||
ax_params.set_xticks(range(len(runs)))
|
||||
ax_params.set_xticklabels([run.get("run_id", "") for run in runs], rotation=90)
|
||||
ax_params.set_xlabel("runs")
|
||||
else:
|
||||
ax_metrics.set_xticks(range(len(runs)))
|
||||
ax_metrics.set_xticklabels([run.get("run_id", "") for run in runs], rotation=90)
|
||||
ax_metrics.set_xlabel("runs")
|
||||
|
||||
if filters:
|
||||
fig.suptitle("Grid Search Heatmap\n" + " | ".join(filters), fontsize=10)
|
||||
else:
|
||||
fig.suptitle("Grid Search Heatmap", fontsize=10)
|
||||
|
||||
if output_path is not None:
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
fig.savefig(output_path, bbox_inches="tight")
|
||||
if show:
|
||||
plt.show()
|
||||
plt.close(fig)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
ap = argparse.ArgumentParser(description="Build an HTML heatmap grid for grid search runs.")
|
||||
ap.add_argument("--analysis-dir", type=Path, default=Path("analysis_data/grid_search"),
|
||||
help="Directory containing run subdirectories")
|
||||
ap.add_argument("--format", choices=["text", "html", "plot"], default="text",
|
||||
help="Output format (default: text)")
|
||||
ap.add_argument("--task", choices=["binary", "multiclass", "all"], default="all",
|
||||
help="Filter runs by task type (default: all)")
|
||||
ap.add_argument("--sort-by", choices=["auc", "acc", "holdout_auc", "holdout_acc"], default=None,
|
||||
help="Metric to sort columns by (default: none)")
|
||||
ap.add_argument("--asc", action="store_true",
|
||||
help="Sort in ascending order (default: descending)")
|
||||
ap.add_argument("--desc", action="store_true",
|
||||
help="Sort in descending order (default: descending)")
|
||||
ap.add_argument("--top", type=int, default=None,
|
||||
help="Limit to the top N runs after sorting")
|
||||
ap.add_argument("--cluster-rows", dest="cluster_rows", action="store_true",
|
||||
help="Order parameter rows by prevalence in the selected runs (default)")
|
||||
ap.add_argument("--no-cluster-rows", dest="cluster_rows", action="store_false",
|
||||
help="Keep parameter rows sorted alphabetically")
|
||||
ap.set_defaults(cluster_rows=True)
|
||||
ap.add_argument("--output", type=Path, default=None,
|
||||
help="Optional path to write output")
|
||||
ap.add_argument("--params", default=None,
|
||||
help="Comma-separated list of parameter keys to include")
|
||||
ap.add_argument("--exclude", default=None,
|
||||
help="Comma-separated list of parameter keys to exclude")
|
||||
ap.add_argument("--match", default=None,
|
||||
help="Only include run directories whose name contains this substring")
|
||||
ap.add_argument("--shallow", action="store_true",
|
||||
help="Only scan directories directly under analysis-dir")
|
||||
ap.add_argument("--no-progress", action="store_true",
|
||||
help="Disable progress output")
|
||||
ap.add_argument("--col-width", type=int, default=13,
|
||||
help="Column width for text output (default: 13)")
|
||||
ap.add_argument("--row-width", type=int, default=36,
|
||||
help="Row label width for text output (default: 36)")
|
||||
ap.add_argument("--color", action="store_true",
|
||||
help="Use ANSI colors in text output")
|
||||
ap.add_argument("--figsize", default=None,
|
||||
help="Figure size as 'width,height' (inches), for plot output")
|
||||
ap.add_argument("--dpi", type=int, default=140,
|
||||
help="Figure DPI for plot output")
|
||||
ap.add_argument("--show", action="store_true",
|
||||
help="Display plot window (only for format=plot)")
|
||||
args = ap.parse_args()
|
||||
|
||||
root = args.analysis_dir
|
||||
if not root.exists():
|
||||
raise SystemExit(f"Analysis directory not found: {root}")
|
||||
|
||||
exclude_keys = set(DEFAULT_EXCLUDE_KEYS)
|
||||
if args.exclude:
|
||||
for item in args.exclude.split(","):
|
||||
item = item.strip()
|
||||
if item:
|
||||
exclude_keys.add(item)
|
||||
|
||||
runs: List[Dict[str, object]] = []
|
||||
values_by_key: Dict[str, List[object]] = {}
|
||||
missing_summary = 0
|
||||
unknown_task = 0
|
||||
missing_cli = 0
|
||||
|
||||
for run_dir in iter_run_dirs(root, shallow=args.shallow, show_progress=not args.no_progress):
|
||||
if args.match and args.match not in run_dir.name:
|
||||
continue
|
||||
summary = read_summary(run_dir)
|
||||
if summary is None:
|
||||
missing_summary += 1
|
||||
continue
|
||||
task_label = task_from_summary(summary)
|
||||
if args.task != "all":
|
||||
if task_label is None:
|
||||
unknown_task += 1
|
||||
continue
|
||||
if task_label != args.task:
|
||||
continue
|
||||
|
||||
metrics = {
|
||||
"auc": mean_metric(summary, "auc_fused"),
|
||||
"acc": mean_metric(summary, "acc_fused"),
|
||||
"holdout_auc": mean_metric(summary, "holdout_auc_fused"),
|
||||
"holdout_acc": mean_metric(summary, "holdout_acc_fused"),
|
||||
}
|
||||
if any(val is None for val in metrics.values()):
|
||||
continue
|
||||
|
||||
cli_args = read_cli_args(run_dir)
|
||||
if cli_args is None:
|
||||
missing_cli += 1
|
||||
config_source = cli_args if cli_args is not None else summary
|
||||
config = flatten_config(config_source, exclude_keys=exclude_keys)
|
||||
run_id = read_run_id(run_dir, summary)
|
||||
runs.append({
|
||||
"run_id": run_id,
|
||||
"relative_path": str(run_dir.relative_to(root)),
|
||||
"task": task_label,
|
||||
"metrics": metrics,
|
||||
"config": config,
|
||||
})
|
||||
for key, value in config.items():
|
||||
values_by_key.setdefault(key, []).append(value)
|
||||
|
||||
if not runs:
|
||||
print("No matching runs found.")
|
||||
return
|
||||
|
||||
if args.params:
|
||||
param_keys = [p.strip() for p in args.params.split(",") if p.strip()]
|
||||
else:
|
||||
param_keys = []
|
||||
for key, values in values_by_key.items():
|
||||
unique_values = {format_value(v) for v in values}
|
||||
if len(unique_values) > 1:
|
||||
param_keys.append(key)
|
||||
param_keys.sort()
|
||||
|
||||
row_specs: List[Tuple[str, object]] = []
|
||||
for key in param_keys:
|
||||
values = values_by_key.get(key, [])
|
||||
unique_values = []
|
||||
seen = set()
|
||||
for val in values:
|
||||
marker = (type(val), val)
|
||||
if marker in seen:
|
||||
continue
|
||||
seen.add(marker)
|
||||
unique_values.append(val)
|
||||
unique_values.sort(key=sort_value_key)
|
||||
for value in unique_values:
|
||||
row_specs.append((key, value))
|
||||
|
||||
if args.asc and args.desc:
|
||||
raise SystemExit("Choose only one of --asc or --desc.")
|
||||
|
||||
if args.sort_by:
|
||||
def sort_key(item: Dict[str, object]) -> float:
|
||||
metrics = item.get("metrics", {})
|
||||
val = metrics.get(args.sort_by) if isinstance(metrics, dict) else None
|
||||
if val is None:
|
||||
return float("inf") if args.asc else float("-inf")
|
||||
return float(val)
|
||||
|
||||
runs.sort(key=sort_key, reverse=not args.asc)
|
||||
else:
|
||||
runs.sort(key=lambda r: str(r.get("run_id", "")))
|
||||
|
||||
if args.top is not None:
|
||||
runs = runs[:args.top]
|
||||
|
||||
if args.cluster_rows and runs:
|
||||
total = len(runs)
|
||||
counts_by_spec: Dict[Tuple[str, object], int] = {}
|
||||
for key, value in row_specs:
|
||||
counts_by_spec[(key, value)] = 0
|
||||
for run in runs:
|
||||
config = run.get("config", {})
|
||||
if not isinstance(config, dict):
|
||||
continue
|
||||
for key, value in row_specs:
|
||||
if config.get(key) == value:
|
||||
counts_by_spec[(key, value)] += 1
|
||||
|
||||
def row_sort(spec: Tuple[str, object]) -> Tuple[float, str, str]:
|
||||
count = counts_by_spec.get(spec, 0)
|
||||
score = count / total if total else 0.0
|
||||
key, value = spec
|
||||
return (-score, str(key), format_value(value))
|
||||
|
||||
row_specs.sort(key=row_sort)
|
||||
|
||||
filters = []
|
||||
if args.task != "all":
|
||||
filters.append(f"Task filter: {args.task}")
|
||||
if args.match:
|
||||
filters.append(f"Name filter: {args.match}")
|
||||
filters.append(f"Runs: {len(runs)}")
|
||||
filters.append(f"Params: {len(row_specs)}")
|
||||
filters.append(f"Row clustering: {'on' if args.cluster_rows else 'off'}")
|
||||
if missing_summary or unknown_task:
|
||||
filters.append(f"Skipped: {missing_summary} missing summary, {unknown_task} unknown task")
|
||||
if missing_cli:
|
||||
filters.append(f"Missing cli_args: {missing_cli}")
|
||||
|
||||
title = "Grid Search Heatmap"
|
||||
if args.format == "html":
|
||||
html = build_html(runs, row_specs, title=title, filters=filters)
|
||||
output_path = args.output or Path("analysis_data/grid_search_heatmap.html")
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
output_path.write_text(html)
|
||||
print(f"Wrote {output_path}")
|
||||
elif args.format == "plot":
|
||||
output_path = args.output
|
||||
if output_path is None and not args.show:
|
||||
output_path = Path("analysis_data/grid_search_heatmap.png")
|
||||
figsize = parse_figsize(args.figsize, n_cols=len(runs), n_rows=len(row_specs))
|
||||
plot_heatmap(
|
||||
runs,
|
||||
row_specs,
|
||||
filters=filters,
|
||||
output_path=output_path,
|
||||
figsize=figsize,
|
||||
dpi=args.dpi,
|
||||
show=args.show,
|
||||
)
|
||||
if output_path is not None:
|
||||
print(f"Wrote {output_path}")
|
||||
else:
|
||||
text = build_text_grid(
|
||||
runs,
|
||||
row_specs,
|
||||
filters=filters,
|
||||
col_width=max(4, args.col_width),
|
||||
row_width=max(12, args.row_width),
|
||||
color=args.color,
|
||||
)
|
||||
if args.output:
|
||||
args.output.parent.mkdir(parents=True, exist_ok=True)
|
||||
args.output.write_text(text)
|
||||
print(f"Wrote {args.output}")
|
||||
else:
|
||||
print(text)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,572 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Plot holdout ROC curves for a specific grid search run."""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import re
|
||||
from pathlib import Path
|
||||
from typing import Dict, List, Tuple
|
||||
|
||||
import matplotlib.pyplot as plt
|
||||
import numpy as np
|
||||
from sklearn.metrics import auc, roc_curve
|
||||
|
||||
|
||||
# ---------------------------
|
||||
# Config (edit in IDE)
|
||||
# ---------------------------
|
||||
RUN_ID = "20251129-0312"
|
||||
ANALYSIS_ROOT = Path("analysis_data/grid_search")
|
||||
HEADS = ["fused", "image", "metadata"]
|
||||
OUTPUT_SUBDIR = Path("plots/holdout_rocs")
|
||||
BEST_OUTPUT_SUBDIR = Path("plots/best_rocs")
|
||||
POSITIVE_CLASS = 1
|
||||
DEBUG = True
|
||||
USE_JSON_ROC = True
|
||||
USE_HOLDOUT_PROBS = True
|
||||
ALLOW_FALLBACK_TO_VALIDATION = False
|
||||
PLOT_VALIDATION_FROM_HOLDOUT_EPOCH = True
|
||||
PLOT_BEST_EPOCH = True
|
||||
PLOT_HOLDOUT_FROM_BEST_EPOCH = True
|
||||
PLOT_ALL_CLASSES = True
|
||||
FORCE_PROBS_FOR_BEST_BINARY = True
|
||||
FORCE_PROBS_FOR_HOLDOUT_BINARY = False
|
||||
|
||||
|
||||
HEAD_FILE_KEYS = {
|
||||
"fused": "fused",
|
||||
"image": "img",
|
||||
"metadata": "md",
|
||||
}
|
||||
|
||||
JSON_DIR_NAMES = [
|
||||
"roc_curves_holdout_best",
|
||||
"roc_curves",
|
||||
]
|
||||
|
||||
|
||||
def _epoch_from_name(path: Path) -> int:
|
||||
m = re.search(r"epoch(\d+)", path.name)
|
||||
return int(m.group(1)) if m else -1
|
||||
|
||||
|
||||
def _load_json(path: Path) -> Dict:
|
||||
try:
|
||||
return json.loads(path.read_text())
|
||||
except Exception:
|
||||
return {}
|
||||
|
||||
|
||||
def _infer_run_info(run_dir: Path) -> Tuple[str | None, int | None, List[str] | None]:
|
||||
cli = _load_json(run_dir / "cli_args.json")
|
||||
summary = _load_json(run_dir / "summary.json")
|
||||
payloads = [cli, summary]
|
||||
eval_mode = None
|
||||
num_classes = None
|
||||
class_names = None
|
||||
for payload in payloads:
|
||||
if not payload:
|
||||
continue
|
||||
if eval_mode is None:
|
||||
em = payload.get("eval_mode")
|
||||
if isinstance(em, str):
|
||||
eval_mode = em.strip().lower()
|
||||
if num_classes is None:
|
||||
nc = payload.get("num_classes")
|
||||
if isinstance(nc, (int, float)):
|
||||
num_classes = int(nc)
|
||||
if class_names is None:
|
||||
cn = payload.get("class_names")
|
||||
if isinstance(cn, list) and cn:
|
||||
class_names = [str(x) for x in cn]
|
||||
if num_classes is None and eval_mode:
|
||||
num_classes = 2 if eval_mode == "binary" else 3
|
||||
return eval_mode, num_classes, class_names
|
||||
|
||||
|
||||
def _collect_holdout_json_files(run_dir: Path, head: str) -> Dict[int, Path]:
|
||||
fold_files: Dict[int, Path] = {}
|
||||
# Fold-scoped folders
|
||||
for folder_name in JSON_DIR_NAMES:
|
||||
for fold_dir in run_dir.glob(f"fold*_{folder_name}"):
|
||||
fold_match = re.search(r"fold(\d+)_", fold_dir.name)
|
||||
if not fold_match:
|
||||
continue
|
||||
fold_idx = int(fold_match.group(1))
|
||||
candidates = list(fold_dir.glob(f"epoch*_holdout_{head}.json"))
|
||||
if not candidates:
|
||||
candidates = list(fold_dir.glob(f"epoch*_{head}.json"))
|
||||
if candidates:
|
||||
candidates.sort(key=_epoch_from_name)
|
||||
fold_files[fold_idx] = candidates[-1]
|
||||
if fold_files:
|
||||
return fold_files
|
||||
# Fallback: unscoped roc_curves in run_dir (single-fold or in-progress)
|
||||
for folder_name in JSON_DIR_NAMES:
|
||||
base_dir = run_dir / folder_name
|
||||
if not base_dir.exists():
|
||||
continue
|
||||
candidates = list(base_dir.glob(f"epoch*_holdout_{head}.json"))
|
||||
if not candidates:
|
||||
candidates = list(base_dir.glob(f"epoch*_{head}.json"))
|
||||
if candidates:
|
||||
candidates.sort(key=_epoch_from_name)
|
||||
fold_files[0] = candidates[-1]
|
||||
break
|
||||
return fold_files
|
||||
|
||||
|
||||
def _collect_validation_json_files(
|
||||
holdout_files: Dict[int, Path], head: str
|
||||
) -> Dict[int, Path]:
|
||||
validation_files: Dict[int, Path] = {}
|
||||
for fold_idx, holdout_path in holdout_files.items():
|
||||
epoch = _epoch_from_name(holdout_path)
|
||||
if epoch < 0:
|
||||
continue
|
||||
candidate = holdout_path.parent / f"epoch{epoch}_{head}.json"
|
||||
if candidate.exists():
|
||||
validation_files[fold_idx] = candidate
|
||||
continue
|
||||
# Fallback: try the same epoch under roc_curves (if holdout_best folder omitted it).
|
||||
for folder_name in JSON_DIR_NAMES:
|
||||
alt_dir = holdout_path.parent.parent / f"fold{fold_idx}_{folder_name}"
|
||||
alt_candidate = alt_dir / f"epoch{epoch}_{head}.json"
|
||||
if alt_candidate.exists():
|
||||
validation_files[fold_idx] = alt_candidate
|
||||
break
|
||||
return validation_files
|
||||
|
||||
|
||||
def _collect_holdout_from_validation_files(
|
||||
validation_files: Dict[int, Path], head: str
|
||||
) -> Dict[int, Path]:
|
||||
holdout_files: Dict[int, Path] = {}
|
||||
for fold_idx, val_path in validation_files.items():
|
||||
epoch = _epoch_from_name(val_path)
|
||||
if epoch < 0:
|
||||
continue
|
||||
candidate = val_path.parent / f"epoch{epoch}_holdout_{head}.json"
|
||||
if candidate.exists():
|
||||
holdout_files[fold_idx] = candidate
|
||||
continue
|
||||
for folder_name in ("roc_curves", "roc_curves_holdout_best"):
|
||||
alt_dir = val_path.parent.parent / f"fold{fold_idx}_{folder_name}"
|
||||
alt_candidate = alt_dir / f"epoch{epoch}_holdout_{head}.json"
|
||||
if alt_candidate.exists():
|
||||
holdout_files[fold_idx] = alt_candidate
|
||||
break
|
||||
return holdout_files
|
||||
|
||||
|
||||
def _collect_best_json_files(run_dir: Path, head: str) -> Dict[int, Path]:
|
||||
fold_files: Dict[int, Path] = {}
|
||||
for fold_dir in run_dir.glob("fold*_roc_curves_best"):
|
||||
fold_match = re.search(r"fold(\d+)_", fold_dir.name)
|
||||
if not fold_match:
|
||||
continue
|
||||
fold_idx = int(fold_match.group(1))
|
||||
candidates = list(fold_dir.glob(f"epoch*_{head}.json"))
|
||||
if candidates:
|
||||
candidates.sort(key=_epoch_from_name)
|
||||
fold_files[fold_idx] = candidates[-1]
|
||||
return fold_files
|
||||
|
||||
|
||||
def _extract_curves(data: Dict) -> Dict[str, Tuple[List[float], List[float], float]]:
|
||||
curves: Dict[str, Tuple[List[float], List[float], float]] = {}
|
||||
per_class = data.get("per_class") if isinstance(data, dict) else None
|
||||
if not isinstance(per_class, dict):
|
||||
return curves
|
||||
for cls, entry in per_class.items():
|
||||
if not isinstance(entry, dict):
|
||||
continue
|
||||
fpr = entry.get("fpr")
|
||||
tpr = entry.get("tpr")
|
||||
auc_val = entry.get("auc")
|
||||
if not isinstance(fpr, list) or not isinstance(tpr, list):
|
||||
continue
|
||||
try:
|
||||
auc_f = float(auc_val) if auc_val is not None else float("nan")
|
||||
except Exception:
|
||||
auc_f = float("nan")
|
||||
curves[str(cls)] = (fpr, tpr, auc_f)
|
||||
return curves
|
||||
|
||||
|
||||
def _derive_positive_from_class0(
|
||||
curves_by_class: Dict[str, List[Tuple[int, List[float], List[float], float]]],
|
||||
positive_class: int,
|
||||
) -> None:
|
||||
zero_key = "0"
|
||||
if zero_key not in curves_by_class:
|
||||
return
|
||||
derived = []
|
||||
for fold_idx, fpr0, tpr0, auc0 in curves_by_class.get(zero_key, []):
|
||||
# The JSON for binary currently stores class-1 labels with class-0 scores,
|
||||
# so invert the curve to recover the true class-1 ROC.
|
||||
fpr1 = [1.0 - float(x) for x in fpr0]
|
||||
tpr1 = [1.0 - float(x) for x in tpr0]
|
||||
# Ensure increasing FPR for plotting.
|
||||
if len(fpr1) > 1 and fpr1[0] > fpr1[-1]:
|
||||
fpr1 = list(reversed(fpr1))
|
||||
tpr1 = list(reversed(tpr1))
|
||||
auc1 = 1.0 - auc0 if auc0 == auc0 else auc0
|
||||
derived.append((fold_idx, fpr1, tpr1, auc1))
|
||||
curves_by_class[str(positive_class)] = derived
|
||||
|
||||
|
||||
def _needs_positive_derivation(
|
||||
curves_by_class: Dict[str, List[Tuple[int, List[float], List[float], float]]],
|
||||
positive_class: int,
|
||||
) -> bool:
|
||||
curves = curves_by_class.get(str(positive_class))
|
||||
if not curves:
|
||||
return True
|
||||
for _, fpr, tpr, auc_val in curves:
|
||||
if auc_val == auc_val and len(fpr) > 2 and len(tpr) > 2:
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def _collect_prob_files(run_dir: Path, suffix: str) -> Dict[int, Dict[str, Path]]:
|
||||
files: Dict[int, Dict[str, Path]] = {}
|
||||
for y_file in run_dir.glob(f"fold*_y_true{suffix}.npy"):
|
||||
fold_str = y_file.stem.split("_")[0].replace("fold", "")
|
||||
try:
|
||||
fold_idx = int(fold_str)
|
||||
except ValueError:
|
||||
continue
|
||||
files.setdefault(fold_idx, {})["y_true"] = y_file
|
||||
for head, key in HEAD_FILE_KEYS.items():
|
||||
for p_file in run_dir.glob(f"fold*_probs_{key}{suffix}.npy"):
|
||||
fold_str = p_file.stem.split("_")[0].replace("fold", "")
|
||||
try:
|
||||
fold_idx = int(fold_str)
|
||||
except ValueError:
|
||||
continue
|
||||
files.setdefault(fold_idx, {})[head] = p_file
|
||||
return files
|
||||
|
||||
|
||||
def _load_array(path: Path) -> np.ndarray | None:
|
||||
try:
|
||||
return np.load(path)
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
def _compute_binary_curve(
|
||||
y_true: np.ndarray, probs: np.ndarray, positive_class: int
|
||||
) -> Tuple[List[float], List[float], float] | None:
|
||||
if probs.ndim == 1:
|
||||
scores = probs
|
||||
elif probs.ndim == 2 and probs.shape[1] > positive_class:
|
||||
scores = probs[:, positive_class]
|
||||
else:
|
||||
return None
|
||||
y_bin = (y_true == positive_class).astype(int)
|
||||
if y_bin.sum() == 0 or y_bin.sum() == len(y_bin):
|
||||
return None
|
||||
fpr, tpr, _ = roc_curve(y_bin, scores)
|
||||
auc_val = float(auc(fpr, tpr))
|
||||
return fpr.tolist(), tpr.tolist(), auc_val
|
||||
|
||||
|
||||
def _compute_multiclass_curves(
|
||||
y_true: np.ndarray, probs: np.ndarray
|
||||
) -> Dict[str, Tuple[List[float], List[float], float]]:
|
||||
curves: Dict[str, Tuple[List[float], List[float], float]] = {}
|
||||
if probs.ndim != 2:
|
||||
return curves
|
||||
num_classes = probs.shape[1]
|
||||
for cls in range(num_classes):
|
||||
y_bin = (y_true == cls).astype(int)
|
||||
if y_bin.sum() == 0 or y_bin.sum() == len(y_bin):
|
||||
continue
|
||||
fpr, tpr, _ = roc_curve(y_bin, probs[:, cls])
|
||||
curves[str(cls)] = (fpr.tolist(), tpr.tolist(), float(auc(fpr, tpr)))
|
||||
return curves
|
||||
|
||||
|
||||
def _plot_overlays(
|
||||
curves_by_fold: List[Tuple[int, List[float], List[float], float]],
|
||||
title: str,
|
||||
out_path: Path,
|
||||
) -> None:
|
||||
fig, ax = plt.subplots(figsize=(6, 5))
|
||||
for fold_idx, fpr, tpr, auc_val in curves_by_fold:
|
||||
label = (
|
||||
f"fold{fold_idx} AUC={auc_val:.3f}"
|
||||
if auc_val == auc_val
|
||||
else f"fold{fold_idx}"
|
||||
)
|
||||
ax.plot(fpr, tpr, lw=1.4, label=label)
|
||||
ax.plot([0, 1], [0, 1], "k--", lw=1)
|
||||
ax.set_xlabel("False Positive Rate")
|
||||
ax.set_ylabel("True Positive Rate")
|
||||
ax.set_title(title)
|
||||
ax.legend(loc="lower right", fontsize="small")
|
||||
ax.grid(True, alpha=0.3, linestyle="--")
|
||||
fig.tight_layout()
|
||||
out_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
fig.savefig(out_path, dpi=170)
|
||||
plt.close(fig)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
run_dir = ANALYSIS_ROOT / RUN_ID
|
||||
if not run_dir.exists():
|
||||
raise SystemExit(f"Run not found: {run_dir}")
|
||||
eval_mode, num_classes, class_names = _infer_run_info(run_dir)
|
||||
is_binary = eval_mode == "binary" or num_classes == 2
|
||||
|
||||
def _plot_set(
|
||||
label: str,
|
||||
head: str,
|
||||
json_files: Dict[int, Path],
|
||||
out_dir: Path,
|
||||
paired_files: Dict[int, Path] | None,
|
||||
paired_suffix: str,
|
||||
class_names: List[str] | None,
|
||||
) -> None:
|
||||
if not json_files:
|
||||
return
|
||||
curves_by_class: Dict[
|
||||
str, List[Tuple[int, List[float], List[float], float]]
|
||||
] = {}
|
||||
paired_curves_by_class: Dict[
|
||||
str, List[Tuple[int, List[float], List[float], float]]
|
||||
] = {}
|
||||
for fold_idx, path in sorted(json_files.items()):
|
||||
if DEBUG:
|
||||
print(f"[debug] {label} head={head} fold={fold_idx} json={path}")
|
||||
data = _load_json(path)
|
||||
curves = _extract_curves(data)
|
||||
for cls, (fpr, tpr, auc_val) in curves.items():
|
||||
curves_by_class.setdefault(cls, []).append(
|
||||
(fold_idx, fpr, tpr, auc_val)
|
||||
)
|
||||
if paired_files:
|
||||
p_path = paired_files.get(fold_idx)
|
||||
if p_path is not None:
|
||||
if DEBUG:
|
||||
print(
|
||||
f"[debug] {label} head={head} fold={fold_idx} paired_json={p_path}"
|
||||
)
|
||||
p_data = _load_json(p_path)
|
||||
p_curves = _extract_curves(p_data)
|
||||
for cls, (fpr, tpr, auc_val) in p_curves.items():
|
||||
paired_curves_by_class.setdefault(cls, []).append(
|
||||
(fold_idx, fpr, tpr, auc_val)
|
||||
)
|
||||
if _needs_positive_derivation(curves_by_class, POSITIVE_CLASS):
|
||||
_derive_positive_from_class0(curves_by_class, POSITIVE_CLASS)
|
||||
if paired_curves_by_class and _needs_positive_derivation(
|
||||
paired_curves_by_class, POSITIVE_CLASS
|
||||
):
|
||||
_derive_positive_from_class0(paired_curves_by_class, POSITIVE_CLASS)
|
||||
|
||||
if PLOT_ALL_CLASSES:
|
||||
classes = list(curves_by_class.keys())
|
||||
else:
|
||||
classes = (
|
||||
[str(POSITIVE_CLASS)]
|
||||
if str(POSITIVE_CLASS) in curves_by_class
|
||||
else list(curves_by_class.keys())
|
||||
)
|
||||
if not classes:
|
||||
return
|
||||
|
||||
for cls in classes:
|
||||
fold_curves = curves_by_class.get(cls, [])
|
||||
if not fold_curves:
|
||||
continue
|
||||
class_label = cls
|
||||
if class_names is not None:
|
||||
try:
|
||||
idx = int(cls)
|
||||
if 0 <= idx < len(class_names):
|
||||
class_label = f"{cls} ({class_names[idx]})"
|
||||
except Exception:
|
||||
pass
|
||||
title = f"{RUN_ID} {label} ROC — head={head} class={class_label}"
|
||||
out_path = out_dir / f"{label}_{head}_class{cls}.png"
|
||||
_plot_overlays(fold_curves, title, out_path)
|
||||
print(f"[ok] {out_path}")
|
||||
|
||||
if paired_curves_by_class:
|
||||
p_curves = paired_curves_by_class.get(cls, [])
|
||||
if p_curves:
|
||||
p_title = f"{RUN_ID} {label} {paired_suffix} ROC — head={head} class={class_label}"
|
||||
p_path = out_dir / f"{label}_{head}_class{cls}_{paired_suffix}.png"
|
||||
_plot_overlays(p_curves, p_title, p_path)
|
||||
print(f"[ok] {p_path}")
|
||||
|
||||
def _plot_from_probs(
|
||||
label: str,
|
||||
head: str,
|
||||
out_dir: Path,
|
||||
suffix: str,
|
||||
class_names: List[str] | None,
|
||||
) -> None:
|
||||
curves_by_class: Dict[
|
||||
str, List[Tuple[int, List[float], List[float], float]]
|
||||
] = {}
|
||||
files = _collect_prob_files(run_dir, suffix)
|
||||
if not files and suffix and ALLOW_FALLBACK_TO_VALIDATION:
|
||||
files = _collect_prob_files(run_dir, "")
|
||||
if files:
|
||||
print(
|
||||
"[warn] Holdout probability dumps not found; using validation probabilities instead."
|
||||
)
|
||||
if not files:
|
||||
return
|
||||
for fold_idx in sorted(files.keys()):
|
||||
fold_files = files[fold_idx]
|
||||
y_path = fold_files.get("y_true")
|
||||
p_path = fold_files.get(head)
|
||||
if y_path is None or p_path is None:
|
||||
continue
|
||||
y_true = _load_array(y_path)
|
||||
probs = _load_array(p_path)
|
||||
if y_true is None or probs is None:
|
||||
continue
|
||||
curves = _compute_multiclass_curves(y_true, probs)
|
||||
for cls, payload in curves.items():
|
||||
curves_by_class.setdefault(cls, []).append((fold_idx, *payload))
|
||||
if not curves_by_class:
|
||||
return
|
||||
classes = list(curves_by_class.keys())
|
||||
for cls in classes:
|
||||
fold_curves = curves_by_class.get(cls, [])
|
||||
if not fold_curves:
|
||||
continue
|
||||
class_label = cls
|
||||
if class_names is not None:
|
||||
try:
|
||||
idx = int(cls)
|
||||
if 0 <= idx < len(class_names):
|
||||
class_label = f"{cls} ({class_names[idx]})"
|
||||
except Exception:
|
||||
pass
|
||||
title = f"{RUN_ID} {label} ROC — head={head} class={class_label}"
|
||||
out_path = out_dir / f"{label}_{head}_class{cls}.png"
|
||||
_plot_overlays(fold_curves, title, out_path)
|
||||
print(f"[ok] {out_path}")
|
||||
|
||||
any_holdout_json = False
|
||||
if USE_JSON_ROC:
|
||||
out_dir = run_dir / OUTPUT_SUBDIR
|
||||
for head in HEADS:
|
||||
files = _collect_holdout_json_files(run_dir, head)
|
||||
if files:
|
||||
any_holdout_json = True
|
||||
paired = (
|
||||
_collect_validation_json_files(files, head)
|
||||
if PLOT_VALIDATION_FROM_HOLDOUT_EPOCH
|
||||
else None
|
||||
)
|
||||
if is_binary and FORCE_PROBS_FOR_HOLDOUT_BINARY:
|
||||
_plot_from_probs("holdout", head, out_dir, "_holdout", class_names)
|
||||
else:
|
||||
_plot_set(
|
||||
"holdout",
|
||||
head,
|
||||
files,
|
||||
out_dir,
|
||||
paired,
|
||||
"validation",
|
||||
class_names,
|
||||
)
|
||||
if not any_holdout_json and DEBUG:
|
||||
print("[debug] no JSON ROC files found; falling back to probs")
|
||||
|
||||
if PLOT_BEST_EPOCH:
|
||||
best_out_dir = run_dir / BEST_OUTPUT_SUBDIR
|
||||
for head in HEADS:
|
||||
if is_binary and FORCE_PROBS_FOR_BEST_BINARY:
|
||||
_plot_from_probs("best", head, best_out_dir, "", class_names)
|
||||
continue
|
||||
best_files = _collect_best_json_files(run_dir, head)
|
||||
if best_files:
|
||||
paired = (
|
||||
_collect_holdout_from_validation_files(best_files, head)
|
||||
if PLOT_HOLDOUT_FROM_BEST_EPOCH
|
||||
else None
|
||||
)
|
||||
_plot_set(
|
||||
"best",
|
||||
head,
|
||||
best_files,
|
||||
best_out_dir,
|
||||
paired,
|
||||
"holdout",
|
||||
class_names,
|
||||
)
|
||||
|
||||
# Fallback to probs for holdout plots if JSON wasn't found.
|
||||
if USE_JSON_ROC and any_holdout_json:
|
||||
return
|
||||
|
||||
out_dir = run_dir / OUTPUT_SUBDIR
|
||||
for head in HEADS:
|
||||
curves_by_class: Dict[
|
||||
str, List[Tuple[int, List[float], List[float], float]]
|
||||
] = {}
|
||||
suffix = "_holdout" if USE_HOLDOUT_PROBS else ""
|
||||
files = _collect_prob_files(run_dir, suffix)
|
||||
if not files and USE_HOLDOUT_PROBS and ALLOW_FALLBACK_TO_VALIDATION:
|
||||
suffix = ""
|
||||
files = _collect_prob_files(run_dir, suffix)
|
||||
if files:
|
||||
print(
|
||||
"[warn] Holdout probability dumps not found; using validation probabilities instead."
|
||||
)
|
||||
if not files:
|
||||
raise SystemExit(
|
||||
"No saved probability dumps found. If you want holdout ROC curves, "
|
||||
"run scripts/rebuild_run_best_plots.py with --use-holdout --overwrite "
|
||||
"to generate fold*_y_true_holdout.npy and fold*_probs_*_holdout.npy files."
|
||||
)
|
||||
for fold_idx in sorted(files.keys()):
|
||||
fold_files = files[fold_idx]
|
||||
y_path = fold_files.get("y_true")
|
||||
p_path = fold_files.get(head)
|
||||
if y_path is None or p_path is None:
|
||||
continue
|
||||
y_true = _load_array(y_path)
|
||||
probs = _load_array(p_path)
|
||||
if y_true is None or probs is None:
|
||||
continue
|
||||
curves = _compute_multiclass_curves(y_true, probs)
|
||||
for cls, payload in curves.items():
|
||||
if payload is None:
|
||||
continue
|
||||
fpr, tpr, auc_val = payload
|
||||
curves_by_class.setdefault(cls, []).append(
|
||||
(fold_idx, fpr, tpr, auc_val)
|
||||
)
|
||||
if not curves_by_class:
|
||||
continue
|
||||
classes = sorted(curves_by_class.keys(), key=lambda x: (float(x), str(x)))
|
||||
for cls in classes:
|
||||
fold_curves = curves_by_class.get(cls, [])
|
||||
if not fold_curves:
|
||||
continue
|
||||
class_label = cls
|
||||
if class_names is not None:
|
||||
try:
|
||||
idx = int(cls)
|
||||
if 0 <= idx < len(class_names):
|
||||
class_label = f"{cls} ({class_names[idx]})"
|
||||
except Exception:
|
||||
pass
|
||||
title = f"{RUN_ID} holdout ROC — head={head} class={class_label}"
|
||||
out_path = out_dir / f"holdout_{head}_class{cls}.png"
|
||||
_plot_overlays(fold_curves, title, out_path)
|
||||
print(f"[ok] {out_path}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,514 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Analyze correlations between grid search parameters and performance metrics.
|
||||
|
||||
Example:
|
||||
python scripts/grid_search_analytics/param_perf_correlations.py \
|
||||
--analysis-dir analysis_data/grid_search \
|
||||
--task binary \
|
||||
--metric holdout_auc \
|
||||
--top 30
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Dict, Iterable, List, Optional, Tuple
|
||||
|
||||
DEFAULT_EXCLUDE_KEYS = {
|
||||
"run_id",
|
||||
"fold_metrics",
|
||||
"best_metric",
|
||||
"best_metric_mode",
|
||||
"best_metric_mean",
|
||||
"best_metric_std",
|
||||
"eval_mode",
|
||||
"n_splits",
|
||||
"num_classes",
|
||||
}
|
||||
|
||||
METRIC_MAP = {
|
||||
"auc": "auc_fused",
|
||||
"acc": "acc_fused",
|
||||
"holdout_auc": "holdout_auc_fused",
|
||||
"holdout_acc": "holdout_acc_fused",
|
||||
}
|
||||
|
||||
|
||||
def to_float(value: Optional[object]) -> Optional[float]:
|
||||
if value is None:
|
||||
return None
|
||||
if isinstance(value, (int, float)) and not isinstance(value, bool):
|
||||
num = float(value)
|
||||
if math.isnan(num):
|
||||
return None
|
||||
return num
|
||||
if not isinstance(value, str):
|
||||
return None
|
||||
value = value.strip()
|
||||
if not value:
|
||||
return None
|
||||
try:
|
||||
num = float(value)
|
||||
except ValueError:
|
||||
return None
|
||||
if math.isnan(num):
|
||||
return None
|
||||
return num
|
||||
|
||||
|
||||
def mean(values: List[float]) -> Optional[float]:
|
||||
return (sum(values) / len(values)) if values else None
|
||||
|
||||
|
||||
def read_json(path: Path) -> Optional[Dict[str, object]]:
|
||||
if not path.exists():
|
||||
return None
|
||||
try:
|
||||
data = json.loads(path.read_text())
|
||||
except Exception:
|
||||
return None
|
||||
if not isinstance(data, dict):
|
||||
return None
|
||||
return data
|
||||
|
||||
|
||||
def read_summary(run_dir: Path) -> Optional[Dict[str, object]]:
|
||||
return read_json(run_dir / "summary.json")
|
||||
|
||||
|
||||
def read_cli_args(run_dir: Path) -> Optional[Dict[str, object]]:
|
||||
return read_json(run_dir / "cli_args.json")
|
||||
|
||||
|
||||
def read_run_id(run_dir: Path, summary: Optional[Dict[str, object]]) -> str:
|
||||
if summary:
|
||||
rid = summary.get("run_id")
|
||||
if isinstance(rid, str) and rid:
|
||||
return rid
|
||||
return run_dir.name
|
||||
|
||||
|
||||
def task_from_summary(summary: Optional[Dict[str, object]]) -> Optional[str]:
|
||||
if not summary:
|
||||
return None
|
||||
eval_mode = summary.get("eval_mode")
|
||||
if isinstance(eval_mode, str):
|
||||
mode = eval_mode.strip().lower()
|
||||
if mode == "binary":
|
||||
return "binary"
|
||||
if mode in {"multiclass", "multi", "multi-class"}:
|
||||
return "multiclass"
|
||||
num_classes = summary.get("num_classes")
|
||||
if isinstance(num_classes, (int, float)):
|
||||
return "binary" if int(num_classes) <= 2 else "multiclass"
|
||||
return None
|
||||
|
||||
|
||||
def metric_from_stats(stats: Dict[str, object], metric: str) -> Optional[float]:
|
||||
if metric.startswith("holdout_") and stats.get("holdout_best_monitor") == metric:
|
||||
best_val = to_float(stats.get("holdout_best_so_far"))
|
||||
if best_val is not None:
|
||||
return best_val
|
||||
return to_float(stats.get(metric))
|
||||
|
||||
|
||||
def mean_metric(summary: Dict[str, object], metric: str) -> Optional[float]:
|
||||
folds = summary.get("fold_metrics") or []
|
||||
if not isinstance(folds, list) or not folds:
|
||||
return None
|
||||
values = []
|
||||
for fold in folds:
|
||||
stats = fold.get("stats") if isinstance(fold, dict) else None
|
||||
if not isinstance(stats, dict):
|
||||
return None
|
||||
val = metric_from_stats(stats, metric)
|
||||
if val is None:
|
||||
return None
|
||||
values.append(val)
|
||||
return mean(values)
|
||||
|
||||
|
||||
def flatten_config(data: Dict[str, object],
|
||||
prefix: str = "",
|
||||
exclude_keys: Optional[Iterable[str]] = None) -> Dict[str, object]:
|
||||
out: Dict[str, object] = {}
|
||||
excludes = set(exclude_keys or [])
|
||||
for key, value in data.items():
|
||||
if key in excludes or key.startswith("best_"):
|
||||
continue
|
||||
full_key = f"{prefix}{key}" if not prefix else f"{prefix}.{key}"
|
||||
if isinstance(value, dict):
|
||||
out.update(flatten_config(value, full_key, exclude_keys=excludes))
|
||||
continue
|
||||
if isinstance(value, list):
|
||||
continue
|
||||
out[full_key] = value
|
||||
return out
|
||||
|
||||
|
||||
def rankdata(values: List[float]) -> List[float]:
|
||||
order = sorted(range(len(values)), key=lambda i: values[i])
|
||||
ranks = [0.0] * len(values)
|
||||
i = 0
|
||||
while i < len(values):
|
||||
j = i
|
||||
while j + 1 < len(values) and values[order[j + 1]] == values[order[i]]:
|
||||
j += 1
|
||||
avg_rank = (i + j) / 2.0 + 1.0
|
||||
for k in range(i, j + 1):
|
||||
ranks[order[k]] = avg_rank
|
||||
i = j + 1
|
||||
return ranks
|
||||
|
||||
|
||||
def pearson(x: List[float], y: List[float]) -> Optional[float]:
|
||||
if len(x) != len(y) or len(x) < 2:
|
||||
return None
|
||||
mean_x = sum(x) / len(x)
|
||||
mean_y = sum(y) / len(y)
|
||||
num = sum((xi - mean_x) * (yi - mean_y) for xi, yi in zip(x, y))
|
||||
den_x = sum((xi - mean_x) ** 2 for xi in x)
|
||||
den_y = sum((yi - mean_y) ** 2 for yi in y)
|
||||
if den_x <= 0 or den_y <= 0:
|
||||
return None
|
||||
return num / math.sqrt(den_x * den_y)
|
||||
|
||||
|
||||
def spearman(x: List[float], y: List[float]) -> Optional[float]:
|
||||
rx = rankdata(x)
|
||||
ry = rankdata(y)
|
||||
return pearson(rx, ry)
|
||||
|
||||
|
||||
def correlation_ratio(categories: List[object], values: List[float]) -> Optional[float]:
|
||||
if len(categories) != len(values) or len(values) < 2:
|
||||
return None
|
||||
overall = mean(values)
|
||||
if overall is None:
|
||||
return None
|
||||
total = sum((v - overall) ** 2 for v in values)
|
||||
if total <= 0:
|
||||
return None
|
||||
sums: Dict[object, List[float]] = {}
|
||||
for cat, val in zip(categories, values):
|
||||
sums.setdefault(cat, []).append(val)
|
||||
between = 0.0
|
||||
for vals in sums.values():
|
||||
avg = mean(vals)
|
||||
if avg is None:
|
||||
continue
|
||||
between += len(vals) * (avg - overall) ** 2
|
||||
return math.sqrt(between / total)
|
||||
|
||||
|
||||
def format_value(value: object) -> str:
|
||||
if value is None:
|
||||
return ""
|
||||
if isinstance(value, bool):
|
||||
return "true" if value else "false"
|
||||
if isinstance(value, int):
|
||||
return str(value)
|
||||
if isinstance(value, float):
|
||||
return f"{value:.6g}"
|
||||
return str(value)
|
||||
|
||||
|
||||
def format_metric(value: Optional[float]) -> str:
|
||||
if value is None:
|
||||
return ""
|
||||
return f"{value:.4f}"
|
||||
|
||||
|
||||
def format_table(rows: List[Dict[str, object]], columns: List[str]) -> str:
|
||||
col_widths = {
|
||||
col: max(len(col), max((len(str(row.get(col, ""))) for row in rows), default=0))
|
||||
for col in columns
|
||||
}
|
||||
header = " | ".join(col.ljust(col_widths[col]) for col in columns)
|
||||
divider = "-+-".join("-" * col_widths[col] for col in columns)
|
||||
body = [
|
||||
" | ".join(str(row.get(col, "")).ljust(col_widths[col]) for col in columns)
|
||||
for row in rows
|
||||
]
|
||||
return "\n".join([header, divider, *body])
|
||||
|
||||
|
||||
def render_progress(current: int, total: Optional[int], matched: int) -> str:
|
||||
if total:
|
||||
width = 30
|
||||
filled = int(width * current / total)
|
||||
bar = "#" * filled + "-" * (width - filled)
|
||||
return f"[{bar}] {current}/{total} matched {matched}"
|
||||
return f"Scanned {current} dirs, matched {matched}"
|
||||
|
||||
|
||||
def iter_run_dirs(root: Path, shallow: bool, show_progress: bool) -> Iterable[Path]:
|
||||
if shallow:
|
||||
entries = [entry for entry in root.iterdir() if entry.is_dir()]
|
||||
entries.sort(key=lambda p: p.name)
|
||||
total = len(entries)
|
||||
matched = 0
|
||||
last_update = 0.0
|
||||
for idx, entry in enumerate(entries, start=1):
|
||||
if show_progress:
|
||||
now = time.monotonic()
|
||||
if now - last_update >= 0.1 or idx == total:
|
||||
msg = render_progress(idx, total, matched)
|
||||
print(f"\rScanning {msg}", end="", file=sys.stderr, flush=True)
|
||||
last_update = now
|
||||
if (entry / "summary.json").is_file():
|
||||
matched += 1
|
||||
yield entry
|
||||
if show_progress:
|
||||
print(file=sys.stderr)
|
||||
return
|
||||
|
||||
matched = 0
|
||||
scanned = 0
|
||||
last_update = 0.0
|
||||
for dirpath, dirnames, filenames in os.walk(root):
|
||||
scanned += 1
|
||||
if show_progress:
|
||||
now = time.monotonic()
|
||||
if now - last_update >= 0.2:
|
||||
msg = render_progress(scanned, None, matched)
|
||||
print(f"\rScanning {msg}", end="", file=sys.stderr, flush=True)
|
||||
last_update = now
|
||||
if "summary.json" in filenames:
|
||||
matched += 1
|
||||
yield Path(dirpath)
|
||||
if show_progress:
|
||||
msg = render_progress(scanned, None, matched)
|
||||
print(f"\rScanning {msg}", end="", file=sys.stderr, flush=True)
|
||||
print(file=sys.stderr)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
ap = argparse.ArgumentParser(description="Correlate grid search parameters with performance.")
|
||||
ap.add_argument("--analysis-dir", type=Path, default=Path("analysis_data/grid_search"),
|
||||
help="Directory containing run subdirectories")
|
||||
ap.add_argument("--task", choices=["binary", "multiclass", "all"], default="all",
|
||||
help="Filter runs by task type (default: all)")
|
||||
ap.add_argument("--metric", choices=sorted(METRIC_MAP.keys()), default="holdout_auc",
|
||||
help="Performance metric to analyze (default: holdout_auc)")
|
||||
ap.add_argument("--sort-by", choices=["score", "abs_rho", "rho", "r", "eta"], default="score",
|
||||
help="Sorting key for results (default: score)")
|
||||
ap.add_argument("--asc", action="store_true",
|
||||
help="Sort ascending (default: descending)")
|
||||
ap.add_argument("--desc", action="store_true",
|
||||
help="Sort descending (default: descending)")
|
||||
ap.add_argument("--top", type=int, default=30,
|
||||
help="Limit output to top N parameters (default: 30)")
|
||||
ap.add_argument("--params", default=None,
|
||||
help="Comma-separated list of parameter keys to include")
|
||||
ap.add_argument("--exclude", default=None,
|
||||
help="Comma-separated list of parameter keys to exclude")
|
||||
ap.add_argument("--min-count", type=int, default=10,
|
||||
help="Minimum runs required to analyze a parameter (default: 10)")
|
||||
ap.add_argument("--min-unique", type=int, default=2,
|
||||
help="Minimum unique values required (default: 2)")
|
||||
ap.add_argument("--match", default=None,
|
||||
help="Only include run directories whose name contains this substring")
|
||||
ap.add_argument("--shallow", action="store_true",
|
||||
help="Only scan directories directly under analysis-dir")
|
||||
ap.add_argument("--no-progress", action="store_true",
|
||||
help="Disable progress output")
|
||||
args = ap.parse_args()
|
||||
|
||||
if args.asc and args.desc:
|
||||
raise SystemExit("Choose only one of --asc or --desc.")
|
||||
|
||||
root = args.analysis_dir
|
||||
if not root.exists():
|
||||
raise SystemExit(f"Analysis directory not found: {root}")
|
||||
|
||||
exclude_keys = set(DEFAULT_EXCLUDE_KEYS)
|
||||
if args.exclude:
|
||||
for item in args.exclude.split(","):
|
||||
item = item.strip()
|
||||
if item:
|
||||
exclude_keys.add(item)
|
||||
|
||||
runs: List[Dict[str, object]] = []
|
||||
values_by_key: Dict[str, List[object]] = {}
|
||||
missing_summary = 0
|
||||
unknown_task = 0
|
||||
missing_cli = 0
|
||||
|
||||
metric_key = METRIC_MAP[args.metric]
|
||||
|
||||
for run_dir in iter_run_dirs(root, shallow=args.shallow, show_progress=not args.no_progress):
|
||||
if args.match and args.match not in run_dir.name:
|
||||
continue
|
||||
summary = read_summary(run_dir)
|
||||
if summary is None:
|
||||
missing_summary += 1
|
||||
continue
|
||||
task_label = task_from_summary(summary)
|
||||
if args.task != "all":
|
||||
if task_label is None:
|
||||
unknown_task += 1
|
||||
continue
|
||||
if task_label != args.task:
|
||||
continue
|
||||
|
||||
metric_value = mean_metric(summary, metric_key)
|
||||
if metric_value is None:
|
||||
continue
|
||||
|
||||
cli_args = read_cli_args(run_dir)
|
||||
if cli_args is None:
|
||||
missing_cli += 1
|
||||
config_source = cli_args if cli_args is not None else summary
|
||||
config = flatten_config(config_source, exclude_keys=exclude_keys)
|
||||
|
||||
run_id = read_run_id(run_dir, summary)
|
||||
runs.append({
|
||||
"run_id": run_id,
|
||||
"metric": metric_value,
|
||||
"config": config,
|
||||
})
|
||||
for key, value in config.items():
|
||||
values_by_key.setdefault(key, []).append(value)
|
||||
|
||||
if not runs:
|
||||
print("No matching runs found.")
|
||||
return
|
||||
|
||||
if args.params:
|
||||
param_keys = [p.strip() for p in args.params.split(",") if p.strip()]
|
||||
else:
|
||||
param_keys = []
|
||||
for key, values in values_by_key.items():
|
||||
unique_values = {format_value(v) for v in values}
|
||||
if len(unique_values) >= args.min_unique:
|
||||
param_keys.append(key)
|
||||
param_keys.sort()
|
||||
|
||||
rows: List[Dict[str, object]] = []
|
||||
for key in param_keys:
|
||||
values = []
|
||||
metrics = []
|
||||
for run in runs:
|
||||
config = run.get("config", {})
|
||||
if key not in config:
|
||||
continue
|
||||
values.append(config[key])
|
||||
metrics.append(run["metric"])
|
||||
|
||||
if len(values) < args.min_count:
|
||||
continue
|
||||
|
||||
unique_values = {format_value(v) for v in values}
|
||||
if len(unique_values) < args.min_unique:
|
||||
continue
|
||||
|
||||
numeric_values: List[float] = []
|
||||
numeric_ok = True
|
||||
for v in values:
|
||||
num = to_float(v)
|
||||
if num is None or isinstance(v, bool):
|
||||
numeric_ok = False
|
||||
break
|
||||
numeric_values.append(num)
|
||||
|
||||
groups: Dict[object, List[float]] = {}
|
||||
for val, metric in zip(values, metrics):
|
||||
groups.setdefault(val, []).append(metric)
|
||||
group_means = {k: mean(v) for k, v in groups.items()}
|
||||
best_group = max(group_means.items(), key=lambda item: item[1] or float("-inf"))
|
||||
worst_group = min(group_means.items(), key=lambda item: item[1] or float("inf"))
|
||||
|
||||
if numeric_ok and len(set(numeric_values)) >= 3:
|
||||
rho = spearman(numeric_values, metrics)
|
||||
r = pearson(numeric_values, metrics)
|
||||
score = abs(rho) if rho is not None else None
|
||||
row = {
|
||||
"param": key,
|
||||
"type": "numeric",
|
||||
"n": len(values),
|
||||
"distinct": len(unique_values),
|
||||
"score": format_metric(score) if score is not None else "",
|
||||
"rho": format_metric(rho),
|
||||
"r": format_metric(r),
|
||||
"best_value": format_value(best_group[0]),
|
||||
"best_mean": format_metric(best_group[1]),
|
||||
"worst_value": format_value(worst_group[0]),
|
||||
"worst_mean": format_metric(worst_group[1]),
|
||||
}
|
||||
else:
|
||||
eta = correlation_ratio(values, metrics)
|
||||
score = eta
|
||||
row = {
|
||||
"param": key,
|
||||
"type": "categorical",
|
||||
"n": len(values),
|
||||
"distinct": len(unique_values),
|
||||
"score": format_metric(score) if score is not None else "",
|
||||
"rho": "",
|
||||
"r": "",
|
||||
"best_value": format_value(best_group[0]),
|
||||
"best_mean": format_metric(best_group[1]),
|
||||
"worst_value": format_value(worst_group[0]),
|
||||
"worst_mean": format_metric(worst_group[1]),
|
||||
}
|
||||
|
||||
rows.append(row)
|
||||
|
||||
if not rows:
|
||||
print("No parameters met the minimum requirements.")
|
||||
return
|
||||
|
||||
def sort_key(row: Dict[str, object]) -> float:
|
||||
raw = row.get(args.sort_by)
|
||||
if isinstance(raw, str):
|
||||
val = to_float(raw)
|
||||
else:
|
||||
val = to_float(raw)
|
||||
if val is None:
|
||||
return float("inf") if args.asc else float("-inf")
|
||||
return float(val)
|
||||
|
||||
rows.sort(key=sort_key, reverse=not args.asc)
|
||||
if args.top is not None:
|
||||
rows = rows[:args.top]
|
||||
|
||||
header_lines = []
|
||||
header_lines.append(f"Metric: {args.metric} (mean over folds)")
|
||||
if args.task != "all":
|
||||
header_lines.append(f"Task filter: {args.task}")
|
||||
if args.match:
|
||||
header_lines.append(f"Name filter: {args.match}")
|
||||
header_lines.append(f"Runs: {len(runs)}")
|
||||
if missing_cli:
|
||||
header_lines.append(f"Missing cli_args: {missing_cli}")
|
||||
if missing_summary or unknown_task:
|
||||
header_lines.append(f"Skipped: {missing_summary} missing summary, {unknown_task} unknown task")
|
||||
header_lines.append("")
|
||||
print("\n".join(header_lines))
|
||||
|
||||
columns = [
|
||||
"param",
|
||||
"type",
|
||||
"n",
|
||||
"distinct",
|
||||
"score",
|
||||
"rho",
|
||||
"r",
|
||||
"best_value",
|
||||
"best_mean",
|
||||
"worst_value",
|
||||
"worst_mean",
|
||||
]
|
||||
print(format_table(rows, columns))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,212 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Re-run a single grid-search configuration into analysis_data/re_runs."""
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import csv
|
||||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import Dict, List
|
||||
|
||||
|
||||
# ---------------------------
|
||||
# Config (edit in IDE)
|
||||
# ---------------------------
|
||||
RUN_ID = "20251129-0063" # fallback if --run-number is not provided
|
||||
OUTPUT_RUN_ID = RUN_ID # fallback output run id
|
||||
SHORTNAME = "re_runs" # output root under analysis_data/ and models/
|
||||
GRID_PLAN = Path("analysis_data/grid_search/grid_plan.csv")
|
||||
MANIFEST = Path("manifest.csv")
|
||||
RUN_SCRIPT = Path("scripts/run_multifold.py")
|
||||
REBUILD_SCRIPT = Path("scripts/rebuild_run_best_plots.py")
|
||||
PLOT_HEADS = ["fused", "image", "metadata"]
|
||||
USE_HOLDOUT_BEST_FOR_PLOTS = True
|
||||
OVERWRITE_HOLDOUT_PROBS = True
|
||||
ALLOW_EXISTING_RUN_DIR = False
|
||||
|
||||
|
||||
def _parse_args() -> argparse.Namespace:
|
||||
ap = argparse.ArgumentParser(description="Re-run a single grid-search item.")
|
||||
ap.add_argument(
|
||||
"--run-number",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Last 4 digits of run_id (e.g., 0063).",
|
||||
)
|
||||
ap.add_argument(
|
||||
"--output-run-id",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Optional output run id; defaults to matched run_id.",
|
||||
)
|
||||
return ap.parse_args()
|
||||
|
||||
|
||||
def _read_plan(path: Path) -> List[Dict[str, str]]:
|
||||
if not path.exists():
|
||||
raise FileNotFoundError(f"Grid plan not found: {path}")
|
||||
with path.open(newline="") as fh:
|
||||
reader = csv.DictReader(fh)
|
||||
return list(reader)
|
||||
|
||||
|
||||
def _find_row(rows: List[Dict[str, str]], run_id: str) -> Dict[str, str]:
|
||||
for row in rows:
|
||||
if row.get("run_id") == run_id:
|
||||
return row
|
||||
raise ValueError(f"run_id not found in grid plan: {run_id}")
|
||||
|
||||
|
||||
def _resolve_run_id(rows: List[Dict[str, str]], run_number: str | None) -> str:
|
||||
if not run_number:
|
||||
return RUN_ID
|
||||
run_number = str(run_number).strip()
|
||||
if run_number.isdigit():
|
||||
run_number = run_number.zfill(4)
|
||||
matches = [
|
||||
r.get("run_id", "")
|
||||
for r in rows
|
||||
if str(r.get("run_id", "")).endswith(f"-{run_number}")
|
||||
]
|
||||
if len(matches) == 1:
|
||||
return matches[0]
|
||||
if len(matches) > 1:
|
||||
raise ValueError(
|
||||
f"Multiple run_ids matched run-number '{run_number}': {matches[:5]}{' ...' if len(matches) > 5 else ''}"
|
||||
)
|
||||
raise ValueError(f"No run_id found ending with '-{run_number}'")
|
||||
|
||||
|
||||
def _build_run_command(row: Dict[str, str], output_run_id: str) -> List[str]:
|
||||
cmd = [
|
||||
sys.executable,
|
||||
str(RUN_SCRIPT),
|
||||
"--backbone",
|
||||
"resnet50",
|
||||
"--fusion-mode",
|
||||
"fused",
|
||||
"--epochs",
|
||||
"40",
|
||||
"--batch-size",
|
||||
"8",
|
||||
"--img-crop-manifest",
|
||||
str(MANIFEST),
|
||||
"--img-crop-weights",
|
||||
row["crop_weights"],
|
||||
"--img-crop-normalize",
|
||||
row["crop_normalize"],
|
||||
"--eval_mode",
|
||||
row["eval_mode"],
|
||||
"--holdout-per-class",
|
||||
"12",
|
||||
"--run-id",
|
||||
output_run_id,
|
||||
"--shortname",
|
||||
SHORTNAME,
|
||||
]
|
||||
if row.get("crop_tta") == "True":
|
||||
cmd.append("--img-crop-tta")
|
||||
|
||||
loss_mode = row.get("loss_mode")
|
||||
if loss_mode == "focal":
|
||||
cmd.extend(["--focal-gamma", "2.0"])
|
||||
elif loss_mode == "balanced":
|
||||
cmd.append("--balanced-sampler")
|
||||
|
||||
thaw_mode = row.get("thaw_mode")
|
||||
if thaw_mode == "gradual":
|
||||
cmd.append("--gradual-thaw")
|
||||
cmd.extend(["--thaw-ratio", "0.33"])
|
||||
cmd.extend(["--thaw-start-epoch", "10"])
|
||||
cmd.extend(["--thaw-target", "image"])
|
||||
|
||||
se_mode = row.get("se_mode")
|
||||
if se_mode == "none":
|
||||
cmd.append("--no-se")
|
||||
else:
|
||||
cmd.extend(["--se-reduction", "16"])
|
||||
cmd.extend(["--se-reduction-tower", "16"])
|
||||
cmd.extend(["--se-where", se_mode])
|
||||
bridge_pre = row.get("se_bridge_pre_norm")
|
||||
tower_pre = row.get("se_tower_pre_norm")
|
||||
if bridge_pre == "True":
|
||||
cmd.append("--se-pre-norm")
|
||||
elif bridge_pre == "False":
|
||||
cmd.append("--no-se-pre-norm")
|
||||
if tower_pre == "True":
|
||||
cmd.append("--se-pre-norm-tower")
|
||||
elif tower_pre == "False":
|
||||
cmd.append("--no-se-pre-norm-tower")
|
||||
|
||||
return cmd
|
||||
|
||||
|
||||
def _swap_in_holdout_best(models_dir: Path) -> None:
|
||||
for fold_dir in sorted(models_dir.glob("fold*")):
|
||||
if not fold_dir.is_dir():
|
||||
continue
|
||||
holdout_best = fold_dir / "model_holdout_best.pt"
|
||||
model_best = fold_dir / "model_best.pt"
|
||||
if not holdout_best.exists():
|
||||
print(f"[warn] {holdout_best} missing; skipping.")
|
||||
continue
|
||||
if model_best.exists():
|
||||
backup = fold_dir / "model_best_from_train.pt"
|
||||
if not backup.exists():
|
||||
try:
|
||||
shutil.copy2(model_best, backup)
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
shutil.copy2(holdout_best, model_best)
|
||||
except Exception as exc:
|
||||
print(f"[warn] failed to replace {model_best}: {exc}")
|
||||
|
||||
|
||||
def _run_rebuild(run_dir: Path) -> None:
|
||||
for head in PLOT_HEADS:
|
||||
cmd = [
|
||||
sys.executable,
|
||||
str(REBUILD_SCRIPT),
|
||||
"--run-dir",
|
||||
str(run_dir),
|
||||
"--head",
|
||||
head,
|
||||
"--use-holdout",
|
||||
]
|
||||
if OVERWRITE_HOLDOUT_PROBS:
|
||||
cmd.append("--overwrite")
|
||||
print("[rerun] Rebuilding holdout ROC plots:", " ".join(cmd))
|
||||
subprocess.run(cmd, check=True)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = _parse_args()
|
||||
rows = _read_plan(GRID_PLAN)
|
||||
run_id = _resolve_run_id(rows, args.run_number)
|
||||
row = _find_row(rows, run_id)
|
||||
output_run_id = args.output_run_id or run_id
|
||||
|
||||
run_dir = Path("analysis_data") / SHORTNAME / output_run_id
|
||||
if run_dir.exists() and not ALLOW_EXISTING_RUN_DIR:
|
||||
raise SystemExit(
|
||||
f"Run directory already exists: {run_dir} (set ALLOW_EXISTING_RUN_DIR=True to reuse)"
|
||||
)
|
||||
|
||||
cmd = _build_run_command(row, output_run_id=output_run_id)
|
||||
print("[rerun] Launching:", " ".join(cmd))
|
||||
subprocess.run(cmd, check=True)
|
||||
|
||||
models_dir = Path("models") / SHORTNAME / output_run_id
|
||||
if USE_HOLDOUT_BEST_FOR_PLOTS:
|
||||
print("[rerun] Swapping in holdout-best checkpoints for plotting.")
|
||||
_swap_in_holdout_best(models_dir)
|
||||
|
||||
_run_rebuild(run_dir)
|
||||
print(f"[rerun] Done. Outputs in {run_dir}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,215 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Re-run a single grid-search configuration using the V2 loader pipeline."""
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import csv
|
||||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import Dict, List
|
||||
|
||||
|
||||
# ---------------------------
|
||||
# Config (edit in IDE)
|
||||
# ---------------------------
|
||||
RUN_ID = "20251129-0063" # fallback if --run-number is not provided
|
||||
OUTPUT_RUN_ID = RUN_ID # fallback output run id
|
||||
SHORTNAME = "re_runs_v2" # output root under analysis_data/ and models/
|
||||
GRID_PLAN = Path("analysis_data/grid_search/grid_plan.csv")
|
||||
MANIFEST = Path("manifest.csv")
|
||||
RUN_SCRIPT = Path("scripts/run_multifold_v2.py")
|
||||
REBUILD_SCRIPT = Path("scripts/rebuild_run_best_plots.py")
|
||||
PLOT_HEADS = ["fused", "image", "metadata"]
|
||||
USE_HOLDOUT_BEST_FOR_PLOTS = True
|
||||
OVERWRITE_HOLDOUT_PROBS = True
|
||||
ALLOW_EXISTING_RUN_DIR = False
|
||||
SAMPLE_MODE = "eye" # eye-level for parity with v1 grid runs
|
||||
|
||||
|
||||
def _parse_args() -> argparse.Namespace:
|
||||
ap = argparse.ArgumentParser(description="Re-run a single grid-search item with V2 loaders.")
|
||||
ap.add_argument(
|
||||
"--run-number",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Last 4 digits of run_id (e.g., 0063).",
|
||||
)
|
||||
ap.add_argument(
|
||||
"--output-run-id",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Optional output run id; defaults to matched run_id.",
|
||||
)
|
||||
return ap.parse_args()
|
||||
|
||||
|
||||
def _read_plan(path: Path) -> List[Dict[str, str]]:
|
||||
if not path.exists():
|
||||
raise FileNotFoundError(f"Grid plan not found: {path}")
|
||||
with path.open(newline="") as fh:
|
||||
reader = csv.DictReader(fh)
|
||||
return list(reader)
|
||||
|
||||
|
||||
def _find_row(rows: List[Dict[str, str]], run_id: str) -> Dict[str, str]:
|
||||
for row in rows:
|
||||
if row.get("run_id") == run_id:
|
||||
return row
|
||||
raise ValueError(f"run_id not found in grid plan: {run_id}")
|
||||
|
||||
|
||||
def _resolve_run_id(rows: List[Dict[str, str]], run_number: str | None) -> str:
|
||||
if not run_number:
|
||||
return RUN_ID
|
||||
run_number = str(run_number).strip()
|
||||
if run_number.isdigit():
|
||||
run_number = run_number.zfill(4)
|
||||
matches = [
|
||||
r.get("run_id", "")
|
||||
for r in rows
|
||||
if str(r.get("run_id", "")).endswith(f"-{run_number}")
|
||||
]
|
||||
if len(matches) == 1:
|
||||
return matches[0]
|
||||
if len(matches) > 1:
|
||||
raise ValueError(
|
||||
f"Multiple run_ids matched run-number '{run_number}': {matches[:5]}{' ...' if len(matches) > 5 else ''}"
|
||||
)
|
||||
raise ValueError(f"No run_id found ending with '-{run_number}'")
|
||||
|
||||
|
||||
def _build_run_command(row: Dict[str, str], output_run_id: str) -> List[str]:
|
||||
cmd = [
|
||||
sys.executable,
|
||||
str(RUN_SCRIPT),
|
||||
"--backbone",
|
||||
"resnet50",
|
||||
"--fusion-mode",
|
||||
"fused",
|
||||
"--epochs",
|
||||
"40",
|
||||
"--batch-size",
|
||||
"8",
|
||||
"--img-crop-manifest",
|
||||
str(MANIFEST),
|
||||
"--img-crop-weights",
|
||||
row["crop_weights"],
|
||||
"--img-crop-normalize",
|
||||
row["crop_normalize"],
|
||||
"--eval_mode",
|
||||
row["eval_mode"],
|
||||
"--holdout-per-class",
|
||||
"12",
|
||||
"--run-id",
|
||||
output_run_id,
|
||||
"--shortname",
|
||||
SHORTNAME,
|
||||
"--sample-mode",
|
||||
SAMPLE_MODE,
|
||||
]
|
||||
if row.get("crop_tta") == "True":
|
||||
cmd.append("--img-crop-tta")
|
||||
|
||||
loss_mode = row.get("loss_mode")
|
||||
if loss_mode == "focal":
|
||||
cmd.extend(["--focal-gamma", "2.0"])
|
||||
elif loss_mode == "balanced":
|
||||
cmd.append("--balanced-sampler")
|
||||
|
||||
thaw_mode = row.get("thaw_mode")
|
||||
if thaw_mode == "gradual":
|
||||
cmd.append("--gradual-thaw")
|
||||
cmd.extend(["--thaw-ratio", "0.33"])
|
||||
cmd.extend(["--thaw-start-epoch", "10"])
|
||||
cmd.extend(["--thaw-target", "image"])
|
||||
|
||||
se_mode = row.get("se_mode")
|
||||
if se_mode == "none":
|
||||
cmd.append("--no-se")
|
||||
else:
|
||||
cmd.extend(["--se-reduction", "16"])
|
||||
cmd.extend(["--se-reduction-tower", "16"])
|
||||
cmd.extend(["--se-where", se_mode])
|
||||
bridge_pre = row.get("se_bridge_pre_norm")
|
||||
tower_pre = row.get("se_tower_pre_norm")
|
||||
if bridge_pre == "True":
|
||||
cmd.append("--se-pre-norm")
|
||||
elif bridge_pre == "False":
|
||||
cmd.append("--no-se-pre-norm")
|
||||
if tower_pre == "True":
|
||||
cmd.append("--se-pre-norm-tower")
|
||||
elif tower_pre == "False":
|
||||
cmd.append("--no-se-pre-norm-tower")
|
||||
|
||||
return cmd
|
||||
|
||||
|
||||
def _swap_in_holdout_best(models_dir: Path) -> None:
|
||||
for fold_dir in sorted(models_dir.glob("fold*")):
|
||||
if not fold_dir.is_dir():
|
||||
continue
|
||||
holdout_best = fold_dir / "model_holdout_best.pt"
|
||||
model_best = fold_dir / "model_best.pt"
|
||||
if not holdout_best.exists():
|
||||
print(f"[warn] {holdout_best} missing; skipping.")
|
||||
continue
|
||||
if model_best.exists():
|
||||
backup = fold_dir / "model_best_from_train.pt"
|
||||
if not backup.exists():
|
||||
try:
|
||||
shutil.copy2(model_best, backup)
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
shutil.copy2(holdout_best, model_best)
|
||||
except Exception as exc:
|
||||
print(f"[warn] failed to replace {model_best}: {exc}")
|
||||
|
||||
|
||||
def _run_rebuild(run_dir: Path) -> None:
|
||||
for head in PLOT_HEADS:
|
||||
cmd = [
|
||||
sys.executable,
|
||||
str(REBUILD_SCRIPT),
|
||||
"--run-dir",
|
||||
str(run_dir),
|
||||
"--head",
|
||||
head,
|
||||
"--use-holdout",
|
||||
]
|
||||
if OVERWRITE_HOLDOUT_PROBS:
|
||||
cmd.append("--overwrite")
|
||||
print("[rerun] Rebuilding holdout ROC plots:", " ".join(cmd))
|
||||
subprocess.run(cmd, check=True)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = _parse_args()
|
||||
rows = _read_plan(GRID_PLAN)
|
||||
run_id = _resolve_run_id(rows, args.run_number)
|
||||
row = _find_row(rows, run_id)
|
||||
output_run_id = args.output_run_id or run_id
|
||||
|
||||
run_dir = Path("analysis_data") / SHORTNAME / output_run_id
|
||||
if run_dir.exists() and not ALLOW_EXISTING_RUN_DIR:
|
||||
raise SystemExit(
|
||||
f"Run directory already exists: {run_dir} (set ALLOW_EXISTING_RUN_DIR=True to reuse)"
|
||||
)
|
||||
|
||||
cmd = _build_run_command(row, output_run_id=output_run_id)
|
||||
print("[rerun] Launching:", " ".join(cmd))
|
||||
subprocess.run(cmd, check=True)
|
||||
|
||||
models_dir = Path("models") / SHORTNAME / output_run_id
|
||||
if USE_HOLDOUT_BEST_FOR_PLOTS:
|
||||
print("[rerun] Swapping in holdout-best checkpoints for plotting.")
|
||||
_swap_in_holdout_best(models_dir)
|
||||
|
||||
_run_rebuild(run_dir)
|
||||
print(f"[rerun] Done. Outputs in {run_dir}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,101 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
from pathlib import Path
|
||||
|
||||
from scripts.grid_search_analytics.derived_analysis import derived_analysis
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
ap = argparse.ArgumentParser(
|
||||
description="Generate derived grid-search analytics artifacts (fusion/error + param-performance)."
|
||||
)
|
||||
ap.add_argument("--analysis-dir", default="analysis_data/grid_search")
|
||||
ap.add_argument("--mode", choices=["binary", "multiclass"], default="multiclass")
|
||||
ap.add_argument("--method", choices=["pearson", "spearman"], default="spearman")
|
||||
ap.add_argument(
|
||||
"--x-metric",
|
||||
choices=["fusion_corrections", "fusion_corrections_per_opportunity"],
|
||||
default="fusion_corrections_per_opportunity",
|
||||
help="Fusion-correlation x-axis metric for summary bar plot.",
|
||||
)
|
||||
ap.add_argument(
|
||||
"--cat-method",
|
||||
choices=["eta", "anova", "kruskal"],
|
||||
default="kruskal",
|
||||
help="Categorical-test method for param-performance correlations.",
|
||||
)
|
||||
ap.add_argument("--top-n", type=int, default=None, help="Optional cap for per-run plots.")
|
||||
ap.add_argument(
|
||||
"--recompute",
|
||||
action="store_true",
|
||||
help="Recompute from run artifacts instead of preferring cached CSVs.",
|
||||
)
|
||||
ap.add_argument(
|
||||
"--deep-scan",
|
||||
action="store_true",
|
||||
help="Scan nested directories instead of direct children only.",
|
||||
)
|
||||
return ap.parse_args()
|
||||
|
||||
|
||||
def main() -> int:
|
||||
args = parse_args()
|
||||
analysis = derived_analysis(
|
||||
Path(args.analysis_dir),
|
||||
classification_mode=args.mode,
|
||||
)
|
||||
|
||||
shallow = not args.deep_scan
|
||||
existing = not args.recompute
|
||||
|
||||
analysis.identify_fusion_corrections(shallow=shallow, existing=existing)
|
||||
analysis.populate_primary_metrics(shallow=shallow, existing=existing)
|
||||
|
||||
analysis.write_fusion_corrections()
|
||||
analysis.write_fusion_errors()
|
||||
analysis.write_primary_metrics()
|
||||
|
||||
analysis.plot_fusion_corrections_errors(
|
||||
shallow=shallow, existing=existing, top_n=args.top_n
|
||||
)
|
||||
analysis.plot_conf_delta_boxplot(
|
||||
shallow=shallow, existing=existing, top_n=args.top_n
|
||||
)
|
||||
|
||||
corr_df = analysis.param_performance_correlations(
|
||||
shallow=shallow,
|
||||
existing=existing,
|
||||
method=args.method,
|
||||
cat_method=args.cat_method,
|
||||
)
|
||||
analysis.plot_param_perf_corr_panels(corr_df)
|
||||
|
||||
corr_acc = analysis.fusion_corrections_correlation(
|
||||
method=args.method, metric_type="acc"
|
||||
)
|
||||
corr_auc = analysis.fusion_corrections_correlation(
|
||||
method=args.method, metric_type="auc"
|
||||
)
|
||||
try:
|
||||
analysis.plot_fusion_perf_summary(
|
||||
corr_acc,
|
||||
corr_auc,
|
||||
method=args.method,
|
||||
x_metric=args.x_metric,
|
||||
)
|
||||
except RuntimeError:
|
||||
analysis.plot_fusion_perf_summary(
|
||||
corr_acc,
|
||||
corr_auc,
|
||||
method=args.method,
|
||||
x_metric="fusion_corrections",
|
||||
)
|
||||
|
||||
print(f"Done. Outputs written under: {Path(args.analysis_dir) / 'plots'}")
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -1,269 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Compare Perkins→Pneumatic IOP conversion methods:
|
||||
- Current: fixed ratio (1.158)
|
||||
- OLS linear regression (minimises MSE)
|
||||
- LAD linear regression (minimises MAE — robust to outliers)
|
||||
- Multiple regression: Perkins + Pachymetry (OLS)
|
||||
|
||||
Also reports on the impact of dropping IOP_raw.
|
||||
|
||||
Run from repo root:
|
||||
python scripts/exploratory/iop_correction_analysis.py
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
from scipy import stats
|
||||
import matplotlib.pyplot as plt
|
||||
import matplotlib.gridspec as gridspec
|
||||
|
||||
REPO = Path(__file__).resolve().parents[2]
|
||||
sys.path.insert(0, str(REPO))
|
||||
|
||||
CURRENT_RATIO = 1.158
|
||||
|
||||
# ── load raw clinical data ────────────────────────────────────────────────────
|
||||
|
||||
def load_raw() -> pd.DataFrame:
|
||||
dfs = []
|
||||
for fname in ("patient_data_od.xlsx", "patient_data_os.xlsx"):
|
||||
df = pd.read_excel(REPO / "Papila/ClinicalData" / fname, header=1)
|
||||
eye = "OD" if "od" in fname else "OS"
|
||||
df["eyeID"] = eye
|
||||
dfs.append(df)
|
||||
return pd.concat(dfs, ignore_index=True)
|
||||
|
||||
df = load_raw()
|
||||
print(f"Total rows: {len(df)}")
|
||||
print(f"Columns with IOP: {[c for c in df.columns if 'iop' in c.lower() or c in ('Pneumatic','Perkins')]}")
|
||||
|
||||
# ── find paired rows ──────────────────────────────────────────────────────────
|
||||
|
||||
paired = df.dropna(subset=["Pneumatic", "Perkins", "Pachymetry"]).copy()
|
||||
pneumatic = paired["Pneumatic"].values.astype(float)
|
||||
perkins = paired["Perkins"].values.astype(float)
|
||||
pachymetry = paired["Pachymetry"].values.astype(float)
|
||||
|
||||
print(f"\nPaired obs (Pneumatic + Perkins + Pachymetry): n={len(paired)}")
|
||||
print(f" Pneumatic: mean={pneumatic.mean():.2f} std={pneumatic.std():.2f} "
|
||||
f"range=[{pneumatic.min():.1f}, {pneumatic.max():.1f}]")
|
||||
print(f" Perkins: mean={perkins.mean():.2f} std={perkins.std():.2f} "
|
||||
f"range=[{perkins.min():.1f}, {perkins.max():.1f}]")
|
||||
|
||||
# ── current ratio ─────────────────────────────────────────────────────────────
|
||||
|
||||
ratio_obs = pneumatic / perkins
|
||||
ratio_mean = ratio_obs.mean()
|
||||
ratio_pred = perkins * CURRENT_RATIO
|
||||
ratio_resid = pneumatic - ratio_pred
|
||||
ratio_mae = np.abs(ratio_resid).mean()
|
||||
ratio_rmse = np.sqrt((ratio_resid ** 2).mean())
|
||||
|
||||
print(f"\n── Current ratio approach ────────────────────────────────")
|
||||
print(f" Observed Pneumatic/Perkins ratio: mean={ratio_mean:.4f} "
|
||||
f"std={ratio_obs.std():.4f} range=[{ratio_obs.min():.3f}, {ratio_obs.max():.3f}]")
|
||||
print(f" Hardcoded ratio used: {CURRENT_RATIO}")
|
||||
print(f" MAE: {ratio_mae:.3f} mmHg")
|
||||
print(f" RMSE: {ratio_rmse:.3f} mmHg")
|
||||
|
||||
# ── OLS regression ────────────────────────────────────────────────────────────
|
||||
|
||||
slope, intercept, r, p, se = stats.linregress(perkins, pneumatic)
|
||||
reg_pred = slope * perkins + intercept
|
||||
reg_resid = pneumatic - reg_pred
|
||||
reg_mae = np.abs(reg_resid).mean()
|
||||
reg_rmse = np.sqrt((reg_resid ** 2).mean())
|
||||
|
||||
print(f"\n── OLS regression (minimises MSE): Pneumatic = slope * Perkins + intercept ──")
|
||||
print(f" slope={slope:.4f} intercept={intercept:.4f}")
|
||||
print(f" R²={r**2:.4f} p={p:.4e}")
|
||||
print(f" MAE: {reg_mae:.3f} mmHg")
|
||||
print(f" RMSE: {reg_rmse:.3f} mmHg")
|
||||
print(f" Improvement over ratio — MAE: {ratio_mae - reg_mae:+.3f} RMSE: {ratio_rmse - reg_rmse:+.3f}")
|
||||
|
||||
# LAD regression (minimises MAE) — more robust to outliers
|
||||
from scipy.optimize import minimize
|
||||
|
||||
def lad_loss(params):
|
||||
a, b = params
|
||||
return np.abs(pneumatic - (a * perkins + b)).mean()
|
||||
|
||||
lad_res = minimize(lad_loss, x0=[slope, intercept], method="Nelder-Mead")
|
||||
lad_slope, lad_intercept = lad_res.x
|
||||
lad_pred = lad_slope * perkins + lad_intercept
|
||||
lad_resid = pneumatic - lad_pred
|
||||
lad_mae = np.abs(lad_resid).mean()
|
||||
lad_rmse = np.sqrt((lad_resid ** 2).mean())
|
||||
|
||||
print(f"\n── LAD regression (minimises MAE — robust to outliers) ──────")
|
||||
print(f" slope={lad_slope:.4f} intercept={lad_intercept:.4f}")
|
||||
print(f" MAE: {lad_mae:.3f} mmHg")
|
||||
print(f" RMSE: {lad_rmse:.3f} mmHg")
|
||||
print(f" Improvement over ratio — MAE: {ratio_mae - lad_mae:+.3f} RMSE: {ratio_rmse - lad_rmse:+.3f}")
|
||||
|
||||
# ── Multiple regression: Perkins + Pachymetry ─────────────────────────────────
|
||||
# Pachymetry affects Perkins (applanation) more than Pneumatic, so including
|
||||
# CCT should absorb some instrument-specific bias.
|
||||
# Design matrix: [Perkins, Pachymetry, 1]
|
||||
from numpy.linalg import lstsq
|
||||
|
||||
X_multi = np.column_stack([perkins, pachymetry, np.ones(len(perkins))])
|
||||
coeffs, _, _, _ = lstsq(X_multi, pneumatic, rcond=None)
|
||||
coef_perkins, coef_pachy, coef_intercept = coeffs
|
||||
multi_pred = X_multi @ coeffs
|
||||
multi_resid = pneumatic - multi_pred
|
||||
multi_mae = np.abs(multi_resid).mean()
|
||||
multi_rmse = np.sqrt((multi_resid ** 2).mean())
|
||||
|
||||
# R² for the multiple model
|
||||
ss_res = (multi_resid ** 2).sum()
|
||||
ss_tot = ((pneumatic - pneumatic.mean()) ** 2).sum()
|
||||
multi_r2 = 1 - ss_res / ss_tot
|
||||
|
||||
# Partial correlation of Pachymetry with residual after removing Perkins effect
|
||||
perkins_resid = pneumatic - (slope * perkins + intercept)
|
||||
pachy_r, pachy_p = stats.pearsonr(pachymetry, perkins_resid)
|
||||
|
||||
print(f"\n── Multiple OLS (Perkins + Pachymetry, n={len(paired)}) ──────────────")
|
||||
print(f" Pneumatic = {coef_perkins:.4f}×Perkins + {coef_pachy:.5f}×Pachymetry + {coef_intercept:.4f}")
|
||||
print(f" R²={multi_r2:.4f} (vs simple OLS R²={r**2:.4f})")
|
||||
print(f" Pachymetry partial corr with OLS residuals: r={pachy_r:.3f} p={pachy_p:.4f}")
|
||||
print(f" MAE: {multi_mae:.3f} mmHg (vs ratio {ratio_mae:.3f})")
|
||||
print(f" RMSE: {multi_rmse:.3f} mmHg (vs ratio {ratio_rmse:.3f})")
|
||||
print(f" Improvement over ratio — MAE: {ratio_mae - multi_mae:+.3f} RMSE: {ratio_rmse - multi_rmse:+.3f}")
|
||||
print(f" NOTE: n={len(paired)} with 3 parameters — interpret with caution.")
|
||||
|
||||
# How much does the intercept matter at typical IOP values?
|
||||
typical_iop = np.array([10, 15, 20, 25])
|
||||
print(f"\n Comparison at typical Perkins values:")
|
||||
print(f" {'Perkins':>8} {'Ratio':>10} {'OLS':>10} {'LAD':>10}")
|
||||
for v in typical_iop:
|
||||
ratio_v = v * CURRENT_RATIO
|
||||
ols_v = slope * v + intercept
|
||||
lad_v = lad_slope * v + lad_intercept
|
||||
print(f" {v:>8.1f} {ratio_v:>10.2f} {ols_v:>10.2f} {lad_v:>10.2f}")
|
||||
|
||||
# ── IOP_raw coverage ──────────────────────────────────────────────────────────
|
||||
|
||||
pneumatic_only = df["Pneumatic"].notna() & df["Perkins"].isna()
|
||||
perkins_only = df["Pneumatic"].isna() & df["Perkins"].notna()
|
||||
both = df["Pneumatic"].notna() & df["Perkins"].notna()
|
||||
neither = df["Pneumatic"].isna() & df["Perkins"].isna()
|
||||
|
||||
print(f"\n── IOP measurement coverage ──────────────────────────────")
|
||||
print(f" Pneumatic only: {pneumatic_only.sum()}")
|
||||
print(f" Perkins only: {perkins_only.sum()}")
|
||||
print(f" Both: {both.sum()}")
|
||||
print(f" Neither: {neither.sum()}")
|
||||
print(f" Rows where IOP_raw == IOP_corr (no pachy correction): ", end="")
|
||||
|
||||
# ── plot ──────────────────────────────────────────────────────────────────────
|
||||
|
||||
fig = plt.figure(figsize=(16, 5), layout="constrained")
|
||||
gs = gridspec.GridSpec(1, 3, figure=fig)
|
||||
|
||||
x_line = np.linspace(perkins.min() - 1, perkins.max() + 1, 100)
|
||||
|
||||
# ── Row 1: conversion fits ────────────────────────────────────────────────────
|
||||
|
||||
# 1a. Scatter with all four fits
|
||||
ax1 = fig.add_subplot(gs[0, :2]) # spans first two columns
|
||||
ax1.scatter(perkins, pneumatic, alpha=0.6, s=35, color="gray", zorder=3, label="Observed pairs (n=41)")
|
||||
ax1.plot(x_line, x_line * CURRENT_RATIO, "r--", lw=2, label=f"Ratio ×{CURRENT_RATIO} MAE={ratio_mae:.2f}")
|
||||
ax1.plot(x_line, slope * x_line + intercept, "b-", lw=2, label=f"OLS (slope={slope:.3f}, int={intercept:.2f}) MAE={reg_mae:.2f}")
|
||||
ax1.plot(x_line, lad_slope * x_line + lad_intercept, "g-", lw=2, label=f"LAD (slope={lad_slope:.3f}, int={lad_intercept:.2f}) MAE={lad_mae:.2f}")
|
||||
|
||||
# Multi-reg projected at mean CCT
|
||||
pachy_mean, pachy_sd = pachymetry.mean(), pachymetry.std()
|
||||
multi_mean_line = coef_perkins * x_line + coef_pachy * pachy_mean + coef_intercept
|
||||
ax1.plot(x_line, multi_mean_line, color="purple", lw=2, ls="-.",
|
||||
label=f"Multi (Perkins+CCT) @ mean CCT MAE={multi_mae:.2f}")
|
||||
|
||||
ax1.set_xlabel("Perkins IOP (mmHg)")
|
||||
ax1.set_ylabel("Pneumatic IOP (mmHg)")
|
||||
ax1.set_title("Perkins → Pneumatic conversion: all methods")
|
||||
ax1.legend(fontsize=8)
|
||||
|
||||
# 1b. MAE / RMSE bar chart
|
||||
ax_bar = fig.add_subplot(gs[0, 2]) # third column
|
||||
bar_methods = ["Ratio", "OLS", "LAD", "Multi\n(+CCT)"]
|
||||
maes = [ratio_mae, reg_mae, lad_mae, multi_mae]
|
||||
rmses = [ratio_rmse, reg_rmse, lad_rmse, multi_rmse]
|
||||
bar_colors = ["#e05c5c", "#5c7de0", "#5cc97c", "#9b59b6"]
|
||||
bx = np.arange(4)
|
||||
w = 0.35
|
||||
ax_bar.bar(bx - w/2, maes, w, label="MAE", color=bar_colors)
|
||||
ax_bar.bar(bx + w/2, rmses, w, label="RMSE", color=bar_colors, alpha=0.5)
|
||||
ax_bar.set_xticks(bx); ax_bar.set_xticklabels(bar_methods, fontsize=8)
|
||||
ax_bar.set_ylabel("Error (mmHg)")
|
||||
ax_bar.set_title("MAE & RMSE comparison")
|
||||
ax_bar.legend(fontsize=9)
|
||||
ax_bar.set_ylim(0, max(rmses) * 1.25)
|
||||
for i, (mae, rmse) in enumerate(zip(maes, rmses)):
|
||||
ax_bar.text(i - w/2, mae + 0.05, f"{mae:.2f}", ha="center", va="bottom", fontsize=8)
|
||||
ax_bar.text(i + w/2, rmse + 0.05, f"{rmse:.2f}", ha="center", va="bottom", fontsize=8)
|
||||
|
||||
out = REPO / "analysis_data/iop_correction_comparison.png"
|
||||
fig.savefig(out, dpi=150)
|
||||
print(f"\nPlot saved to {out}")
|
||||
|
||||
# ── Figure 2: Pachymetry analysis ─────────────────────────────────────────────
|
||||
|
||||
fig2, axes2 = plt.subplots(1, 3, figsize=(15, 5))
|
||||
|
||||
pachy_norm = (pachymetry - pachymetry.mean()) / pachymetry.std()
|
||||
|
||||
# 2a. Pachymetry vs Pneumatic–Perkins difference
|
||||
ax = axes2[0]
|
||||
diff = pneumatic - perkins
|
||||
m, b, rr, pp, _ = stats.linregress(pachymetry, diff)
|
||||
ax.scatter(pachymetry, diff, alpha=0.6, s=35, color="steelblue")
|
||||
px = np.linspace(pachymetry.min() - 5, pachymetry.max() + 5, 100)
|
||||
ax.plot(px, m * px + b, "r-", lw=2, label=f"r={rr:.2f} p={pp:.3f}")
|
||||
ax.axhline(0, color="k", lw=0.7, ls="--")
|
||||
ax.set_xlabel("Pachymetry (μm)")
|
||||
ax.set_ylabel("Pneumatic − Perkins (mmHg)")
|
||||
ax.set_title("Does CCT predict the instrument gap?")
|
||||
ax.legend(fontsize=9)
|
||||
|
||||
# 2b. Residuals from ratio vs Pachymetry (coloured by size)
|
||||
ax = axes2[1]
|
||||
sc = ax.scatter(pachymetry, ratio_resid, c=perkins, cmap="viridis", alpha=0.7, s=35)
|
||||
plt.colorbar(sc, ax=ax, label="Perkins IOP")
|
||||
m2, b2, rr2, pp2, _ = stats.linregress(pachymetry, ratio_resid)
|
||||
ax.plot(px, m2 * px + b2, "r-", lw=2, label=f"r={rr2:.2f} p={pp2:.3f}")
|
||||
ax.axhline(0, color="k", lw=0.7, ls="--")
|
||||
ax.set_xlabel("Pachymetry (μm)")
|
||||
ax.set_ylabel("Ratio residual (mmHg)")
|
||||
ax.set_title("Ratio residuals vs Pachymetry\n(colour = Perkins IOP)")
|
||||
ax.legend(fontsize=9)
|
||||
|
||||
# 2c. Multiple regression residuals vs Pachymetry (should be flat if absorbed)
|
||||
ax = axes2[2]
|
||||
m3, b3, rr3, pp3, _ = stats.linregress(pachymetry, multi_resid)
|
||||
ax.scatter(pachymetry, multi_resid, alpha=0.6, s=35, color="darkorange")
|
||||
ax.plot(px, m3 * px + b3, "r-", lw=2, label=f"r={rr3:.2f} p={pp3:.3f}")
|
||||
ax.axhline(0, color="k", lw=0.7, ls="--")
|
||||
ax.set_xlabel("Pachymetry (μm)")
|
||||
ax.set_ylabel("Multi-reg residual (mmHg)")
|
||||
ax.set_title("Multi-reg residuals vs Pachymetry\n(flat = Pachymetry effect absorbed)")
|
||||
ax.legend(fontsize=9)
|
||||
|
||||
# shared y-axis
|
||||
ylim2 = max(np.abs(diff).max(), np.abs(ratio_resid).max(), np.abs(multi_resid).max()) + 1
|
||||
for ax in axes2[1:]:
|
||||
ax.set_ylim(-ylim2, ylim2)
|
||||
|
||||
fig2.suptitle(
|
||||
f"Pachymetry as a covariate (n={len(paired)}, CCT mean={pachymetry.mean():.0f}±{pachymetry.std():.0f} μm)",
|
||||
fontsize=11,
|
||||
)
|
||||
fig2.tight_layout()
|
||||
out2 = REPO / "analysis_data/iop_pachymetry_analysis.png"
|
||||
fig2.savefig(out2, dpi=150)
|
||||
print(f"Plot saved to {out2}")
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,140 +0,0 @@
|
||||
"""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()
|
||||
@@ -1,138 +0,0 @@
|
||||
"""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()
|
||||
@@ -1,966 +0,0 @@
|
||||
"""REFUGE training/evaluation helper.
|
||||
|
||||
Usage examples (after activating .venv_refuge):
|
||||
|
||||
python refuge_build.py --train-clf
|
||||
python refuge_build.py --eval --with-ttt
|
||||
|
||||
The script expects the REFUGE folder and writes checkpoints under
|
||||
models/v2/refuge/segmentation and models/v2/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/v2/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)
|
||||
if getattr(args, "clear_clf_cache", False):
|
||||
papila_clf.clear_disk_cache()
|
||||
|
||||
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_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,
|
||||
)
|
||||
if args.clear_clf_cache:
|
||||
clf.clear_disk_cache()
|
||||
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)
|
||||
if args.clear_clf_cache:
|
||||
clf.clear_disk_cache()
|
||||
|
||||
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.seg_weights is None:
|
||||
raise SystemExit(
|
||||
"--seg-weights must be specified for --eval-seg; "
|
||||
"e.g. --seg-weights models/v2/refuge/segmentation/per_image_refuge_build/best.pt"
|
||||
)
|
||||
ckpt = args.seg_weights
|
||||
if not ckpt.exists():
|
||||
raise FileNotFoundError(f"Segmentation weights not found at {ckpt}")
|
||||
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}")
|
||||
|
||||
if args.in_memory_cache:
|
||||
segmenter.prebuild_in_memory_cache(
|
||||
cache_workers=max(0, int(args.cache_workers)),
|
||||
include_train=False,
|
||||
include_val="val" in args.eval_seg_splits,
|
||||
include_holdout="holdout" in args.eval_seg_splits,
|
||||
)
|
||||
|
||||
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-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("cache_data/classifier_cache"),
|
||||
help="Directory to cache classifier preprocessing artifacts",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--clear-clf-cache",
|
||||
action="store_true",
|
||||
help="Delete all cached geometry/mask files before running (use when segmenter weights have changed)",
|
||||
)
|
||||
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/v2/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"],
|
||||
default=["holdout"],
|
||||
help="Segmentation splits to evaluate (default: holdout)",
|
||||
)
|
||||
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_unet_seg,
|
||||
args.train_clf,
|
||||
args.eval,
|
||||
args.eval_seg,
|
||||
args.export_backbone,
|
||||
]
|
||||
):
|
||||
raise SystemExit(
|
||||
"Specify at least one action: --train-unet-seg, --train-clf, --eval, --eval-seg, or --export-backbone"
|
||||
)
|
||||
|
||||
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()
|
||||
@@ -1,156 +0,0 @@
|
||||
#!/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.v2.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()
|
||||
@@ -1,27 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
"""CLI wrapper that delegates to classes.frontend.Multifold."""
|
||||
|
||||
from pathlib import Path
|
||||
import sys
|
||||
|
||||
# ensure repo root on path
|
||||
REPO_ROOT = Path(__file__).resolve().parents[2]
|
||||
if str(REPO_ROOT) not in sys.path:
|
||||
sys.path.insert(0, str(REPO_ROOT))
|
||||
|
||||
from classes.frontend import Multifold
|
||||
|
||||
|
||||
def run_cli(cli_args=None):
|
||||
parser = Multifold.build_parser()
|
||||
args = parser.parse_args(cli_args)
|
||||
runner = Multifold(args)
|
||||
runner.run()
|
||||
|
||||
|
||||
def main():
|
||||
run_cli()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,431 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Grid-search runner for run_multifold experiments.
|
||||
|
||||
Features:
|
||||
* Enumerates the requested configuration grid and writes grid_plan.csv.
|
||||
* Picks the next incomplete run, marks it running, executes run_multifold.py.
|
||||
* Records AUC/accuracy metrics per fold into grid_report.csv.
|
||||
* Removes model checkpoints for runs dominated (80%+ metrics worse) by others.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import csv
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
from contextlib import contextmanager
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from typing import Dict, List, Optional
|
||||
|
||||
import fcntl
|
||||
|
||||
REPO_ROOT = Path(__file__).resolve().parents[2]
|
||||
RUN_SCRIPT = REPO_ROOT / "scripts" / "run_multifold.py"
|
||||
MANIFEST = REPO_ROOT / "manifest.csv"
|
||||
GRID_DIR = REPO_ROOT / "analysis_data" / "grid_search"
|
||||
PLAN_PATH = GRID_DIR / "grid_plan.csv"
|
||||
REPORT_PATH = GRID_DIR / "grid_report.csv"
|
||||
LOCK_PATH = GRID_DIR / ".grid_lock"
|
||||
MODELS_ROOT = REPO_ROOT / "models" / "grid_search"
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
ap = argparse.ArgumentParser(description="Grid-search orchestrator for run_multifold.")
|
||||
ap.add_argument("--plan-date", default=datetime.now().strftime("%Y%m%d"),
|
||||
help="Date prefix used when generating run IDs (default: today).")
|
||||
ap.add_argument("--regen-plan", action="store_true",
|
||||
help="Rebuild the grid plan from scratch (overwrites existing plan).")
|
||||
ap.add_argument("--manifest", type=Path, default=MANIFEST,
|
||||
help="UNet manifest CSV for cropper.")
|
||||
ap.add_argument("--weights-dir", type=Path, default=REPO_ROOT / "models" / "unet_segmenter",
|
||||
help="Directory containing norm_* subfolders with best.pt.")
|
||||
ap.add_argument("--dry-run", action="store_true", help="Enumerate next run without executing.")
|
||||
ap.add_argument("--max-runs", type=int, default=1,
|
||||
help="Maximum runs to execute in this invocation (default: 1).")
|
||||
ap.add_argument("--run-all", action="store_true",
|
||||
help="Execute runs sequentially until plan is exhausted (overrides --max-runs).")
|
||||
return ap.parse_args()
|
||||
|
||||
|
||||
@contextmanager
|
||||
def file_lock(lock_path: Path):
|
||||
lock_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
with open(lock_path, "w") as lock_file:
|
||||
fcntl.flock(lock_file, fcntl.LOCK_EX)
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
fcntl.flock(lock_file, fcntl.LOCK_UN)
|
||||
|
||||
|
||||
def read_csv(path: Path) -> List[Dict[str, str]]:
|
||||
if not path.exists():
|
||||
return []
|
||||
with path.open(newline="") as fh:
|
||||
reader = csv.DictReader(fh)
|
||||
return list(reader)
|
||||
|
||||
|
||||
def write_csv(path: Path, rows: List[Dict[str, str]], headers: List[str]) -> None:
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
with path.open("w", newline="") as fh:
|
||||
writer = csv.DictWriter(fh, fieldnames=headers)
|
||||
writer.writeheader()
|
||||
for row in rows:
|
||||
writer.writerow(row)
|
||||
|
||||
|
||||
def grid_configs(base_date: str, weights_dir: Path) -> List[Dict[str, str]]:
|
||||
eval_modes = ["binary", "multiclass"]
|
||||
crop_variants = [
|
||||
("norm_imagenet", "imagenet"),
|
||||
("normalize_none", "none"),
|
||||
("norm_per_image", "per_image"),
|
||||
]
|
||||
tta_opts = [False, True]
|
||||
loss_modes = ["focal", "balanced", "none"]
|
||||
thaw_modes = ["none", "gradual"]
|
||||
|
||||
se_configs = []
|
||||
# none
|
||||
se_configs.append(("none", {"se_enabled": False}))
|
||||
# bridge only
|
||||
for pre in (True, False):
|
||||
se_configs.append((
|
||||
"bridge",
|
||||
{"se_enabled": True, "se_where": "bridge", "bridge_pre_norm": pre, "tower_pre_norm": None},
|
||||
))
|
||||
# tower only
|
||||
for pre in (True, False):
|
||||
se_configs.append((
|
||||
"tower",
|
||||
{"se_enabled": True, "se_where": "tower", "bridge_pre_norm": None, "tower_pre_norm": pre},
|
||||
))
|
||||
# both (four combos)
|
||||
for b_pre in (True, False):
|
||||
for t_pre in (True, False):
|
||||
se_configs.append((
|
||||
"both",
|
||||
{
|
||||
"se_enabled": True,
|
||||
"se_where": "both",
|
||||
"bridge_pre_norm": b_pre,
|
||||
"tower_pre_norm": t_pre,
|
||||
},
|
||||
))
|
||||
|
||||
combos = []
|
||||
idx = 0
|
||||
for eval_mode in eval_modes:
|
||||
for variant, norm in crop_variants:
|
||||
weights_path = weights_dir / variant / "best.pt"
|
||||
for tta in tta_opts:
|
||||
for loss in loss_modes:
|
||||
for thaw in thaw_modes:
|
||||
for se_name, se_opts in se_configs:
|
||||
run_id = f"{base_date}-{idx:04d}"
|
||||
combos.append({
|
||||
"run_id": run_id,
|
||||
"status": "incomplete",
|
||||
"eval_mode": eval_mode,
|
||||
"crop_variant": variant,
|
||||
"crop_normalize": norm,
|
||||
"crop_weights": str(weights_path),
|
||||
"crop_tta": str(tta),
|
||||
"loss_mode": loss,
|
||||
"thaw_mode": thaw,
|
||||
"se_mode": se_name,
|
||||
"se_bridge_pre_norm": str(se_opts.get("bridge_pre_norm")),
|
||||
"se_tower_pre_norm": str(se_opts.get("tower_pre_norm")),
|
||||
})
|
||||
idx += 1
|
||||
return combos
|
||||
|
||||
|
||||
PLAN_HEADERS = [
|
||||
"run_id",
|
||||
"status",
|
||||
"eval_mode",
|
||||
"crop_variant",
|
||||
"crop_normalize",
|
||||
"crop_weights",
|
||||
"crop_tta",
|
||||
"loss_mode",
|
||||
"thaw_mode",
|
||||
"se_mode",
|
||||
"se_bridge_pre_norm",
|
||||
"se_tower_pre_norm",
|
||||
]
|
||||
|
||||
|
||||
def ensure_plan(args: argparse.Namespace) -> None:
|
||||
if args.regen_plan or not PLAN_PATH.exists():
|
||||
combos = grid_configs(args.plan_date, args.weights_dir)
|
||||
write_csv(PLAN_PATH, combos, PLAN_HEADERS)
|
||||
print(f"[grid] Plan created with {len(combos)} runs at {PLAN_PATH}")
|
||||
|
||||
|
||||
def select_next_run() -> Optional[Dict[str, str]]:
|
||||
rows = read_csv(PLAN_PATH)
|
||||
for row in rows:
|
||||
if row["status"] == "incomplete":
|
||||
row["status"] = "running"
|
||||
write_csv(PLAN_PATH, rows, PLAN_HEADERS)
|
||||
return row
|
||||
return None
|
||||
|
||||
|
||||
def update_run_status(run_id: str, new_status: str) -> None:
|
||||
rows = read_csv(PLAN_PATH)
|
||||
for row in rows:
|
||||
if row["run_id"] == run_id:
|
||||
row["status"] = new_status
|
||||
break
|
||||
write_csv(PLAN_PATH, rows, PLAN_HEADERS)
|
||||
|
||||
|
||||
def build_run_command(row: Dict[str, str], manifest: Path) -> List[str]:
|
||||
cmd = [
|
||||
sys.executable,
|
||||
str(RUN_SCRIPT),
|
||||
"--backbone",
|
||||
"resnet50",
|
||||
"--fusion-mode",
|
||||
"fused",
|
||||
"--epochs",
|
||||
"40",
|
||||
"--batch-size",
|
||||
"8",
|
||||
"--img-crop-manifest",
|
||||
str(manifest),
|
||||
"--img-crop-weights",
|
||||
row["crop_weights"],
|
||||
"--img-crop-normalize",
|
||||
row["crop_normalize"],
|
||||
"--eval_mode",
|
||||
row["eval_mode"],
|
||||
"--holdout-per-class",
|
||||
"12",
|
||||
"--run-id",
|
||||
row["run_id"],
|
||||
"--shortname",
|
||||
"grid_search",
|
||||
]
|
||||
if row["crop_tta"] == "True":
|
||||
cmd.append("--img-crop-tta")
|
||||
|
||||
# Loss/balancing modes
|
||||
if row["loss_mode"] == "focal":
|
||||
cmd.extend(["--focal-gamma", "2.0"])
|
||||
elif row["loss_mode"] == "balanced":
|
||||
cmd.append("--balanced-sampler")
|
||||
|
||||
# Thaw schedule
|
||||
if row["thaw_mode"] == "gradual":
|
||||
cmd.append("--gradual-thaw")
|
||||
cmd.extend(["--thaw-ratio", "0.33"])
|
||||
cmd.extend(["--thaw-start-epoch", "10"])
|
||||
cmd.extend(["--thaw-target", "image"])
|
||||
|
||||
# SE settings
|
||||
if row["se_mode"] == "none":
|
||||
cmd.append("--no-se")
|
||||
else:
|
||||
cmd.extend(["--se-reduction", "16"])
|
||||
cmd.extend(["--se-reduction-tower", "16"])
|
||||
cmd.extend(["--se-where", row["se_mode"]])
|
||||
bridge_pre = row["se_bridge_pre_norm"]
|
||||
tower_pre = row["se_tower_pre_norm"]
|
||||
if bridge_pre == "True":
|
||||
cmd.append("--se-pre-norm")
|
||||
elif bridge_pre == "False":
|
||||
cmd.append("--no-se-pre-norm")
|
||||
if tower_pre == "True":
|
||||
cmd.append("--se-pre-norm-tower")
|
||||
elif tower_pre == "False":
|
||||
cmd.append("--no-se-pre-norm-tower")
|
||||
return cmd
|
||||
|
||||
|
||||
def run_command(cmd: List[str]) -> None:
|
||||
print("[grid] Launching:", " ".join(cmd))
|
||||
subprocess.run(cmd, check=True)
|
||||
|
||||
|
||||
METRIC_KEYS = ["auc_fused", "auc_img", "auc_md", "acc_fused", "acc_img", "acc_md"]
|
||||
|
||||
|
||||
def extract_metrics(run_id: str) -> Dict[str, str]:
|
||||
summary_path = REPO_ROOT / "analysis_data" / "grid_search" / run_id / "summary.json"
|
||||
if not summary_path.exists():
|
||||
raise FileNotFoundError(f"Missing summary.json for run {run_id}")
|
||||
with summary_path.open() as fh:
|
||||
summary = json.load(fh)
|
||||
|
||||
rows = {}
|
||||
for fold in summary.get("fold_metrics", []):
|
||||
if not isinstance(fold, dict):
|
||||
continue
|
||||
f_idx = fold.get("fold")
|
||||
stats = fold.get("stats") or {}
|
||||
if not isinstance(stats, dict):
|
||||
continue
|
||||
for key in METRIC_KEYS:
|
||||
val = stats.get(key)
|
||||
if val is None:
|
||||
continue
|
||||
rows[f"metric_fold{f_idx}_{key}"] = str(val)
|
||||
best_mean = summary.get("best_metric_mean")
|
||||
if best_mean is not None:
|
||||
rows["metric_best_mean"] = str(best_mean)
|
||||
return rows
|
||||
|
||||
|
||||
def update_report(row: Dict[str, str], metrics: Dict[str, str]) -> None:
|
||||
existing = read_csv(REPORT_PATH)
|
||||
# Remove existing entry for run_id
|
||||
existing = [r for r in existing if r.get("run_id") != row["run_id"]]
|
||||
record = {**row, **metrics}
|
||||
existing.append(record)
|
||||
headers = sorted({key for r in existing for key in r.keys()})
|
||||
write_csv(REPORT_PATH, existing, headers)
|
||||
|
||||
|
||||
def load_report_rows() -> List[Dict[str, str]]:
|
||||
return read_csv(REPORT_PATH)
|
||||
|
||||
|
||||
def metric_columns(rows: List[Dict[str, str]]) -> List[str]:
|
||||
keys = set()
|
||||
for row in rows:
|
||||
for key in row:
|
||||
if key.startswith("metric_"):
|
||||
keys.add(key)
|
||||
return sorted(keys)
|
||||
|
||||
|
||||
def _to_float(val: str) -> Optional[float]:
|
||||
try:
|
||||
f = float(val)
|
||||
if math.isnan(f):
|
||||
return None
|
||||
return f
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
def prune_dominated(rows: List[Dict[str, str]]) -> None:
|
||||
"""
|
||||
Remove model directories for runs that are clearly dominated by another run.
|
||||
A run is dominated if:
|
||||
* Another run has a strictly higher metric_best_mean, OR
|
||||
* Another run is >= on >=80% of overlapping metrics and strictly better on at least one.
|
||||
"""
|
||||
metrics = metric_columns(rows)
|
||||
if not metrics:
|
||||
return
|
||||
|
||||
dominated = set()
|
||||
for row in rows:
|
||||
run_id = row["run_id"]
|
||||
row_vals = {m: row.get(m) for m in metrics}
|
||||
row_best = _to_float(row_vals.get("metric_best_mean"))
|
||||
|
||||
for other in rows:
|
||||
if other["run_id"] == run_id:
|
||||
continue
|
||||
|
||||
other_vals = {m: other.get(m) for m in metrics}
|
||||
other_best = _to_float(other_vals.get("metric_best_mean"))
|
||||
|
||||
# Fast path: compare aggregate best mean if both have it
|
||||
if row_best is not None and other_best is not None and other_best > row_best:
|
||||
dominated.add(run_id)
|
||||
break
|
||||
|
||||
# Fallback: overlap-wise dominance
|
||||
comparisons = []
|
||||
better = 0
|
||||
for key in metrics:
|
||||
v1 = _to_float(row_vals.get(key))
|
||||
v2 = _to_float(other_vals.get(key))
|
||||
if v1 is None or v2 is None:
|
||||
continue
|
||||
comparisons.append(v2 >= v1)
|
||||
if v2 > v1:
|
||||
better += 1
|
||||
if not comparisons:
|
||||
continue
|
||||
fraction = sum(comparisons) / len(comparisons)
|
||||
if fraction >= 0.8 and better > 0:
|
||||
dominated.add(run_id)
|
||||
break
|
||||
|
||||
for run_id in dominated:
|
||||
model_dir = MODELS_ROOT / run_id
|
||||
if model_dir.exists():
|
||||
print(f"[grid] Removing dominated model artifacts for {run_id}")
|
||||
try:
|
||||
shutil.rmtree(model_dir)
|
||||
except OSError as exc:
|
||||
# Don't fail the grid run if cleanup isn't permitted (e.g., locked SMB dirs).
|
||||
print(f"[grid] Warning: could not remove {model_dir}: {exc}")
|
||||
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
ensure_plan(args)
|
||||
if args.dry_run:
|
||||
with file_lock(LOCK_PATH):
|
||||
next_run = select_next_run()
|
||||
if next_run is None:
|
||||
print("[grid] No incomplete runs remaining.")
|
||||
return
|
||||
update_run_status(next_run["run_id"], "incomplete")
|
||||
print("[grid] Next run:", next_run)
|
||||
return
|
||||
|
||||
max_runs = None if args.run_all else args.max_runs
|
||||
runs_done = 0
|
||||
|
||||
while True:
|
||||
with file_lock(LOCK_PATH):
|
||||
next_run = select_next_run()
|
||||
if next_run is None:
|
||||
if runs_done == 0:
|
||||
print("[grid] All runs completed.")
|
||||
else:
|
||||
print(f"[grid] No more runs remaining after {runs_done} run(s).")
|
||||
return
|
||||
|
||||
run_id = next_run["run_id"]
|
||||
try:
|
||||
cmd = build_run_command(next_run, args.manifest)
|
||||
run_command(cmd)
|
||||
metrics = extract_metrics(run_id)
|
||||
with file_lock(LOCK_PATH):
|
||||
update_run_status(run_id, "completed")
|
||||
update_report(next_run, metrics)
|
||||
report_rows = load_report_rows()
|
||||
prune_dominated(report_rows)
|
||||
print(f"[grid] Run {run_id} completed.")
|
||||
except Exception as exc:
|
||||
with file_lock(LOCK_PATH):
|
||||
update_run_status(run_id, "incomplete")
|
||||
raise SystemExit(f"[grid] Run {run_id} failed: {exc}") from exc
|
||||
|
||||
runs_done += 1
|
||||
if max_runs is not None and runs_done >= max_runs:
|
||||
print(f"[grid] Reached run limit ({max_runs}); stopping.")
|
||||
return
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,60 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Thin CLI wrapper that runs V2 hypertower modes sequentially."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
import sys
|
||||
import argparse
|
||||
|
||||
REPO_ROOT = Path(__file__).resolve().parents[3]
|
||||
if str(REPO_ROOT) not in sys.path:
|
||||
sys.path.insert(0, str(REPO_ROOT))
|
||||
|
||||
from classes.v2.v2_hypertower import V2HyperTower
|
||||
|
||||
|
||||
def parse_args():
|
||||
ap = argparse.ArgumentParser(
|
||||
description="Run selected eval/tower mode combinations sequentially."
|
||||
)
|
||||
ap.add_argument(
|
||||
"--eval-modes",
|
||||
nargs="+",
|
||||
choices=["binary", "multiclass"],
|
||||
default=["binary", "multiclass"],
|
||||
)
|
||||
ap.add_argument(
|
||||
"--tower-modes",
|
||||
nargs="+",
|
||||
choices=["single", "ensemble", "bilateral", "classic"],
|
||||
default=["single", "ensemble", "bilateral"],
|
||||
)
|
||||
return ap.parse_known_args()
|
||||
|
||||
|
||||
def main():
|
||||
seq_args, remaining = parse_args()
|
||||
base_parser = V2HyperTower.build_parser()
|
||||
first_run = True
|
||||
for eval_mode in seq_args.eval_modes:
|
||||
for tower_mode in seq_args.tower_modes:
|
||||
tower_mode = "single" if tower_mode == "classic" else tower_mode
|
||||
cli = list(remaining) + ["--eval-mode", eval_mode, "--tower-mode", tower_mode]
|
||||
# Clear cache only on the first run; reuse it for all subsequent runs.
|
||||
if not first_run:
|
||||
cli.append("--persist-img-crop-cache")
|
||||
args = base_parser.parse_args(cli)
|
||||
# Skip if this mode is already fully complete.
|
||||
if args.run_name:
|
||||
tm_dir = Path(args.output_root) / args.run_name / eval_mode / tower_mode
|
||||
if (tm_dir / "summary.json").exists():
|
||||
print(f"[compare] {eval_mode}:{tower_mode} already complete — skipping.")
|
||||
first_run = False # treat as done so cache is preserved for later runs
|
||||
continue
|
||||
V2HyperTower(args).run()
|
||||
first_run = False
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,107 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
10× repeated 5-fold CV runner for the best hypertower configuration
|
||||
(nocrop, ensemble mode, both binary and multiclass).
|
||||
|
||||
Each repetition uses a different fold-seed so the 5 folds are split
|
||||
differently, giving 50 folds per eval-mode total. Holdout composition
|
||||
is kept identical across repetitions (same --holdout-seed).
|
||||
|
||||
Results land under:
|
||||
{output-root}/rep{N:02d}/{eval_mode}/ensemble/fold{K}/
|
||||
|
||||
Usage
|
||||
-----
|
||||
python scripts/main/v2/run_10x5cv.py \
|
||||
--n-reps 10 \
|
||||
--eval-modes binary multiclass \
|
||||
--output-root analysis_data/pipeline_10x5 \
|
||||
--epochs 40 --fused-head \
|
||||
--backbone refugelike
|
||||
|
||||
Any extra flags are forwarded directly to V2HyperTower.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
REPO_ROOT = Path(__file__).resolve().parents[3]
|
||||
if str(REPO_ROOT) not in sys.path:
|
||||
sys.path.insert(0, str(REPO_ROOT))
|
||||
|
||||
from classes.v2.v2_hypertower import V2HyperTower
|
||||
|
||||
# Base fold seed for rep 0; rep N uses BASE_SEED + N * SEED_STRIDE
|
||||
_BASE_SEED = 100
|
||||
_SEED_STRIDE = 100
|
||||
|
||||
|
||||
def _parse_own(argv=None):
|
||||
ap = argparse.ArgumentParser(
|
||||
description=__doc__,
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||||
add_help=False,
|
||||
)
|
||||
ap.add_argument("--n-reps", type=int, default=10,
|
||||
help="Number of repetitions (default: 10).")
|
||||
ap.add_argument("--eval-modes", nargs="+",
|
||||
choices=["binary", "multiclass"],
|
||||
default=["binary", "multiclass"])
|
||||
ap.add_argument("--output-root", default="analysis_data/pipeline_10x5",
|
||||
help="Parent directory for all rep sub-runs.")
|
||||
ap.add_argument("-h", "--help", action="store_true")
|
||||
return ap.parse_known_args(argv)
|
||||
|
||||
|
||||
def main(argv=None):
|
||||
own, remaining = _parse_own(argv)
|
||||
|
||||
if own.help:
|
||||
print(__doc__)
|
||||
base_parser = V2HyperTower.build_parser()
|
||||
base_parser.print_help()
|
||||
return
|
||||
|
||||
base_parser = V2HyperTower.build_parser()
|
||||
output_root = Path(own.output_root)
|
||||
first_run = True
|
||||
|
||||
for rep in range(own.n_reps):
|
||||
fold_seed = _BASE_SEED + rep * _SEED_STRIDE
|
||||
rep_label = f"rep{rep:02d}"
|
||||
|
||||
for eval_mode in own.eval_modes:
|
||||
tower_mode = "ensemble"
|
||||
|
||||
# Skip if already fully complete
|
||||
tm_dir = output_root / rep_label / eval_mode / tower_mode
|
||||
if (tm_dir / "summary.json").exists():
|
||||
print(f"[10x5cv] {rep_label} {eval_mode}:{tower_mode} — already done, skipping.")
|
||||
first_run = False
|
||||
continue
|
||||
|
||||
cli = list(remaining) + [
|
||||
"--eval-mode", eval_mode,
|
||||
"--tower-mode", tower_mode,
|
||||
"--fold-seed", str(fold_seed),
|
||||
"--run-name", rep_label,
|
||||
"--output-root", str(output_root),
|
||||
]
|
||||
|
||||
# Reuse crop cache across runs after the first
|
||||
if not first_run:
|
||||
cli.append("--persist-img-crop-cache")
|
||||
|
||||
print(f"\n[10x5cv] Starting {rep_label} {eval_mode}:{tower_mode} "
|
||||
f"(fold_seed={fold_seed})")
|
||||
args = base_parser.parse_args(cli)
|
||||
V2HyperTower(args).run()
|
||||
first_run = False
|
||||
|
||||
print(f"\n[10x5cv] All done. Results in: {output_root}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,90 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Run single-mode binary + multiclass for each IOP correction method and
|
||||
collect all results under one output root for easy comparison.
|
||||
|
||||
Output layout:
|
||||
analysis_data/iop_corr_comparison/
|
||||
ratio/binary/single/ ratio/multiclass/single/
|
||||
ols/binary/single/ ols/multiclass/single/
|
||||
lad/binary/single/ lad/multiclass/single/
|
||||
multi/binary/single/ multi/multiclass/single/
|
||||
|
||||
Usage
|
||||
-----
|
||||
python scripts/main/v2/run_iop_corr_comparison.py [V2HyperTower args...]
|
||||
|
||||
Any extra args (backbone, epochs, img-crop-*, etc.) are forwarded to every run.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
REPO_ROOT = Path(__file__).resolve().parents[3]
|
||||
if str(REPO_ROOT) not in sys.path:
|
||||
sys.path.insert(0, str(REPO_ROOT))
|
||||
|
||||
from classes.v2.v2_hypertower import V2HyperTower
|
||||
|
||||
IOP_METHODS = ["ratio", "ols", "lad", "multi"]
|
||||
EVAL_MODES = ["binary", "multiclass"]
|
||||
OUTPUT_ROOT = "analysis_data/iop_corr_comparison"
|
||||
|
||||
|
||||
def main() -> None:
|
||||
base_parser = V2HyperTower.build_parser()
|
||||
# Consume only the remaining (forwarded) args — iop-corr-method and
|
||||
# run-name are set by this script; eval-mode and tower-mode likewise.
|
||||
_, remaining = base_parser.parse_known_args()
|
||||
|
||||
first_run = True
|
||||
for method in IOP_METHODS:
|
||||
for eval_mode in EVAL_MODES:
|
||||
run_name = method # one sub-folder per method
|
||||
tm_dir = (Path(OUTPUT_ROOT) / run_name / eval_mode / "single")
|
||||
if (tm_dir / "summary.json").exists():
|
||||
print(f"[iop_corr] {method}/{eval_mode}/single — already done, skipping.")
|
||||
first_run = False
|
||||
continue
|
||||
|
||||
cli = list(remaining) + [
|
||||
"--eval-mode", eval_mode,
|
||||
"--tower-mode", "single",
|
||||
"--iop-corr-method", method,
|
||||
"--output-root", OUTPUT_ROOT,
|
||||
"--run-name", run_name,
|
||||
]
|
||||
if not first_run:
|
||||
cli.append("--persist-img-crop-cache")
|
||||
|
||||
print(f"\n[iop_corr] Starting {method}/{eval_mode}/single ...")
|
||||
args = base_parser.parse_args(cli)
|
||||
V2HyperTower(args).run()
|
||||
first_run = False
|
||||
|
||||
# ── summary table ──────────────────────────────────────────────────────
|
||||
import json
|
||||
print("\n" + "=" * 60)
|
||||
print("IOP correction method comparison — single mode")
|
||||
print("=" * 60)
|
||||
header = f"{'Method':<8} {'Mode':<12} {'Val AUC':>10} {'Hld AUC':>10}"
|
||||
print(header)
|
||||
print("-" * len(header))
|
||||
for method in IOP_METHODS:
|
||||
for eval_mode in EVAL_MODES:
|
||||
p = Path(OUTPUT_ROOT) / method / eval_mode / "single" / "summary.json"
|
||||
if not p.exists():
|
||||
print(f"{method:<8} {eval_mode:<12} {'missing':>10} {'missing':>10}")
|
||||
continue
|
||||
ms = json.loads(p.read_text()).get("mode_summary", {})
|
||||
val = ms.get("classic_best_val", {})
|
||||
hld = ms.get("classic_holdout", {})
|
||||
val_s = f"{val['auc_mean']:.3f}±{val['auc_std']:.3f}" if val.get("auc_mean") else "—"
|
||||
hld_s = f"{hld['auc_mean']:.3f}±{hld['auc_std']:.3f}" if hld.get("auc_mean") else "—"
|
||||
print(f"{method:<8} {eval_mode:<12} {val_s:>10} {hld_s:>10}")
|
||||
print("=" * 60)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,484 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
V2-parity metadata-only runner.
|
||||
|
||||
Goal:
|
||||
- Match V2HyperTower single-model metadata-only behavior as closely as possible.
|
||||
- Avoid image tower/image IO overhead in forward/training.
|
||||
|
||||
How parity is achieved:
|
||||
- Uses PatientFirstSplitManager (same split policy).
|
||||
- Uses PAPILA profile builders + V2 filters:
|
||||
- eye_train = filter_eye_samples(...)
|
||||
- bilat_val/test = filter_bilateral_samples(...)
|
||||
- Uses V2 training/eval helpers directly:
|
||||
- train_single_epoch(...)
|
||||
- collect_probs_single_components(...)
|
||||
- Uses bridge_mode="metadata_only".
|
||||
|
||||
Implementation detail:
|
||||
- Batch dictionaries still include image slots to satisfy shared V2 helpers,
|
||||
but these are tiny dummy tensors and are never consumed in metadata-only mode.
|
||||
|
||||
Outputs:
|
||||
analysis_data/{run_name}/{eval_mode}/{tower_mode}/fold{N}/
|
||||
y_true.npy
|
||||
probs_classic.npy or probs_ensemble.npy
|
||||
y_true_holdout.npy
|
||||
probs_classic_holdout.npy or probs_ensemble_holdout.npy
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import copy
|
||||
import json
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
|
||||
# ensure repo root is on sys.path when run directly
|
||||
_REPO_ROOT = Path(__file__).resolve().parents[3]
|
||||
if str(_REPO_ROOT) not in sys.path:
|
||||
sys.path.insert(0, str(_REPO_ROOT))
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from torch import nn
|
||||
from torch.utils.data import DataLoader, Dataset
|
||||
|
||||
from classes.v2.bridges import Bridge
|
||||
from classes.v2.loader_factory import (
|
||||
build_balanced_sampler,
|
||||
filter_bilateral_samples,
|
||||
filter_eye_samples,
|
||||
)
|
||||
from classes.v2.metrics import _score_arrays
|
||||
from classes.v2.models import collect_probs_single_components, train_single_epoch
|
||||
from classes.v2.papila_builders import build_papila_data
|
||||
from classes.v2.profiles import build_papila_profile
|
||||
from classes.v2.split_manager import PatientFirstSplitManager
|
||||
from classes.v2.towers import MDTower
|
||||
from classes.v2.utils import choose_device, seed_everything
|
||||
|
||||
|
||||
class MetadataOnlySingleHT(nn.Module):
|
||||
"""SingleEyeHT-compatible shell without real image tower usage."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
clinical_data,
|
||||
num_classes: int,
|
||||
md_hidden_dim: int,
|
||||
fusion_dim: int,
|
||||
dropout: float,
|
||||
use_se: bool,
|
||||
se_reduction: int,
|
||||
se_pre_norm: bool,
|
||||
):
|
||||
super().__init__()
|
||||
# Placeholder module to satisfy phase toggling logic.
|
||||
self.img_tower = nn.Identity()
|
||||
self.md_tower = MDTower(
|
||||
clinical_data=clinical_data,
|
||||
hidden_dim=md_hidden_dim,
|
||||
dropout=dropout,
|
||||
use_se=use_se,
|
||||
se_reduction=se_reduction,
|
||||
se_pre_norm=se_pre_norm,
|
||||
)
|
||||
# img_dim is irrelevant in metadata_only mode, but Bridge defines img head params.
|
||||
self.bridge = Bridge(
|
||||
img_dim=1,
|
||||
meta_dim=self.md_tower.out_dim,
|
||||
num_classes=num_classes,
|
||||
fusion_dim=fusion_dim,
|
||||
mode="metadata_only",
|
||||
use_se=False,
|
||||
se_reduction=16,
|
||||
se_pre_norm=True,
|
||||
)
|
||||
|
||||
|
||||
class EyeMetaDataset(Dataset):
|
||||
"""Eye-level dataset for V2 train_single_epoch input contract."""
|
||||
|
||||
def __init__(self, samples: list[dict]):
|
||||
self.samples = samples
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self.samples)
|
||||
|
||||
def __getitem__(self, idx: int) -> dict[str, torch.Tensor]:
|
||||
s = self.samples[idx]
|
||||
return {
|
||||
"image_1": torch.zeros(1, dtype=torch.float32),
|
||||
"matrix_1": torch.as_tensor(s["matrix_1"], dtype=torch.float32),
|
||||
"label_1": torch.tensor(int(s["label_1"]), dtype=torch.long),
|
||||
}
|
||||
|
||||
|
||||
class BilatMetaDataset(Dataset):
|
||||
"""Patient-level bilateral dataset for collect_probs_single_components."""
|
||||
|
||||
def __init__(self, samples: list[dict]):
|
||||
self.samples = samples
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self.samples)
|
||||
|
||||
def __getitem__(self, idx: int) -> dict[str, torch.Tensor]:
|
||||
s = self.samples[idx]
|
||||
return {
|
||||
"image_1": torch.zeros(1, dtype=torch.float32),
|
||||
"image_2": torch.zeros(1, dtype=torch.float32),
|
||||
"matrix_1": torch.as_tensor(s["matrix_1"], dtype=torch.float32),
|
||||
"matrix_2": torch.as_tensor(s["matrix_2"], dtype=torch.float32),
|
||||
"label_1": torch.tensor(int(s["label_1"]), dtype=torch.long),
|
||||
}
|
||||
|
||||
|
||||
def _phase_for_epoch(epoch_idx: int, warm_tower: int, warm_fused: int, main_epochs: int) -> tuple[str, int]:
|
||||
if epoch_idx < warm_tower:
|
||||
return "tower_warmup", 0
|
||||
if epoch_idx < (warm_tower + warm_fused):
|
||||
return "fused_warmup", 0
|
||||
if epoch_idx < (warm_tower + warm_fused + main_epochs):
|
||||
main_ep = epoch_idx - warm_tower - warm_fused + 1
|
||||
return "main", main_ep
|
||||
return "done", main_epochs
|
||||
|
||||
|
||||
def _evaluate_single(
|
||||
model: nn.Module,
|
||||
loader: DataLoader,
|
||||
device: torch.device,
|
||||
num_classes: int,
|
||||
aggregate_patient: bool,
|
||||
) -> tuple[np.ndarray, np.ndarray, float, float]:
|
||||
y, p_fused, _, p_md = collect_probs_single_components(
|
||||
model, loader, device, aggregate_patient=aggregate_patient
|
||||
)
|
||||
# In metadata_only mode p_fused == p_md; keep md explicitly for clarity.
|
||||
probs = p_md if p_md.size else p_fused
|
||||
acc, auc, _ = _score_arrays(y, probs, num_classes)
|
||||
return y, probs, float(auc), float(acc)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
ap = argparse.ArgumentParser(description="V2-parity metadata-only runner.")
|
||||
ap.add_argument("--eval-mode", required=True, choices=["binary", "multiclass"])
|
||||
ap.add_argument("--tower-mode", default="single", choices=["single", "ensemble"])
|
||||
ap.add_argument("--run-name", required=True)
|
||||
|
||||
ap.add_argument("--epochs", type=int, default=40, help="Main-phase epochs.")
|
||||
ap.add_argument("--warmup-tower-epochs", type=int, default=None)
|
||||
ap.add_argument("--warmup-fused-epochs", type=int, default=None)
|
||||
ap.add_argument("--batch-size", type=int, default=8)
|
||||
ap.add_argument("--lr", type=float, default=1e-4)
|
||||
ap.add_argument("--weight-decay", type=float, default=0.0)
|
||||
ap.add_argument("--bcd-prob", type=float, default=0.5)
|
||||
|
||||
ap.add_argument("--md-hidden-dim", type=int, default=128)
|
||||
ap.add_argument("--fusion-dim", type=int, default=256)
|
||||
ap.add_argument("--dropout", type=float, default=0.1)
|
||||
ap.add_argument("--use-se", action="store_true")
|
||||
ap.add_argument("--se-reduction", type=int, default=16)
|
||||
ap.add_argument("--se-pre-norm", action="store_true")
|
||||
|
||||
ap.add_argument("--n-splits", type=int, default=5)
|
||||
ap.add_argument("--holdout-per-class", type=int, default=5)
|
||||
ap.add_argument("--holdout-seed", type=int, default=123)
|
||||
ap.add_argument("--fold-seed", type=int, default=42)
|
||||
ap.add_argument("--seed", type=int, default=1234)
|
||||
ap.add_argument("--balanced-sampling", action=argparse.BooleanOptionalAction, default=False)
|
||||
|
||||
ap.add_argument("--analysis-dir", default="analysis_data")
|
||||
ap.add_argument("--image-dir", default="Papila/FundusImages")
|
||||
ap.add_argument("--clinical-dir", default="Papila/ClinicalData")
|
||||
ap.add_argument("--label-col", default="Diagnosis")
|
||||
ap.add_argument("--patient-col", default="Patient ID")
|
||||
ap.add_argument("--cat-cols", nargs="*", default=["Gender", "Phakic/Pseudophakic"])
|
||||
|
||||
args = ap.parse_args()
|
||||
|
||||
seed_everything(args.seed)
|
||||
device = choose_device(None)
|
||||
print(f"Device: {device}", flush=True)
|
||||
|
||||
print("Loading PAPILA data...", flush=True)
|
||||
data = build_papila_data(
|
||||
image_dir=args.image_dir,
|
||||
clinical_dir=args.clinical_dir,
|
||||
label_col=args.label_col,
|
||||
cat_cols=list(args.cat_cols),
|
||||
n_splits=args.n_splits,
|
||||
random_seed=args.fold_seed,
|
||||
)
|
||||
print(f"Loaded: {len(data.df)} rows feature_dim={data.feature_dim}", flush=True)
|
||||
|
||||
num_classes = 2 if args.eval_mode == "binary" else 3
|
||||
df_mode = data.df.copy()
|
||||
if args.eval_mode == "binary":
|
||||
df_mode = df_mode[df_mode[args.label_col].isin([0, 1])].reset_index(drop=True)
|
||||
print(f"[{args.eval_mode}] rows={len(df_mode)}", flush=True)
|
||||
|
||||
class _ClinicalShim:
|
||||
label_col = args.label_col
|
||||
|
||||
def __init__(self, df):
|
||||
self.df = df
|
||||
|
||||
split_args = SimpleNamespace(
|
||||
eval_mode=args.eval_mode,
|
||||
holdout_per_class=args.holdout_per_class,
|
||||
holdout_seed=args.holdout_seed,
|
||||
n_splits=args.n_splits,
|
||||
fold_seed=args.fold_seed,
|
||||
)
|
||||
splitter = PatientFirstSplitManager(patient_col=args.patient_col, label_col=args.label_col)
|
||||
plans = splitter.build_plans(clinical=_ClinicalShim(df_mode), args=split_args)
|
||||
|
||||
profile_eye = build_papila_profile(
|
||||
patient_col=args.patient_col,
|
||||
label_col=args.label_col,
|
||||
sample_mode="eye",
|
||||
)
|
||||
profile_patient = build_papila_profile(
|
||||
patient_col=args.patient_col,
|
||||
label_col=args.label_col,
|
||||
sample_mode="patient",
|
||||
)
|
||||
|
||||
warm_tower = int(args.warmup_tower_epochs) if args.warmup_tower_epochs is not None else 2
|
||||
warm_fused = int(args.warmup_fused_epochs) if args.warmup_fused_epochs is not None else 2
|
||||
total_epochs = warm_tower + warm_fused + int(args.epochs)
|
||||
|
||||
out_root = Path(args.analysis_dir) / args.run_name / args.eval_mode / args.tower_mode
|
||||
out_root.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
fold_metrics = []
|
||||
aggregate_patient = args.tower_mode == "ensemble"
|
||||
|
||||
for fold_idx, split in enumerate(plans[: args.n_splits]):
|
||||
fold_dir = out_root / f"fold{fold_idx}"
|
||||
fold_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
eye_train = filter_eye_samples(profile_eye.build_samples(df=split.train, clinical=data))
|
||||
bilat_val = filter_bilateral_samples(profile_patient.build_samples(df=split.val, clinical=data))
|
||||
|
||||
holdout_bilat = []
|
||||
if split.holdout is not None and not split.holdout.empty:
|
||||
holdout_bilat = filter_bilateral_samples(
|
||||
profile_patient.build_samples(df=split.holdout, clinical=data)
|
||||
)
|
||||
|
||||
if not eye_train or not bilat_val:
|
||||
print(f"[fold {fold_idx+1}] skipped (eye_train={len(eye_train)} bilat_val={len(bilat_val)})", flush=True)
|
||||
fold_metrics.append(
|
||||
{
|
||||
"fold": fold_idx,
|
||||
"best_epoch": None,
|
||||
"best_phase": None,
|
||||
"val_auc": float("nan"),
|
||||
"val_acc": float("nan"),
|
||||
"hld_auc": float("nan"),
|
||||
"hld_acc": float("nan"),
|
||||
"eye_train_n": len(eye_train),
|
||||
"bilat_val_n": len(bilat_val),
|
||||
"bilat_holdout_n": len(holdout_bilat),
|
||||
}
|
||||
)
|
||||
continue
|
||||
|
||||
sampler = build_balanced_sampler(eye_train) if args.balanced_sampling else None
|
||||
train_loader = DataLoader(
|
||||
EyeMetaDataset(eye_train),
|
||||
batch_size=args.batch_size,
|
||||
shuffle=(sampler is None),
|
||||
sampler=sampler,
|
||||
)
|
||||
val_loader = DataLoader(BilatMetaDataset(bilat_val), batch_size=args.batch_size, shuffle=False)
|
||||
holdout_loader = (
|
||||
DataLoader(BilatMetaDataset(holdout_bilat), batch_size=args.batch_size, shuffle=False)
|
||||
if holdout_bilat
|
||||
else None
|
||||
)
|
||||
|
||||
model = MetadataOnlySingleHT(
|
||||
clinical_data=data,
|
||||
num_classes=num_classes,
|
||||
md_hidden_dim=args.md_hidden_dim,
|
||||
fusion_dim=args.fusion_dim,
|
||||
dropout=args.dropout,
|
||||
use_se=bool(args.use_se),
|
||||
se_reduction=int(args.se_reduction),
|
||||
se_pre_norm=bool(args.se_pre_norm),
|
||||
).to(device)
|
||||
opt = torch.optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay)
|
||||
|
||||
best_auc = -1.0
|
||||
best_epoch = 0
|
||||
best_phase = ""
|
||||
best_state = None
|
||||
epoch_log_rows = []
|
||||
|
||||
print(
|
||||
f"\n[fold {fold_idx+1}/{args.n_splits}] "
|
||||
f"eye_train_n={len(eye_train)} bilat_val_n={len(bilat_val)} "
|
||||
f"holdout_n={len(holdout_bilat)} warmup={warm_tower}+{warm_fused} total={total_epochs}",
|
||||
flush=True,
|
||||
)
|
||||
|
||||
for ep in range(total_epochs):
|
||||
phase, main_ep = _phase_for_epoch(ep, warm_tower, warm_fused, int(args.epochs))
|
||||
tr_loss, tr_acc = train_single_epoch(
|
||||
model,
|
||||
train_loader,
|
||||
opt,
|
||||
device,
|
||||
phase=phase,
|
||||
bcd_prob=float(args.bcd_prob),
|
||||
)
|
||||
|
||||
_, p_val, val_auc, val_acc = _evaluate_single(
|
||||
model,
|
||||
val_loader,
|
||||
device,
|
||||
num_classes,
|
||||
aggregate_patient=aggregate_patient,
|
||||
)
|
||||
|
||||
hld_auc_ep = float("nan")
|
||||
hld_acc_ep = float("nan")
|
||||
if holdout_loader is not None:
|
||||
_, _, hld_auc_ep, hld_acc_ep = _evaluate_single(
|
||||
model, holdout_loader, device, num_classes,
|
||||
aggregate_patient=aggregate_patient,
|
||||
)
|
||||
|
||||
is_main = phase == "main"
|
||||
if is_main and (not np.isnan(val_auc)) and val_auc > best_auc:
|
||||
best_auc = float(val_auc)
|
||||
best_state = copy.deepcopy(model.state_dict())
|
||||
best_epoch = ep + 1
|
||||
best_phase = phase
|
||||
|
||||
epoch_log_rows.append({
|
||||
"epoch": ep + 1,
|
||||
"phase": phase,
|
||||
"train_loss": float(tr_loss),
|
||||
"train_acc": float(tr_acc),
|
||||
"val_auc": float(val_auc),
|
||||
"val_acc": float(val_acc),
|
||||
"hld_auc": float(hld_auc_ep),
|
||||
"hld_acc": float(hld_acc_ep),
|
||||
})
|
||||
|
||||
if ep == 0 or (ep + 1) % 10 == 0 or (ep + 1) == total_epochs:
|
||||
print(
|
||||
f" ep {ep+1:>3}/{total_epochs} [{phase}:{main_ep}/{args.epochs}] "
|
||||
f"loss={tr_loss:.4f} acc={tr_acc:.4f} "
|
||||
f"val_auc={val_auc:.4f} val_acc={val_acc:.4f} "
|
||||
f"best_auc={best_auc:.4f}",
|
||||
flush=True,
|
||||
)
|
||||
|
||||
import pandas as _pd
|
||||
_pd.DataFrame(epoch_log_rows).to_csv(fold_dir / "epoch_log.csv", index=False)
|
||||
|
||||
if best_state is not None:
|
||||
model.load_state_dict(best_state)
|
||||
|
||||
y_val, p_val, val_auc, val_acc = _evaluate_single(
|
||||
model,
|
||||
val_loader,
|
||||
device,
|
||||
num_classes,
|
||||
aggregate_patient=aggregate_patient,
|
||||
)
|
||||
|
||||
if args.tower_mode == "single":
|
||||
probs_name = "probs_classic.npy"
|
||||
probs_h_name = "probs_classic_holdout.npy"
|
||||
else:
|
||||
probs_name = "probs_ensemble.npy"
|
||||
probs_h_name = "probs_ensemble_holdout.npy"
|
||||
|
||||
np.save(fold_dir / "y_true.npy", y_val)
|
||||
np.save(fold_dir / probs_name, p_val)
|
||||
|
||||
hld_auc = float("nan")
|
||||
hld_acc = float("nan")
|
||||
if holdout_loader is not None:
|
||||
y_h, p_h, hld_auc, hld_acc = _evaluate_single(
|
||||
model,
|
||||
holdout_loader,
|
||||
device,
|
||||
num_classes,
|
||||
aggregate_patient=aggregate_patient,
|
||||
)
|
||||
np.save(fold_dir / "y_true_holdout.npy", y_h)
|
||||
np.save(fold_dir / probs_h_name, p_h)
|
||||
|
||||
print(
|
||||
f" [fold {fold_idx+1}] best_epoch={best_epoch} best_auc={best_auc:.4f} "
|
||||
f"val_auc={val_auc:.4f} val_acc={val_acc:.4f} "
|
||||
f"hld_auc={hld_auc:.4f} hld_acc={hld_acc:.4f}",
|
||||
flush=True,
|
||||
)
|
||||
|
||||
fold_metrics.append(
|
||||
{
|
||||
"fold": fold_idx,
|
||||
"best_epoch": best_epoch,
|
||||
"best_phase": best_phase,
|
||||
"val_auc": float(val_auc),
|
||||
"val_acc": float(val_acc),
|
||||
"hld_auc": float(hld_auc),
|
||||
"hld_acc": float(hld_acc),
|
||||
"eye_train_n": len(eye_train),
|
||||
"bilat_val_n": len(bilat_val),
|
||||
"bilat_holdout_n": len(holdout_bilat),
|
||||
}
|
||||
)
|
||||
|
||||
val_aucs = [m["val_auc"] for m in fold_metrics if not np.isnan(m["val_auc"])]
|
||||
hld_aucs = [m["hld_auc"] for m in fold_metrics if not np.isnan(m["hld_auc"])]
|
||||
if val_aucs:
|
||||
print(f"\nMean val AUC: {np.mean(val_aucs):.4f} ± {np.std(val_aucs):.4f}", flush=True)
|
||||
if hld_aucs:
|
||||
print(f"Mean hld AUC: {np.mean(hld_aucs):.4f} ± {np.std(hld_aucs):.4f}", flush=True)
|
||||
|
||||
summary = {
|
||||
"run_name": args.run_name,
|
||||
"eval_mode": args.eval_mode,
|
||||
"tower_mode": args.tower_mode,
|
||||
"bridge_mode": "metadata_only",
|
||||
"model": "MetadataOnlySingleHT",
|
||||
"epochs": int(args.epochs),
|
||||
"warmup_tower_epochs": warm_tower,
|
||||
"warmup_fused_epochs": warm_fused,
|
||||
"md_hidden_dim": int(args.md_hidden_dim),
|
||||
"fusion_dim": int(args.fusion_dim),
|
||||
"dropout": float(args.dropout),
|
||||
"lr": float(args.lr),
|
||||
"weight_decay": float(args.weight_decay),
|
||||
"bcd_prob": float(args.bcd_prob),
|
||||
"balanced_sampling": bool(args.balanced_sampling),
|
||||
"feature_dim": int(data.feature_dim),
|
||||
"timestamp": time.strftime("%Y%m%d_%H%M%S"),
|
||||
"fold_metrics": fold_metrics,
|
||||
"val_auc_mean": float(np.mean(val_aucs)) if val_aucs else None,
|
||||
"val_auc_std": float(np.std(val_aucs)) if val_aucs else None,
|
||||
"hld_auc_mean": float(np.mean(hld_aucs)) if hld_aucs else None,
|
||||
"hld_auc_std": float(np.std(hld_aucs)) if hld_aucs else None,
|
||||
}
|
||||
(out_root / "summary.json").write_text(json.dumps(summary, indent=2))
|
||||
print(f"\nOutputs written to: {out_root}", flush=True)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,28 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
"""CLI wrapper for the V2 hypertower pipeline using V2HyperTower directly."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
import sys
|
||||
|
||||
# ensure repo root on path
|
||||
REPO_ROOT = Path(__file__).resolve().parents[3]
|
||||
if str(REPO_ROOT) not in sys.path:
|
||||
sys.path.insert(0, str(REPO_ROOT))
|
||||
|
||||
from classes.v2.v2_hypertower import V2HyperTower
|
||||
|
||||
|
||||
def run_cli(cli_args=None):
|
||||
parser = V2HyperTower.build_parser()
|
||||
args = parser.parse_args(cli_args)
|
||||
V2HyperTower(args).run()
|
||||
|
||||
|
||||
def main():
|
||||
run_cli()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,21 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
"""CLI wrapper for the V2 three-mode comparison (classic/ensemble/bilateral)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
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.v2.v2_hypertower import V2ModeComparator
|
||||
|
||||
|
||||
def main():
|
||||
V2ModeComparator.run()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,496 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Aggregate and visualise results from a 10× repeated 5-fold CV run.
|
||||
|
||||
Reads probs / y_true from every rep/fold directory, computes per-fold
|
||||
metrics, and produces:
|
||||
|
||||
outputs/
|
||||
fold_metrics.csv — one row per (rep, fold, eval_mode)
|
||||
rep_metrics.csv — one row per (rep, eval_mode): mean over 5 folds
|
||||
overall_summary.txt — mean ± SD and 95% CI printed to console + file
|
||||
{eval_mode}_auc_violin.png
|
||||
{eval_mode}_roc_mean.png — mean ± 1 SD OVR ROC (all classes or class 1)
|
||||
{eval_mode}_holdout_roc_mean.png
|
||||
|
||||
Holdout metrics are extracted from the rep-level predictions.npz using the
|
||||
best_epoch recorded in summary.json, ensemble-averaged over od_fused + os_fused
|
||||
heads, giving 50 fold-level holdout AUC values (5 folds × 10 reps).
|
||||
|
||||
Usage
|
||||
-----
|
||||
python scripts/output_analysis/aggregate_10x5cv.py \
|
||||
--run-root analysis_data/pipeline_10x5 \
|
||||
--eval-modes binary multiclass \
|
||||
--out analysis_data/pipeline_10x5/aggregate
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
import matplotlib
|
||||
matplotlib.use("Agg")
|
||||
import matplotlib.pyplot as plt
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
from sklearn.metrics import roc_auc_score, accuracy_score, roc_curve, auc as sk_auc
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Config
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_CLASS_NAMES = {
|
||||
"binary": ["Healthy", "Glaucoma"],
|
||||
"multiclass": ["Healthy", "Glaucoma", "Suspect"],
|
||||
}
|
||||
|
||||
# Probe files in preference order (first found wins)
|
||||
# probs_fused = simple OD/OS softmax average (ensemble head — primary metric)
|
||||
# probs_fused_head = learned logit-level fusion head (worse on average; kept as fallback)
|
||||
_PROBS_PRIORITY = ["probs_fused.npy", "probs_fused_head.npy"]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _find_probs(fold_dir: Path) -> Path | None:
|
||||
for name in _PROBS_PRIORITY:
|
||||
p = fold_dir / name
|
||||
if p.exists():
|
||||
return p
|
||||
return None
|
||||
|
||||
|
||||
def _load_holdout_summary(mode_dir: Path) -> dict | None:
|
||||
"""
|
||||
Read rep-level holdout metrics from summary.json.
|
||||
|
||||
Returns dict with keys auc_mean, auc_std, acc_mean (may be None if missing).
|
||||
"""
|
||||
summary_path = mode_dir / "summary.json"
|
||||
if not summary_path.exists():
|
||||
return None
|
||||
try:
|
||||
summary = json.loads(summary_path.read_text())
|
||||
return summary.get("mode_summary", {}).get("ensemble_holdout")
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
def _auc_macro(y: np.ndarray, p: np.ndarray, num_classes: int) -> float:
|
||||
try:
|
||||
if num_classes == 2:
|
||||
return float(roc_auc_score(y, p[:, 1]))
|
||||
return float(roc_auc_score(y, p, multi_class="ovr", average="macro"))
|
||||
except Exception:
|
||||
return float("nan")
|
||||
|
||||
|
||||
def _per_class_roc(y: np.ndarray, p: np.ndarray) -> dict[int, dict]:
|
||||
out: dict[int, dict] = {}
|
||||
for k in range(p.shape[1]):
|
||||
yb = (y == k).astype(np.uint8)
|
||||
if yb.sum() == 0 or yb.sum() == len(yb):
|
||||
continue
|
||||
fpr, tpr, _ = roc_curve(yb, p[:, k])
|
||||
out[k] = {"fpr": fpr, "tpr": tpr, "auc": sk_auc(fpr, tpr)}
|
||||
return out
|
||||
|
||||
|
||||
def _ci95(values: np.ndarray) -> tuple[float, float]:
|
||||
"""95% CI via t-distribution (two-sided)."""
|
||||
from scipy import stats as scipy_stats
|
||||
if len(values) < 2:
|
||||
return (float("nan"), float("nan"))
|
||||
ci = scipy_stats.t.interval(0.95, df=len(values) - 1,
|
||||
loc=np.mean(values), scale=scipy_stats.sem(values))
|
||||
return float(ci[0]), float(ci[1])
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Data loading
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def load_all_folds(run_root: Path, eval_modes: list[str]) -> pd.DataFrame:
|
||||
rows = []
|
||||
rep_dirs = sorted(
|
||||
[d for d in run_root.iterdir() if d.is_dir() and d.name.startswith("rep")],
|
||||
key=lambda d: d.name,
|
||||
)
|
||||
if not rep_dirs:
|
||||
raise SystemExit(f"No rep* directories found in {run_root}")
|
||||
|
||||
for rep_dir in rep_dirs:
|
||||
for eval_mode in eval_modes:
|
||||
tower_mode = "ensemble"
|
||||
mode_dir = rep_dir / eval_mode / tower_mode
|
||||
if not mode_dir.exists():
|
||||
print(f" [skip] {mode_dir} not found")
|
||||
continue
|
||||
num_classes = 2 if eval_mode == "binary" else 3
|
||||
|
||||
fold_dirs = sorted(
|
||||
[d for d in mode_dir.iterdir() if d.is_dir() and d.name.startswith("fold")],
|
||||
key=lambda d: int(d.name[4:]),
|
||||
)
|
||||
for fold_dir in fold_dirs:
|
||||
fold_idx = int(fold_dir.name[4:])
|
||||
y_path = fold_dir / "y_true.npy"
|
||||
p_path = _find_probs(fold_dir)
|
||||
|
||||
if y_path is None or not y_path.exists() or p_path is None:
|
||||
print(f" [skip] {rep_dir.name}/{eval_mode}/fold{fold_idx}: missing files")
|
||||
continue
|
||||
|
||||
y = np.load(y_path)
|
||||
p = np.load(p_path)
|
||||
|
||||
if eval_mode == "binary":
|
||||
mask = np.isin(y, [0, 1])
|
||||
y, p = y[mask], p[mask]
|
||||
if p.shape[1] > 2:
|
||||
p = p[:, :2]
|
||||
|
||||
auc_macro = _auc_macro(y, p, num_classes)
|
||||
acc = float(accuracy_score(y, p.argmax(1)))
|
||||
|
||||
row = {
|
||||
"rep": rep_dir.name,
|
||||
"fold": fold_idx,
|
||||
"eval_mode": eval_mode,
|
||||
"probs_file": p_path.name,
|
||||
"auc_macro": auc_macro,
|
||||
"acc": acc,
|
||||
"n": len(y),
|
||||
}
|
||||
|
||||
# Per-class AUC
|
||||
for k in range(num_classes):
|
||||
yb = (y == k).astype(np.uint8)
|
||||
if yb.sum() > 0 and yb.sum() < len(yb):
|
||||
try:
|
||||
row[f"auc_class{k}"] = float(roc_auc_score(yb, p[:, k]))
|
||||
except Exception:
|
||||
row[f"auc_class{k}"] = float("nan")
|
||||
else:
|
||||
row[f"auc_class{k}"] = float("nan")
|
||||
|
||||
rows.append(row)
|
||||
|
||||
# ---- holdout metrics from rep-level summary.json ----
|
||||
# Holdout probs are not stored per-fold; only aggregated stats are saved.
|
||||
# We attach the rep-level mean to each fold row (same value repeated),
|
||||
# and also add a single rep-level summary row (fold=-1).
|
||||
hld_summary = _load_holdout_summary(mode_dir)
|
||||
if hld_summary:
|
||||
hld_auc = hld_summary.get("auc_mean", float("nan"))
|
||||
hld_auc_std = hld_summary.get("auc_std", float("nan"))
|
||||
hld_acc = hld_summary.get("acc_mean", float("nan"))
|
||||
for row in rows:
|
||||
if row["rep"] == rep_dir.name and row["eval_mode"] == eval_mode:
|
||||
row["hld_auc_macro"] = hld_auc
|
||||
row["hld_acc"] = hld_acc
|
||||
# Also store a rep-level holdout row (fold=-1) for direct rep-level analysis
|
||||
rows.append({
|
||||
"rep": rep_dir.name,
|
||||
"fold": -1,
|
||||
"eval_mode": eval_mode,
|
||||
"probs_file": "summary.json",
|
||||
"auc_macro": float("nan"),
|
||||
"acc": float("nan"),
|
||||
"n": float("nan"),
|
||||
"hld_auc_macro": hld_auc,
|
||||
"hld_auc_std_within_rep": hld_auc_std,
|
||||
"hld_acc": hld_acc,
|
||||
})
|
||||
|
||||
return pd.DataFrame(rows)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Plotting
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _violin(fold_df: pd.DataFrame, eval_mode: str, out_dir: Path) -> None:
|
||||
sub = fold_df[fold_df["eval_mode"] == eval_mode].copy()
|
||||
num_classes = 2 if eval_mode == "binary" else 3
|
||||
class_names = _CLASS_NAMES[eval_mode]
|
||||
|
||||
auc_cols = ["auc_macro"] + [f"auc_class{k}" for k in range(num_classes)]
|
||||
labels = ["Macro AUC"] + [f"AUC {class_names[k]}" for k in range(num_classes)]
|
||||
present = [(c, l) for c, l in zip(auc_cols, labels) if c in sub.columns]
|
||||
|
||||
data = [sub[c].dropna().values for c, _ in present]
|
||||
labels = [l for _, l in present]
|
||||
|
||||
fig, ax = plt.subplots(figsize=(max(6, 2 * len(data)), 5))
|
||||
parts = ax.violinplot(data, showmedians=True, showextrema=True)
|
||||
for pc in parts["bodies"]:
|
||||
pc.set_alpha(0.7)
|
||||
|
||||
# Overlay individual rep means
|
||||
rep_means = sub.groupby("rep")[auc_cols[0]].mean().values
|
||||
ax.scatter(np.ones(len(rep_means)), rep_means, zorder=3,
|
||||
color="k", s=18, alpha=0.6, label="rep mean")
|
||||
|
||||
ax.set_xticks(range(1, len(labels) + 1))
|
||||
ax.set_xticklabels(labels, rotation=15, ha="right")
|
||||
ax.set_ylabel("AUC")
|
||||
ax.set_title(f"AUC distribution — {eval_mode} (10 × 5-fold, n={len(sub)})")
|
||||
ax.set_ylim(max(0, sub[auc_cols[0]].min() - 0.05), 1.02)
|
||||
ax.grid(True, axis="y", linewidth=0.4, alpha=0.5)
|
||||
ax.legend(fontsize=8)
|
||||
fig.tight_layout()
|
||||
path = out_dir / f"{eval_mode}_auc_violin.png"
|
||||
fig.savefig(path, dpi=160, bbox_inches="tight")
|
||||
plt.close(fig)
|
||||
print(f" Saved: {path}")
|
||||
|
||||
|
||||
def _mean_roc(fold_df: pd.DataFrame, eval_mode: str,
|
||||
run_root: Path, out_dir: Path) -> None:
|
||||
"""Mean ± 1 SD OVR ROC across all 50 folds."""
|
||||
sub = fold_df[fold_df["eval_mode"] == eval_mode]
|
||||
num_classes = 2 if eval_mode == "binary" else 3
|
||||
class_names = _CLASS_NAMES[eval_mode]
|
||||
|
||||
# Classes to plot (binary: class 1 only)
|
||||
plot_classes = [1] if eval_mode == "binary" else list(range(num_classes))
|
||||
|
||||
grid = np.linspace(0, 1, 501)
|
||||
fig, ax = plt.subplots(figsize=(9, 7))
|
||||
ax.plot([0, 1], [0, 1], linestyle="--", linewidth=1, color="grey")
|
||||
|
||||
for k in plot_classes:
|
||||
tprs, aucs = [], []
|
||||
for _, row in sub.iterrows():
|
||||
rep_dir = run_root / row["rep"]
|
||||
fold_dir = rep_dir / eval_mode / "ensemble" / f"fold{int(row['fold'])}"
|
||||
y_path = fold_dir / "y_true.npy"
|
||||
p_path = _find_probs(fold_dir)
|
||||
if not y_path.exists() or p_path is None:
|
||||
continue
|
||||
y = np.load(y_path)
|
||||
p = np.load(p_path)
|
||||
if eval_mode == "binary":
|
||||
mask = np.isin(y, [0, 1])
|
||||
y, p = y[mask], p[mask]
|
||||
if p.shape[1] > 2:
|
||||
p = p[:, :2]
|
||||
yb = (y == k).astype(np.uint8)
|
||||
if yb.sum() == 0 or yb.sum() == len(yb):
|
||||
continue
|
||||
fpr, tpr, _ = roc_curve(yb, p[:, k])
|
||||
tprs.append(np.interp(grid, fpr, tpr))
|
||||
aucs.append(sk_auc(fpr, tpr))
|
||||
|
||||
if not tprs:
|
||||
continue
|
||||
arr = np.vstack(tprs)
|
||||
mean = arr.mean(0)
|
||||
std = arr.std(0)
|
||||
cname = class_names[k]
|
||||
lbl = f"{cname} AUC {np.nanmean(aucs):.3f} ± {np.nanstd(aucs):.3f}"
|
||||
line, = ax.plot(grid, mean, linewidth=2, label=lbl)
|
||||
ax.fill_between(grid,
|
||||
np.clip(mean - std, 0, 1),
|
||||
np.clip(mean + std, 0, 1),
|
||||
alpha=0.15, color=line.get_color())
|
||||
|
||||
ax.set_xlabel("False Positive Rate")
|
||||
ax.set_ylabel("True Positive Rate")
|
||||
ax.set_title(f"Mean ± 1 SD OVR ROC — {eval_mode} (10 × 5-fold)")
|
||||
ax.legend(loc="lower right", fontsize=9)
|
||||
fig.tight_layout()
|
||||
path = out_dir / f"{eval_mode}_roc_mean.png"
|
||||
fig.savefig(path, dpi=160, bbox_inches="tight")
|
||||
plt.close(fig)
|
||||
print(f" Saved: {path}")
|
||||
|
||||
|
||||
def _holdout_stability(fold_df: pd.DataFrame, eval_mode: str, out_dir: Path) -> None:
|
||||
"""Bar chart of per-rep holdout AUC (mean across folds within rep ± within-rep SD)."""
|
||||
# use the rep-level rows (fold == -1) which have hld_auc_std_within_rep
|
||||
rep_rows = fold_df[(fold_df["eval_mode"] == eval_mode) & (fold_df["fold"] == -1)].copy()
|
||||
if rep_rows.empty or "hld_auc_macro" not in rep_rows.columns:
|
||||
print(f" [skip] no holdout data for {eval_mode}")
|
||||
return
|
||||
rep_rows = rep_rows.sort_values("rep")
|
||||
|
||||
fig, ax = plt.subplots(figsize=(max(6, len(rep_rows) * 0.9), 4))
|
||||
x = np.arange(len(rep_rows))
|
||||
yerr = rep_rows.get("hld_auc_std_within_rep", pd.Series([0]*len(rep_rows))).fillna(0).values
|
||||
ax.bar(x, rep_rows["hld_auc_macro"].values, yerr=yerr,
|
||||
capsize=4, color="darkorange", alpha=0.8)
|
||||
grand_mean = rep_rows["hld_auc_macro"].mean()
|
||||
ax.axhline(grand_mean, linestyle="--", color="crimson",
|
||||
linewidth=1.2, label=f"grand mean = {grand_mean:.3f}")
|
||||
ax.set_xticks(x)
|
||||
ax.set_xticklabels(rep_rows["rep"].values, rotation=30, ha="right")
|
||||
ax.set_ylabel("Holdout macro AUC (mean ± within-rep SD)")
|
||||
ax.set_title(f"Per-rep holdout stability — {eval_mode}")
|
||||
ymin = max(0, rep_rows["hld_auc_macro"].min() - 0.05)
|
||||
ax.set_ylim(ymin, 1.02)
|
||||
ax.legend(fontsize=9)
|
||||
ax.grid(True, axis="y", linewidth=0.4, alpha=0.5)
|
||||
fig.tight_layout()
|
||||
path = out_dir / f"{eval_mode}_holdout_stability.png"
|
||||
fig.savefig(path, dpi=160, bbox_inches="tight")
|
||||
plt.close(fig)
|
||||
print(f" Saved: {path}")
|
||||
|
||||
|
||||
def _rep_stability(fold_df: pd.DataFrame, eval_mode: str, out_dir: Path) -> None:
|
||||
"""Bar chart of per-rep mean macro AUC with ± 1 SD error bars."""
|
||||
sub = fold_df[fold_df["eval_mode"] == eval_mode]
|
||||
rep_stats = sub.groupby("rep")["auc_macro"].agg(["mean", "std"]).reset_index()
|
||||
rep_stats = rep_stats.sort_values("rep")
|
||||
|
||||
fig, ax = plt.subplots(figsize=(max(6, len(rep_stats) * 0.9), 4))
|
||||
x = np.arange(len(rep_stats))
|
||||
ax.bar(x, rep_stats["mean"], yerr=rep_stats["std"],
|
||||
capsize=4, color="steelblue", alpha=0.8)
|
||||
ax.axhline(rep_stats["mean"].mean(), linestyle="--", color="crimson",
|
||||
linewidth=1.2, label=f"grand mean = {rep_stats['mean'].mean():.3f}")
|
||||
ax.set_xticks(x)
|
||||
ax.set_xticklabels(rep_stats["rep"], rotation=30, ha="right")
|
||||
ax.set_ylabel("Mean macro AUC (5 folds)")
|
||||
ax.set_title(f"Per-rep stability — {eval_mode}")
|
||||
ymin = max(0, rep_stats["mean"].min() - rep_stats["std"].max() - 0.02)
|
||||
ax.set_ylim(ymin, 1.02)
|
||||
ax.legend(fontsize=9)
|
||||
ax.grid(True, axis="y", linewidth=0.4, alpha=0.5)
|
||||
fig.tight_layout()
|
||||
path = out_dir / f"{eval_mode}_rep_stability.png"
|
||||
fig.savefig(path, dpi=160, bbox_inches="tight")
|
||||
plt.close(fig)
|
||||
print(f" Saved: {path}")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Summary text
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _print_summary(fold_df: pd.DataFrame, eval_modes: list[str]) -> str:
|
||||
lines = ["=" * 60, "10 × 5-fold CV — aggregate summary", "=" * 60]
|
||||
for eval_mode in eval_modes:
|
||||
sub = fold_df[(fold_df["eval_mode"] == eval_mode) & (fold_df["fold"] >= 0)]
|
||||
if sub.empty:
|
||||
continue
|
||||
num_classes = 2 if eval_mode == "binary" else 3
|
||||
class_names = _CLASS_NAMES[eval_mode]
|
||||
lines.append(f"\n--- {eval_mode.upper()} ---")
|
||||
lines.append(f" n_folds = {len(sub)}")
|
||||
|
||||
lines.append(" [Validation]")
|
||||
for col, label in [("auc_macro", "Macro AUC"), ("acc", "Accuracy")]:
|
||||
if col not in sub.columns:
|
||||
continue
|
||||
vals = sub[col].dropna().values
|
||||
ci_lo, ci_hi = _ci95(vals)
|
||||
lines.append(
|
||||
f" {label:18s}: {vals.mean():.4f} ± {vals.std():.4f}"
|
||||
f" 95% CI [{ci_lo:.4f}, {ci_hi:.4f}]"
|
||||
)
|
||||
|
||||
for k in range(num_classes):
|
||||
col = f"auc_class{k}"
|
||||
if col not in sub.columns:
|
||||
continue
|
||||
vals = sub[col].dropna().values
|
||||
if len(vals) == 0:
|
||||
continue
|
||||
ci_lo, ci_hi = _ci95(vals)
|
||||
lines.append(
|
||||
f" AUC {class_names[k]:12s}: {vals.mean():.4f} ± {vals.std():.4f}"
|
||||
f" 95% CI [{ci_lo:.4f}, {ci_hi:.4f}]"
|
||||
)
|
||||
|
||||
# Between-rep variance (val)
|
||||
rep_means = sub.groupby("rep")["auc_macro"].mean().values
|
||||
lines.append(
|
||||
f" Rep-mean AUC (n={len(rep_means)}): "
|
||||
f"{rep_means.mean():.4f} ± {rep_means.std():.4f}"
|
||||
f" (between-rep SD = {rep_means.std():.4f})"
|
||||
)
|
||||
|
||||
# Holdout — use rep-level rows (fold == -1)
|
||||
rep_hld = fold_df[
|
||||
(fold_df["eval_mode"] == eval_mode) &
|
||||
(fold_df["fold"] == -1) &
|
||||
fold_df["hld_auc_macro"].notna()
|
||||
]["hld_auc_macro"].values if "hld_auc_macro" in fold_df.columns else np.array([])
|
||||
|
||||
if len(rep_hld) > 0:
|
||||
lines.append(" [Holdout] (rep-level means, n_reps={})".format(len(rep_hld)))
|
||||
ci_lo, ci_hi = _ci95(rep_hld)
|
||||
lines.append(
|
||||
f" {'Macro AUC':18s}: {rep_hld.mean():.4f} ± {rep_hld.std():.4f}"
|
||||
f" 95% CI [{ci_lo:.4f}, {ci_hi:.4f}]"
|
||||
)
|
||||
|
||||
lines.append("=" * 60)
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# CLI
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def main() -> None:
|
||||
ap = argparse.ArgumentParser(
|
||||
description=__doc__,
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||||
)
|
||||
ap.add_argument("--run-root", default="analysis_data/pipeline_10x5",
|
||||
help="Root directory containing rep* sub-directories.")
|
||||
ap.add_argument("--eval-modes", nargs="+",
|
||||
choices=["binary", "multiclass"],
|
||||
default=["binary", "multiclass"])
|
||||
ap.add_argument("--out", default=None,
|
||||
help="Output directory for plots and CSVs "
|
||||
"(default: {run-root}/aggregate).")
|
||||
args = ap.parse_args()
|
||||
|
||||
run_root = Path(args.run_root)
|
||||
out_dir = Path(args.out) if args.out else run_root / "aggregate"
|
||||
out_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
print("Loading fold data...")
|
||||
fold_df = load_all_folds(run_root, args.eval_modes)
|
||||
if fold_df.empty:
|
||||
raise SystemExit("No data loaded — check --run-root.")
|
||||
|
||||
fold_df.to_csv(out_dir / "fold_metrics.csv", index=False)
|
||||
print(f" Saved fold_metrics.csv ({len(fold_df)} rows)")
|
||||
|
||||
rep_df = (fold_df.groupby(["rep", "eval_mode"])
|
||||
[["auc_macro", "acc"] +
|
||||
[c for c in fold_df.columns if c.startswith("auc_class")]]
|
||||
.mean()
|
||||
.reset_index())
|
||||
rep_df.to_csv(out_dir / "rep_metrics.csv", index=False)
|
||||
print(f" Saved rep_metrics.csv ({len(rep_df)} rows)")
|
||||
|
||||
summary_text = _print_summary(fold_df, args.eval_modes)
|
||||
print("\n" + summary_text)
|
||||
(out_dir / "overall_summary.txt").write_text(summary_text + "\n")
|
||||
print(f"\n Saved overall_summary.txt")
|
||||
|
||||
print("\nGenerating plots...")
|
||||
for eval_mode in args.eval_modes:
|
||||
if fold_df[fold_df["eval_mode"] == eval_mode].empty:
|
||||
continue
|
||||
_violin(fold_df, eval_mode, out_dir)
|
||||
_mean_roc(fold_df, eval_mode, run_root, out_dir)
|
||||
_rep_stability(fold_df, eval_mode, out_dir)
|
||||
_holdout_stability(fold_df, eval_mode, out_dir)
|
||||
|
||||
print(f"\nAll outputs written to: {out_dir}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,165 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Inspect saved validation/holdout logits for a multifold run.
|
||||
Prints per-class AUCs and sample counts so we can sanity-check unusually high scores.
|
||||
Can also print per-fold confusion matrices.
|
||||
|
||||
Example:
|
||||
python scripts/fold_confusion_matrix.py \
|
||||
--run-dir analysis_data/1030_Balanced_Unet_Perimg_Resnet_SE16NormB_SE16NormT_multi_fused/1030_Balanced_Unet_Perimg_Resnet_SE16NormB_SE16NormT_multi_fused_20251030_091842 \
|
||||
--head fused
|
||||
python scripts/fold_confusion_matrix.py --run-dir ... --head fused --use-holdout --confusion
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import math
|
||||
from pathlib import Path
|
||||
from typing import Dict, List
|
||||
|
||||
import numpy as np
|
||||
from sklearn.metrics import roc_auc_score, confusion_matrix
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
ap = argparse.ArgumentParser(description="Inspect saved logits for a run and report per-class AUCs.")
|
||||
ap.add_argument("--run-dir", required=True, type=Path, help="Path to the run directory under analysis_data.")
|
||||
ap.add_argument("--head", choices=["fused", "image", "metadata"], default="fused",
|
||||
help="Which prediction head's saved probabilities to load.")
|
||||
ap.add_argument("--use-holdout", action="store_true",
|
||||
help="Look for *_holdout.npy dumps instead of validation splits.")
|
||||
ap.add_argument("--class-names", nargs="*", default=None,
|
||||
help="Optional override for class labels (order should match numeric labels).")
|
||||
ap.add_argument("--macro", action="store_true", help="Also print macro-average AUC across classes.")
|
||||
ap.add_argument("--confusion", action="store_true", help="Print confusion matrix for each fold.")
|
||||
return ap.parse_args()
|
||||
|
||||
|
||||
def load_cli_args(run_dir: Path) -> Dict:
|
||||
path = run_dir / "cli_args.json"
|
||||
if not path.exists():
|
||||
raise FileNotFoundError(f"Missing cli_args.json in {run_dir}")
|
||||
with path.open("r", encoding="utf-8") as fh:
|
||||
return json.load(fh)
|
||||
|
||||
|
||||
def find_fold_files(run_dir: Path, suffix: str) -> Dict[int, Dict[str, Path]]:
|
||||
files: Dict[int, Dict[str, Path]] = {}
|
||||
for y_file in run_dir.glob(f"fold*_y_true{suffix}.npy"):
|
||||
fold_str = y_file.stem.split("_")[0].replace("fold", "")
|
||||
try:
|
||||
fold_idx = int(fold_str)
|
||||
except ValueError:
|
||||
continue
|
||||
files.setdefault(fold_idx, {})["y_true"] = y_file
|
||||
for head_key, glob_pat in [
|
||||
("fused", f"fold*_probs_fused{suffix}.npy"),
|
||||
("image", f"fold*_probs_img{suffix}.npy"),
|
||||
("metadata", f"fold*_probs_md{suffix}.npy"),
|
||||
]:
|
||||
for p_file in run_dir.glob(glob_pat):
|
||||
fold_str = p_file.stem.split("_")[0].replace("fold", "")
|
||||
try:
|
||||
fold_idx = int(fold_str)
|
||||
except ValueError:
|
||||
continue
|
||||
files.setdefault(fold_idx, {})[head_key] = p_file
|
||||
return files
|
||||
|
||||
|
||||
def compute_auc(y_true: np.ndarray, probs: np.ndarray, class_names: List[str], macro: bool) -> List[int]:
|
||||
num_classes = probs.shape[1]
|
||||
unique = np.unique(y_true)
|
||||
print(f" classes present: {sorted(unique.tolist())}")
|
||||
|
||||
aucs = []
|
||||
seen_classes: List[int] = []
|
||||
for cls in range(num_classes):
|
||||
name = class_names[cls] if cls < len(class_names) else f"class_{cls}"
|
||||
mask = (y_true == cls)
|
||||
pos = int(mask.sum())
|
||||
neg = len(y_true) - pos
|
||||
if pos == 0 or neg == 0:
|
||||
print(f" {name:<15} -> insufficient positives/negatives (pos={pos}, neg={neg}); skipping AUC")
|
||||
continue
|
||||
try:
|
||||
auc = roc_auc_score((y_true == cls).astype(int), probs[:, cls])
|
||||
except ValueError as exc:
|
||||
print(f" {name:<15} -> AUC error: {exc}")
|
||||
continue
|
||||
aucs.append(auc)
|
||||
seen_classes.append(cls)
|
||||
print(f" {name:<15} -> AUC={auc:.4f} (pos={pos}, neg={neg})")
|
||||
|
||||
if macro and aucs:
|
||||
mean = float(np.mean(aucs))
|
||||
std = float(np.std(aucs, ddof=0)) if len(aucs) > 1 else math.nan
|
||||
print(f" macro AUC across reported classes: {mean:.4f} (std={std:.4f})")
|
||||
return seen_classes
|
||||
|
||||
|
||||
def print_confusion(y_true: np.ndarray, probs: np.ndarray, class_names: List[str]) -> None:
|
||||
num_classes = probs.shape[1]
|
||||
preds = probs.argmax(axis=1)
|
||||
labels = list(range(num_classes))
|
||||
cm = confusion_matrix(y_true, preds, labels=labels)
|
||||
names = [class_names[i] if i < len(class_names) else f"class_{i}" for i in labels]
|
||||
header = " " * 14 + "".join(f"{name:>12}" for name in names)
|
||||
print(" Confusion matrix (rows=true, cols=pred):")
|
||||
print(header)
|
||||
for idx, row in enumerate(cm):
|
||||
label = names[idx]
|
||||
row_str = "".join(f"{int(val):>12}" for val in row)
|
||||
print(f" {label:<12}{row_str}")
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
run_dir = args.run_dir.resolve()
|
||||
if not run_dir.exists():
|
||||
raise FileNotFoundError(run_dir)
|
||||
|
||||
cli_args = load_cli_args(run_dir)
|
||||
eval_mode = cli_args.get("eval_mode", "multiclass")
|
||||
if args.class_names:
|
||||
class_names = args.class_names
|
||||
else:
|
||||
if eval_mode == "binary":
|
||||
class_names = ["Healthy", "Glaucoma"]
|
||||
else:
|
||||
class_names = cli_args.get("class_names") or ["Healthy", "Glaucoma", "Suspect"]
|
||||
|
||||
suffix = "_holdout" if args.use_holdout else ""
|
||||
files = find_fold_files(run_dir, suffix)
|
||||
if not files:
|
||||
raise SystemExit(f"No saved probability files matching suffix '{suffix}' found in {run_dir}. "
|
||||
"Run scripts/rebuild_run_best_plots.py first if needed.")
|
||||
|
||||
print(f"[info] Inspecting head='{args.head}' ({'holdout' if args.use_holdout else 'validation'})")
|
||||
for fold_idx in sorted(files.keys()):
|
||||
fold = files[fold_idx]
|
||||
if "y_true" not in fold:
|
||||
print(f"[warning] Fold {fold_idx}: missing y_true file; skipping.")
|
||||
continue
|
||||
head_key = {
|
||||
"fused": "fused",
|
||||
"image": "image",
|
||||
"metadata": "metadata",
|
||||
}[args.head]
|
||||
prob_path = fold.get(head_key)
|
||||
if prob_path is None:
|
||||
print(f"[warning] Fold {fold_idx}: missing probability file for head '{args.head}'; skipping.")
|
||||
continue
|
||||
|
||||
y_true = np.load(fold["y_true"])
|
||||
probs = np.load(prob_path)
|
||||
print(f"\n Fold {fold_idx} -> samples={len(y_true)} file={prob_path.name}")
|
||||
compute_auc(y_true, probs, class_names, args.macro)
|
||||
if args.confusion:
|
||||
print_confusion(y_true, probs, class_names)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,241 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Aggregate raw GradCAM heatmaps across all folds for a run.
|
||||
|
||||
For each combination of (eye, class, correct/incorrect) computes:
|
||||
- mean heatmap
|
||||
- std heatmap
|
||||
- count
|
||||
|
||||
Also computes a scalar per patient: fraction of GradCAM attention mass that
|
||||
falls within the expert-segmented optic disc region (from GT contour files),
|
||||
using the manifest.csv to locate the contour for each patient/eye.
|
||||
|
||||
Outputs
|
||||
-------
|
||||
{out_dir}/mean_heatmaps.npz
|
||||
Keys: {eye}_{class_name}_{correct|incorrect}_{mean|std|count}
|
||||
e.g. OD_Glaucoma_correct_mean shape (224, 224)
|
||||
|
||||
{out_dir}/attention_stats.csv
|
||||
per-patient scalars: patient_id, fold, eye, true_name, pred_name,
|
||||
correct, confidence, disc_frac, entropy
|
||||
|
||||
Usage
|
||||
-----
|
||||
python scripts/output_analysis/explainability/aggregate_gradcam.py \
|
||||
--run-dir analysis_data/pipeline_nocrop \
|
||||
--eval-mode binary \
|
||||
--tower-mode single
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
from PIL import Image, ImageDraw
|
||||
|
||||
|
||||
def load_disc_mask(contour_path: Path, orig_size: tuple[int, int],
|
||||
cam_h: int, cam_w: int) -> np.ndarray | None:
|
||||
"""
|
||||
Load a PAPILA disc contour TXT file, polygon-fill at original image
|
||||
dimensions, then resize to (cam_h, cam_w). Returns a bool array or
|
||||
None if the contour cannot be loaded.
|
||||
"""
|
||||
try:
|
||||
arr = np.loadtxt(str(contour_path), dtype=np.float32)
|
||||
except Exception:
|
||||
return None
|
||||
if arr.ndim == 1:
|
||||
arr = arr.reshape(-1, 2)
|
||||
if arr.shape[0] < 3 or arr.shape[1] < 2:
|
||||
return None
|
||||
|
||||
# orig_size is (W, H) as PIL convention
|
||||
img = Image.new("L", orig_size, 0)
|
||||
draw = ImageDraw.Draw(img)
|
||||
draw.polygon([tuple(pt) for pt in arr[:, :2]], fill=1)
|
||||
mask = np.array(img.resize((cam_w, cam_h), Image.NEAREST), dtype=bool)
|
||||
return mask
|
||||
|
||||
|
||||
def build_disc_lookup(manifest_path: Path) -> dict[tuple[int, str], tuple[Path, tuple[int, int]]]:
|
||||
"""
|
||||
Returns {(patient_id_int, eye): (disc_contour_path, (img_W, img_H))}.
|
||||
Only PAPILA rows are included.
|
||||
"""
|
||||
mf = pd.read_csv(manifest_path)
|
||||
lookup: dict[tuple[int, str], tuple[Path, tuple[int, int]]] = {}
|
||||
for _, row in mf.iterrows():
|
||||
sid = str(row["sample_id"])
|
||||
if not sid.startswith("papila_RET"):
|
||||
continue
|
||||
# sample_id: papila_RET002OD or papila_RET002OS
|
||||
suffix = sid[len("papila_RET"):] # e.g. "002OD"
|
||||
eye = suffix[-2:] # "OD" or "OS"
|
||||
pid = int(suffix[:-2]) # 2
|
||||
disc_path = Path(str(row["annotation_disc"]))
|
||||
img_path = Path(str(row["image_path"]))
|
||||
if not disc_path.exists():
|
||||
continue
|
||||
# read original image size once
|
||||
try:
|
||||
with Image.open(img_path) as im:
|
||||
orig_size = im.size # (W, H)
|
||||
except Exception:
|
||||
continue
|
||||
lookup[(pid, eye)] = (disc_path, orig_size)
|
||||
return lookup
|
||||
|
||||
|
||||
def attention_entropy(cam: np.ndarray) -> float:
|
||||
flat = cam.flatten().astype(np.float64)
|
||||
flat = flat / (flat.sum() + 1e-12)
|
||||
return float(-np.sum(flat * np.log(flat + 1e-12)))
|
||||
|
||||
|
||||
def load_fold(gradcam_dir: Path):
|
||||
idx_path = gradcam_dir / "gradcam_index.csv"
|
||||
if not idx_path.exists():
|
||||
return None
|
||||
idx = pd.read_csv(idx_path)
|
||||
records = []
|
||||
for _, row in idx.iterrows():
|
||||
pid = row["patient_id"]
|
||||
for eye in ("OD", "OS"):
|
||||
npy = gradcam_dir / f"patient_{pid}_{eye}_cam.npy"
|
||||
if not npy.exists():
|
||||
continue
|
||||
cam = np.load(npy)
|
||||
records.append({
|
||||
"patient_id": pid,
|
||||
"eye": eye,
|
||||
"true_label": int(row["true_label"]),
|
||||
"true_name": row["true_name"],
|
||||
"pred_label": int(row["pred_label"]),
|
||||
"pred_name": row["pred_name"],
|
||||
"confidence": float(row["confidence"]),
|
||||
"correct": bool(row["correct"]),
|
||||
"cam": cam,
|
||||
})
|
||||
return records
|
||||
|
||||
|
||||
def main():
|
||||
ap = argparse.ArgumentParser()
|
||||
ap.add_argument("--run-dir", default="analysis_data/pipeline_nocrop")
|
||||
ap.add_argument("--eval-mode", default="binary")
|
||||
ap.add_argument("--tower-mode", default="single")
|
||||
ap.add_argument("--manifest", default="manifest.csv")
|
||||
ap.add_argument("--out", default=None)
|
||||
args = ap.parse_args()
|
||||
|
||||
run_dir = Path(args.run_dir)
|
||||
mode_dir = run_dir / args.eval_mode / args.tower_mode
|
||||
out_dir = mode_dir / "gradcam_aggregate"
|
||||
if args.out:
|
||||
out_dir = Path(args.out)
|
||||
out_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# ---- build disc mask lookup ----
|
||||
manifest_path = Path(args.manifest)
|
||||
disc_lookup = build_disc_lookup(manifest_path)
|
||||
print(f"Disc mask lookup: {len(disc_lookup)} entries from {manifest_path}")
|
||||
|
||||
# ---- collect all records ----
|
||||
all_records = []
|
||||
stat_rows = []
|
||||
fold_dirs = sorted(
|
||||
[d for d in mode_dir.iterdir() if d.is_dir() and d.name.startswith("fold")],
|
||||
key=lambda p: int(p.name.replace("fold", "")),
|
||||
)
|
||||
if not fold_dirs:
|
||||
print(f"No fold dirs found under {mode_dir}")
|
||||
return
|
||||
|
||||
for fd in fold_dirs:
|
||||
gcam_dir = fd / "explainability" / "gradcam"
|
||||
records = load_fold(gcam_dir)
|
||||
if records is None:
|
||||
print(f" [skip] {fd.name}: no gradcam_index.csv")
|
||||
continue
|
||||
print(f" {fd.name}: {len(records)} eye records")
|
||||
for r in records:
|
||||
r["fold"] = fd.name
|
||||
all_records.append(r)
|
||||
|
||||
if not all_records:
|
||||
print("No records found — re-run explain_fold.py first.")
|
||||
return
|
||||
|
||||
print(f"\nTotal eye records: {len(all_records)}")
|
||||
|
||||
h, w = all_records[0]["cam"].shape
|
||||
|
||||
# ---- per-record stats ----
|
||||
n_missing = 0
|
||||
for r in all_records:
|
||||
cam = r["cam"]
|
||||
total = cam.sum() + 1e-12
|
||||
pid = int(r["patient_id"])
|
||||
eye = r["eye"]
|
||||
|
||||
disc_mask = None
|
||||
key = (pid, eye)
|
||||
if key in disc_lookup:
|
||||
disc_path, orig_size = disc_lookup[key]
|
||||
disc_mask = load_disc_mask(disc_path, orig_size, h, w)
|
||||
if disc_mask is None:
|
||||
n_missing += 1
|
||||
disc_frac = float("nan")
|
||||
else:
|
||||
disc_frac = float(cam[disc_mask].sum() / total)
|
||||
|
||||
stat_rows.append({
|
||||
"patient_id": r["patient_id"],
|
||||
"fold": r["fold"],
|
||||
"eye": r["eye"],
|
||||
"true_name": r["true_name"],
|
||||
"pred_name": r["pred_name"],
|
||||
"correct": r["correct"],
|
||||
"confidence": r["confidence"],
|
||||
"disc_frac": disc_frac,
|
||||
"entropy": attention_entropy(cam),
|
||||
})
|
||||
|
||||
if n_missing:
|
||||
print(f" Warning: {n_missing} records had no disc mask (disc_frac=NaN)")
|
||||
|
||||
stats_df = pd.DataFrame(stat_rows)
|
||||
stats_path = out_dir / "attention_stats.csv"
|
||||
stats_df.to_csv(stats_path, index=False)
|
||||
print(f"Saved attention stats → {stats_path}")
|
||||
|
||||
# ---- mean heatmaps ----
|
||||
npz_arrays = {}
|
||||
groups: dict[tuple, list[np.ndarray]] = {}
|
||||
for r in all_records:
|
||||
key = (r["eye"], r["true_name"], "correct" if r["correct"] else "incorrect")
|
||||
groups.setdefault(key, []).append(r["cam"])
|
||||
for r in all_records:
|
||||
key = (r["eye"], r["true_name"], "all")
|
||||
groups.setdefault(key, []).append(r["cam"])
|
||||
|
||||
for (eye, cls, split), cams in groups.items():
|
||||
stack = np.stack(cams, axis=0)
|
||||
key_base = f"{eye}_{cls}_{split}"
|
||||
npz_arrays[f"{key_base}_mean"] = stack.mean(axis=0).astype(np.float32)
|
||||
npz_arrays[f"{key_base}_std"] = stack.std(axis=0).astype(np.float32)
|
||||
npz_arrays[f"{key_base}_count"] = np.array(len(cams))
|
||||
print(f" {key_base}: N={len(cams)}")
|
||||
|
||||
npz_path = out_dir / "mean_heatmaps.npz"
|
||||
np.savez_compressed(npz_path, **npz_arrays)
|
||||
print(f"Saved mean heatmaps → {npz_path}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,933 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Post-hoc explainability for a single saved fold.
|
||||
|
||||
Phase 1 — MD permutation feature importance (bar chart + CSV).
|
||||
Phase 2 — GradCAM overlays on all holdout (or val) patients.
|
||||
|
||||
Usage:
|
||||
python scripts/output_analysis/explainability/explain_fold.py \
|
||||
--fold-dir analysis_data/.../binary/single/fold0 \
|
||||
[--checkpoint best_single.pt | best_holdout_single.pt] \
|
||||
[--split holdout] # falls back to val if no holdout
|
||||
[--image-dir Papila/FundusImages] \
|
||||
[--clinical-dir Papila/ClinicalData] \
|
||||
[--n-permutations 30] \
|
||||
[--seed 0] \
|
||||
[--alpha 0.45]
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
|
||||
import matplotlib
|
||||
|
||||
matplotlib.use("Agg")
|
||||
import matplotlib.cm as cm
|
||||
import matplotlib.patches as mpatches
|
||||
import matplotlib.pyplot as plt
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from PIL import Image
|
||||
from sklearn.metrics import roc_auc_score
|
||||
|
||||
REPO_ROOT = Path(__file__).resolve().parents[3]
|
||||
if str(REPO_ROOT) not in sys.path:
|
||||
sys.path.insert(0, str(REPO_ROOT))
|
||||
|
||||
from classes.v2.data_bundle import DataBundle
|
||||
from classes.v2.papila_builders import build_papila_data
|
||||
from classes.v2.profiles.papila import build_papila_profile
|
||||
from classes.v2.split_manager import PatientFirstSplitManager
|
||||
from classes.v2.loader_factory import filter_bilateral_samples, make_loader
|
||||
from classes.v2.metrics import _score_arrays
|
||||
from classes.v2.models import SingleEyeHT
|
||||
from classes.v2.transforms import build_eval_transform
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Label display helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
BINARY_LABELS = {0: "Normal", 1: "Glaucoma"}
|
||||
MULTICLASS_LABELS = {0: "Normal", 1: "Glaucoma", 2: "Suspect"}
|
||||
|
||||
|
||||
def label_name(label: int, eval_mode: str) -> str:
|
||||
mapping = BINARY_LABELS if eval_mode == "binary" else MULTICLASS_LABELS
|
||||
return mapping.get(int(label), str(label))
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# GradCAM
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class GradCAM:
|
||||
"""Minimal GradCAM using forward/backward hooks. No extra dependencies."""
|
||||
|
||||
def __init__(self, target_layer: torch.nn.Module) -> None:
|
||||
self._acts: torch.Tensor | None = None
|
||||
self._grads: torch.Tensor | None = None
|
||||
self._h1 = target_layer.register_forward_hook(self._save_acts)
|
||||
self._h2 = target_layer.register_full_backward_hook(self._save_grads)
|
||||
|
||||
def _save_acts(self, _m, _i, output):
|
||||
self._acts = output.detach()
|
||||
|
||||
def _save_grads(self, _m, _gi, grad_output):
|
||||
self._grads = grad_output[0].detach()
|
||||
|
||||
def compute(
|
||||
self,
|
||||
img: torch.Tensor,
|
||||
meta: torch.Tensor,
|
||||
model: torch.nn.Module,
|
||||
target_class: int | None = None,
|
||||
) -> tuple[np.ndarray, int]:
|
||||
"""Return (cam [H,W] in [0,1], predicted_class_index)."""
|
||||
model.eval()
|
||||
with torch.enable_grad():
|
||||
out = model(img, meta)
|
||||
pred = int(out.argmax(1).item())
|
||||
tc = pred if target_class is None else target_class
|
||||
model.zero_grad()
|
||||
out[0, tc].backward()
|
||||
|
||||
if self._acts is None or self._grads is None:
|
||||
raise RuntimeError("GradCAM hooks did not fire — check target_layer.")
|
||||
|
||||
weights = self._grads.mean(dim=(2, 3), keepdim=True) # [1,C,1,1]
|
||||
cam = F.relu((weights * self._acts).sum(dim=1, keepdim=True)) # [1,1,h,w]
|
||||
cam = F.interpolate(cam, img.shape[-2:], mode="bilinear", align_corners=False)
|
||||
cam_np = cam.squeeze().cpu().numpy()
|
||||
lo, hi = cam_np.min(), cam_np.max()
|
||||
cam_np = (cam_np - lo) / (hi - lo + 1e-8)
|
||||
return cam_np, pred
|
||||
|
||||
def remove(self) -> None:
|
||||
self._h1.remove()
|
||||
self._h2.remove()
|
||||
|
||||
|
||||
def get_gradcam_layer(model: SingleEyeHT, backbone: str) -> torch.nn.Module:
|
||||
"""Return the final spatial feature map layer for GradCAM."""
|
||||
bb = model.img_tower.backbone
|
||||
key = backbone.lower()
|
||||
if key in ("refugelike",) or "resnet" in key:
|
||||
return bb.layer4[-1]
|
||||
if "efficientnet" in key or "refuge_efficient" in key:
|
||||
return bb.features[-1]
|
||||
if "densenet" in key or key == "refuge_densenet":
|
||||
return bb.features.denseblock4
|
||||
if "mobilenet" in key:
|
||||
return bb.features[-1]
|
||||
if "vgg" in key:
|
||||
return bb.features[-1]
|
||||
raise ValueError(f"Unknown backbone for GradCAM target layer: {backbone!r}")
|
||||
|
||||
|
||||
def overlay_gradcam(
|
||||
original_pil: Image.Image, cam: np.ndarray, alpha: float = 0.45
|
||||
) -> Image.Image:
|
||||
"""Blend a jet-coloured GradCAM map onto the original image."""
|
||||
cam_u8 = (cam * 255).astype(np.uint8)
|
||||
cam_resized = (
|
||||
np.array(Image.fromarray(cam_u8).resize(original_pil.size, Image.BILINEAR))
|
||||
/ 255.0
|
||||
)
|
||||
colored = (cm.jet(cam_resized)[:, :, :3] * 255).astype(np.uint8)
|
||||
return Image.blend(original_pil.convert("RGB"), Image.fromarray(colored), alpha)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Feature index map
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def build_feature_index_map(data: DataBundle) -> dict[str, dict]:
|
||||
"""
|
||||
Return a mapping feature_name → {"value_dims": [...], "missing_dims": [...]}
|
||||
that covers every input dimension of the MD tower vector.
|
||||
|
||||
Layout (from DataBundle.vectorize_row):
|
||||
[scalar_0..scalar_n-1 | cat_onehot | scalar_missing_0..scalar_missing_n-1]
|
||||
"""
|
||||
n_scalar = len(data.scalar_cols)
|
||||
cat_expanded = sum(len(m) for m in data.cat_maps.values())
|
||||
|
||||
feature_map: dict[str, dict] = {}
|
||||
idx = 0
|
||||
|
||||
# Scalar features: value_dim + corresponding missing flag
|
||||
for i, col in enumerate(data.scalar_cols):
|
||||
missing_dim = n_scalar + cat_expanded + i
|
||||
feature_map[col] = {"value_dims": [i], "missing_dims": [missing_dim]}
|
||||
idx += 1
|
||||
|
||||
# Categorical features: permute the entire one-hot block
|
||||
cat_offset = n_scalar
|
||||
for col in data.cat_cols:
|
||||
n_cats = len(data.cat_maps[col])
|
||||
dims = list(range(cat_offset, cat_offset + n_cats))
|
||||
feature_map[col] = {"value_dims": dims, "missing_dims": []}
|
||||
cat_offset += n_cats
|
||||
|
||||
return feature_map
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Phase 1 — MD permutation importance
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def run_permutation_importance(
|
||||
model: SingleEyeHT,
|
||||
loader,
|
||||
data: DataBundle,
|
||||
num_classes: int,
|
||||
device: torch.device,
|
||||
n_permutations: int,
|
||||
seed: int,
|
||||
out_dir: Path,
|
||||
) -> None:
|
||||
print("\n[Phase 1] MD permutation importance ...", flush=True)
|
||||
|
||||
# ---- cache bilateral image embeddings + metadata tensors + labels ----
|
||||
img1_feats_list, img2_feats_list = [], []
|
||||
md1_list, md2_list, label_list = [], [], []
|
||||
model.eval()
|
||||
with torch.no_grad():
|
||||
for batch in loader:
|
||||
img1 = batch["image_1"].to(device)
|
||||
img2 = batch["image_2"].to(device)
|
||||
md1 = batch["matrix_1"].to(device)
|
||||
md2 = batch["matrix_2"].to(device)
|
||||
labels = batch["label_1"]
|
||||
img1_feats_list.append(model.img_tower(img1))
|
||||
img2_feats_list.append(model.img_tower(img2))
|
||||
md1_list.append(md1)
|
||||
md2_list.append(md2)
|
||||
if isinstance(labels, torch.Tensor):
|
||||
label_list.append(labels)
|
||||
else:
|
||||
label_list.append(torch.tensor(labels, dtype=torch.long))
|
||||
|
||||
img1_feats = torch.cat(img1_feats_list) # [N, img_dim]
|
||||
img2_feats = torch.cat(img2_feats_list) # [N, img_dim]
|
||||
md1_tensor = torch.cat(md1_list) # [N, feature_dim]
|
||||
md2_tensor = torch.cat(md2_list) # [N, feature_dim]
|
||||
y_true = torch.cat(label_list).numpy()
|
||||
N = len(y_true)
|
||||
|
||||
if N == 0:
|
||||
print(" [Phase 1] No samples — skipping.", flush=True)
|
||||
return
|
||||
|
||||
# ---- baseline AUC (patient-level: average OD/OS fused probabilities) ----
|
||||
with torch.no_grad():
|
||||
md1_feats = model.md_tower(md1_tensor)
|
||||
md2_feats = model.md_tower(md2_tensor)
|
||||
fused1, _, _ = model.bridge(img1_feats, md1_feats)
|
||||
fused2, _, _ = model.bridge(img2_feats, md2_feats)
|
||||
probs_baseline = (
|
||||
0.5 * (torch.softmax(fused1, dim=1) + torch.softmax(fused2, dim=1))
|
||||
).cpu().numpy()
|
||||
_, baseline_auc, _ = _score_arrays(y_true, probs_baseline, num_classes)
|
||||
print(f" Baseline AUC: {baseline_auc:.4f} (N={N})", flush=True)
|
||||
|
||||
# ---- feature index map ----
|
||||
feat_map = build_feature_index_map(data)
|
||||
rng = np.random.default_rng(seed)
|
||||
|
||||
results = []
|
||||
for feat_name, dims in feat_map.items():
|
||||
all_dims = dims["value_dims"] + dims["missing_dims"]
|
||||
drops = []
|
||||
for _ in range(n_permutations):
|
||||
perm1 = md1_tensor.clone()
|
||||
perm2 = md2_tensor.clone()
|
||||
perm_idx = torch.from_numpy(rng.permutation(N)).to(device)
|
||||
# Apply the same donor patient permutation to both eyes to preserve
|
||||
# within-patient coherence while breaking feature-label association.
|
||||
perm1[:, all_dims] = perm1[perm_idx][:, all_dims]
|
||||
perm2[:, all_dims] = perm2[perm_idx][:, all_dims]
|
||||
with torch.no_grad():
|
||||
md1_p = model.md_tower(perm1)
|
||||
md2_p = model.md_tower(perm2)
|
||||
fused1_p, _, _ = model.bridge(img1_feats, md1_p)
|
||||
fused2_p, _, _ = model.bridge(img2_feats, md2_p)
|
||||
probs_p = (
|
||||
0.5
|
||||
* (
|
||||
torch.softmax(fused1_p, dim=1)
|
||||
+ torch.softmax(fused2_p, dim=1)
|
||||
)
|
||||
).cpu().numpy()
|
||||
_, auc_p, _ = _score_arrays(y_true, probs_p, num_classes)
|
||||
drops.append(baseline_auc - auc_p)
|
||||
|
||||
mean_drop = float(np.mean(drops))
|
||||
std_drop = float(np.std(drops))
|
||||
results.append({"feature": feat_name, "importance": mean_drop, "std": std_drop})
|
||||
print(
|
||||
f" {feat_name:30s} Δ AUC = {mean_drop:+.4f} ± {std_drop:.4f}", flush=True
|
||||
)
|
||||
|
||||
results.sort(key=lambda r: r["importance"], reverse=True)
|
||||
|
||||
# ---- total MD ablation (all features permuted simultaneously) ----
|
||||
print(" Running total MD ablation ...", flush=True)
|
||||
total_drops = []
|
||||
for _ in range(n_permutations):
|
||||
perm_idx = torch.from_numpy(rng.permutation(N)).to(device)
|
||||
perm1_all = md1_tensor[perm_idx]
|
||||
perm2_all = md2_tensor[perm_idx]
|
||||
with torch.no_grad():
|
||||
md1_all = model.md_tower(perm1_all)
|
||||
md2_all = model.md_tower(perm2_all)
|
||||
f1, _, _ = model.bridge(img1_feats, md1_all)
|
||||
f2, _, _ = model.bridge(img2_feats, md2_all)
|
||||
probs_all = (
|
||||
0.5 * (torch.softmax(f1, dim=1) + torch.softmax(f2, dim=1))
|
||||
).cpu().numpy()
|
||||
_, auc_all, _ = _score_arrays(y_true, probs_all, num_classes)
|
||||
total_drops.append(baseline_auc - auc_all)
|
||||
total_mean = float(np.mean(total_drops))
|
||||
total_std = float(np.std(total_drops))
|
||||
print(
|
||||
f" Total MD ablation Δ AUC = {total_mean:+.4f} ± {total_std:.4f}", flush=True
|
||||
)
|
||||
|
||||
# ---- Gaussian noise ablation (tests architectural vs informational benefit) ----
|
||||
print(" Running Gaussian noise ablation ...", flush=True)
|
||||
noise_drops = []
|
||||
for _ in range(n_permutations):
|
||||
noise1 = torch.randn_like(md1_tensor)
|
||||
noise2 = torch.randn_like(md2_tensor)
|
||||
with torch.no_grad():
|
||||
md1_noise = model.md_tower(noise1)
|
||||
md2_noise = model.md_tower(noise2)
|
||||
f1, _, _ = model.bridge(img1_feats, md1_noise)
|
||||
f2, _, _ = model.bridge(img2_feats, md2_noise)
|
||||
probs_noise = (
|
||||
0.5 * (torch.softmax(f1, dim=1) + torch.softmax(f2, dim=1))
|
||||
).cpu().numpy()
|
||||
_, auc_noise, _ = _score_arrays(y_true, probs_noise, num_classes)
|
||||
noise_drops.append(baseline_auc - auc_noise)
|
||||
noise_mean = float(np.mean(noise_drops))
|
||||
noise_std = float(np.std(noise_drops))
|
||||
print(
|
||||
f" Gaussian noise ablation Δ AUC = {noise_mean:+.4f} ± {noise_std:.4f}", flush=True
|
||||
)
|
||||
print(
|
||||
f" [interpretation] permutation Δ={total_mean:+.4f} noise Δ={noise_mean:+.4f} "
|
||||
f"informational gain = {total_mean - noise_mean:+.4f}",
|
||||
flush=True,
|
||||
)
|
||||
|
||||
# ---- save CSV ----
|
||||
import csv
|
||||
|
||||
csv_path = out_dir / "md_permutation_importance.csv"
|
||||
with csv_path.open("w", newline="") as f:
|
||||
writer = csv.DictWriter(f, fieldnames=["feature", "importance", "std"])
|
||||
writer.writeheader()
|
||||
writer.writerows(results)
|
||||
writer.writerow({"feature": "TOTAL_MD_ABLATION", "importance": total_mean, "std": total_std})
|
||||
writer.writerow({"feature": "GAUSSIAN_NOISE_ABLATION", "importance": noise_mean, "std": noise_std})
|
||||
|
||||
# ---- bar chart ----
|
||||
names = [r["feature"] for r in results]
|
||||
imps = [r["importance"] for r in results]
|
||||
stds = [r["std"] for r in results]
|
||||
colors = ["#e05c5c" if v >= 0 else "#5c9ee0" for v in imps]
|
||||
|
||||
fig, ax = plt.subplots(figsize=(9, max(4, (len(names) + 3) * 0.45)))
|
||||
y_pos = np.arange(len(names))
|
||||
ax.barh(
|
||||
y_pos, imps, xerr=stds, color=colors, ecolor="grey", capsize=3, height=0.6
|
||||
)
|
||||
ax.axhline(len(names) - 0.25, color="grey", linewidth=0.6, linestyle="--")
|
||||
# total ablation
|
||||
ax.barh(
|
||||
len(names) + 0.5, total_mean, xerr=total_std,
|
||||
color="#c45ce0" if total_mean >= 0 else "#5c9ee0",
|
||||
ecolor="grey", capsize=3, height=0.6,
|
||||
)
|
||||
# gaussian noise ablation
|
||||
ax.barh(
|
||||
len(names) + 1.5, noise_mean, xerr=noise_std,
|
||||
color="#e08c2a" if noise_mean >= 0 else "#5c9ee0",
|
||||
ecolor="grey", capsize=3, height=0.6,
|
||||
)
|
||||
ax.set_yticks(list(y_pos) + [len(names) + 0.5, len(names) + 1.5])
|
||||
ax.set_yticklabels(names + ["ALL MD (permute)", "ALL MD (noise)"], fontsize=9)
|
||||
ax.invert_yaxis()
|
||||
ax.axvline(0, color="black", linewidth=0.8)
|
||||
ax.set_xlabel("Mean AUC drop (baseline − permuted)", fontsize=10)
|
||||
ax.set_title(
|
||||
f"MD Tower — Permutation Feature Importance\n"
|
||||
f"baseline AUC={baseline_auc:.4f} N={N} repeats={n_permutations}",
|
||||
fontsize=11,
|
||||
)
|
||||
fig.tight_layout()
|
||||
fig.savefig(out_dir / "md_permutation_importance.png", dpi=150)
|
||||
plt.close(fig)
|
||||
print(f" Saved → {out_dir / 'md_permutation_importance.png'}", flush=True)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Phase 2 — GradCAM overlays
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def run_gradcam(
|
||||
model: SingleEyeHT,
|
||||
loader,
|
||||
data: DataBundle,
|
||||
eval_df,
|
||||
eval_mode: str,
|
||||
backbone: str,
|
||||
device: torch.device,
|
||||
alpha: float,
|
||||
out_dir: Path,
|
||||
) -> None:
|
||||
print("\n[Phase 2] GradCAM overlays ...", flush=True)
|
||||
gradcam_dir = out_dir / "gradcam"
|
||||
gradcam_dir.mkdir(exist_ok=True)
|
||||
|
||||
target_layer = get_gradcam_layer(model, backbone)
|
||||
gcam = GradCAM(target_layer)
|
||||
|
||||
num_classes = model.bridge.classifier_fused[-1].out_features
|
||||
|
||||
overlay_grid_items: list[
|
||||
tuple[Image.Image | None, Image.Image | None, str, bool]
|
||||
] = []
|
||||
index_rows: list[dict] = []
|
||||
|
||||
model.eval()
|
||||
for batch in loader:
|
||||
img_od = batch["image_1"].to(device) # [1, 3, H, W]
|
||||
img_os = batch["image_2"].to(device) # [1, 3, H, W]
|
||||
meta_od = batch["matrix_1"].to(device) # [1, feature_dim]
|
||||
meta_os = batch["matrix_2"].to(device)
|
||||
lbl_raw = batch["label_1"][0]
|
||||
label = int(lbl_raw.item() if isinstance(lbl_raw, torch.Tensor) else lbl_raw)
|
||||
pid = batch["id_1"][0]
|
||||
|
||||
# GradCAM for each eye (OD drives the prediction label)
|
||||
cam_od, pred = gcam.compute(img_od, meta_od, model)
|
||||
cam_os, _ = gcam.compute(img_os, meta_os, model)
|
||||
|
||||
# Confidence of predicted class
|
||||
with torch.no_grad():
|
||||
out_od = model(img_od, meta_od)
|
||||
conf = float(torch.softmax(out_od, dim=1)[0, pred].item())
|
||||
|
||||
# Load original (un-normalised) images from disk
|
||||
row_od = eval_df[
|
||||
(eval_df["Patient ID"] == int(pid)) & (eval_df["eyeID"] == "OD")
|
||||
]
|
||||
row_os = eval_df[
|
||||
(eval_df["Patient ID"] == int(pid)) & (eval_df["eyeID"] == "OS")
|
||||
]
|
||||
orig_od = (
|
||||
Image.open(data.get_image_path(row_od.iloc[0])).convert("RGB")
|
||||
if len(row_od)
|
||||
else None
|
||||
)
|
||||
orig_os = (
|
||||
Image.open(data.get_image_path(row_os.iloc[0])).convert("RGB")
|
||||
if len(row_os)
|
||||
else None
|
||||
)
|
||||
|
||||
true_name = label_name(label, eval_mode)
|
||||
pred_name = label_name(pred, eval_mode)
|
||||
correct = label == pred
|
||||
title = (
|
||||
f"Patient {pid} | True: {true_name} | Pred: {pred_name} "
|
||||
f"| conf={conf:.2f} {'✓' if correct else '✗'}"
|
||||
)
|
||||
|
||||
# ---- per-patient 2×2 figure (OD raw | OD overlay / OS raw | OS overlay) ----
|
||||
fig, axes = plt.subplots(2, 2, figsize=(10, 9))
|
||||
fig.suptitle(
|
||||
title, fontsize=11, fontweight="bold", color="green" if correct else "red"
|
||||
)
|
||||
|
||||
# Row 0: OD
|
||||
if orig_od is not None:
|
||||
axes[0, 0].imshow(orig_od)
|
||||
axes[0, 0].set_title("OD — original", fontsize=9)
|
||||
axes[0, 1].imshow(overlay_gradcam(orig_od, cam_od, alpha))
|
||||
axes[0, 1].set_title("OD — GradCAM", fontsize=9)
|
||||
else:
|
||||
axes[0, 0].set_title("OD — (missing)", fontsize=9)
|
||||
axes[0, 0].axis("off")
|
||||
axes[0, 1].axis("off")
|
||||
|
||||
# Row 1: OS
|
||||
if orig_os is not None:
|
||||
axes[1, 0].imshow(orig_os)
|
||||
axes[1, 0].set_title("OS — original", fontsize=9)
|
||||
axes[1, 1].imshow(overlay_gradcam(orig_os, cam_os, alpha))
|
||||
axes[1, 1].set_title("OS — GradCAM", fontsize=9)
|
||||
else:
|
||||
axes[1, 0].set_title("OS — (missing)", fontsize=9)
|
||||
axes[1, 0].axis("off")
|
||||
axes[1, 1].axis("off")
|
||||
|
||||
fig.tight_layout()
|
||||
out_path = gradcam_dir / f"patient_{pid}_OD_OS.png"
|
||||
fig.savefig(out_path, dpi=120)
|
||||
plt.close(fig)
|
||||
print(
|
||||
f" Patient {pid}: {true_name} → {pred_name} ({conf:.2f}) → {out_path.name}",
|
||||
flush=True,
|
||||
)
|
||||
|
||||
# Save raw CAM arrays
|
||||
np.save(gradcam_dir / f"patient_{pid}_OD_cam.npy", cam_od)
|
||||
np.save(gradcam_dir / f"patient_{pid}_OS_cam.npy", cam_os)
|
||||
index_rows.append({
|
||||
"patient_id": pid,
|
||||
"true_label": label,
|
||||
"true_name": true_name,
|
||||
"pred_label": pred,
|
||||
"pred_name": pred_name,
|
||||
"confidence": conf,
|
||||
"correct": correct,
|
||||
})
|
||||
|
||||
# Accumulate for summary grid
|
||||
od_overlay = overlay_gradcam(orig_od, cam_od, alpha) if orig_od else None
|
||||
os_overlay = overlay_gradcam(orig_os, cam_os, alpha) if orig_os else None
|
||||
short_lbl = f"P{pid} {true_name[:3]}→{pred_name[:3]} {'✓' if correct else '✗'}"
|
||||
overlay_grid_items.append((od_overlay, os_overlay, short_lbl, correct))
|
||||
|
||||
gcam.remove()
|
||||
|
||||
# ---- save index CSV ----
|
||||
if index_rows:
|
||||
import csv
|
||||
idx_path = gradcam_dir / "gradcam_index.csv"
|
||||
with idx_path.open("w", newline="") as f:
|
||||
writer = csv.DictWriter(f, fieldnames=list(index_rows[0].keys()))
|
||||
writer.writeheader()
|
||||
writer.writerows(index_rows)
|
||||
print(f" Index CSV → {idx_path}", flush=True)
|
||||
|
||||
# ---- summary grid: N_patients rows × 2 cols (OD overlay | OS overlay) ----
|
||||
n = len(overlay_grid_items)
|
||||
if n == 0:
|
||||
print(" [Phase 2] No patients to visualise.", flush=True)
|
||||
return
|
||||
|
||||
fig, axes = plt.subplots(n, 2, figsize=(8, n * 3.2 + 0.8))
|
||||
if n == 1:
|
||||
axes = axes[np.newaxis, :]
|
||||
fig.suptitle("GradCAM Summary Grid — all holdout patients", fontsize=12)
|
||||
|
||||
for i, (od_ov, os_ov, lbl, correct) in enumerate(overlay_grid_items):
|
||||
color = "green" if correct else "red"
|
||||
for j in range(2):
|
||||
axes[i, j].axis("off")
|
||||
if od_ov is not None:
|
||||
axes[i, 0].imshow(od_ov)
|
||||
axes[i, 0].set_title(f"{lbl}\nOD", fontsize=7, color=color)
|
||||
if os_ov is not None:
|
||||
axes[i, 1].imshow(os_ov)
|
||||
axes[i, 1].set_title(f"{lbl}\nOS", fontsize=7, color=color)
|
||||
|
||||
fig.tight_layout()
|
||||
grid_path = out_dir / "gradcam_summary_grid.png"
|
||||
fig.savefig(grid_path, dpi=120)
|
||||
plt.close(fig)
|
||||
print(f" Summary grid → {grid_path}", flush=True)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Main
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def parse_args():
|
||||
ap = argparse.ArgumentParser(
|
||||
description="Post-hoc explainability for a saved fold."
|
||||
)
|
||||
ap.add_argument(
|
||||
"--fold-dir",
|
||||
type=Path,
|
||||
required=True,
|
||||
help="Path to fold directory, e.g. analysis_data/.../binary/single/fold0",
|
||||
)
|
||||
ap.add_argument(
|
||||
"--checkpoint",
|
||||
default="best_single.pt",
|
||||
help="Checkpoint filename inside fold_dir (default: best_single.pt; "
|
||||
"use best_holdout_single.pt for holdout-selected model)",
|
||||
)
|
||||
ap.add_argument(
|
||||
"--split",
|
||||
choices=["holdout", "val"],
|
||||
default="holdout",
|
||||
help="Which patient set to analyse (default: holdout, falls back to val)",
|
||||
)
|
||||
ap.add_argument("--image-dir", default="Papila/FundusImages")
|
||||
ap.add_argument("--clinical-dir", default="Papila/ClinicalData")
|
||||
ap.add_argument("--label-col", default="Diagnosis")
|
||||
ap.add_argument("--cat-cols", nargs="*", default=["Gender", "Phakic/Pseudophakic"])
|
||||
ap.add_argument("--fold-seed", type=int, default=42)
|
||||
ap.add_argument("--holdout-seed", type=int, default=123)
|
||||
ap.add_argument("--holdout-per-class", type=int, default=5)
|
||||
ap.add_argument("--n-splits", type=int, default=5)
|
||||
ap.add_argument(
|
||||
"--n-permutations",
|
||||
type=int,
|
||||
default=30,
|
||||
help="Repetitions per feature for permutation importance (default: 30)",
|
||||
)
|
||||
ap.add_argument("--seed", type=int, default=0)
|
||||
ap.add_argument(
|
||||
"--alpha",
|
||||
type=float,
|
||||
default=0.45,
|
||||
help="GradCAM overlay opacity (default: 0.45)",
|
||||
)
|
||||
ap.add_argument("--batch-size", type=int, default=1)
|
||||
ap.add_argument("--no-phase1", action="store_true", help="Skip MD importance")
|
||||
ap.add_argument("--no-phase2", action="store_true", help="Skip GradCAM")
|
||||
ap.add_argument("--no-phase3", action="store_true", help="Skip fusion event analysis")
|
||||
return ap.parse_args()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Phase 3 — Fusion event analysis
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def run_fusion_event_analysis(
|
||||
model: SingleEyeHT,
|
||||
loader,
|
||||
device: torch.device,
|
||||
out_dir: Path,
|
||||
) -> None:
|
||||
print("\n[Phase 3] Fusion event analysis ...", flush=True)
|
||||
|
||||
from classes.v2.models import collect_probs_single_components
|
||||
|
||||
y_true, pf, pi, pm = collect_probs_single_components(
|
||||
model, loader, device, aggregate_patient=True
|
||||
)
|
||||
N = len(y_true)
|
||||
if N == 0:
|
||||
print(" [Phase 3] No samples — skipping.", flush=True)
|
||||
return
|
||||
|
||||
pred_f = pf.argmax(axis=1)
|
||||
pred_i = pi.argmax(axis=1)
|
||||
pred_m = pm.argmax(axis=1)
|
||||
|
||||
# confidence of the predicted class for each head
|
||||
conf_f = np.take_along_axis(pf, pred_f[:, None], axis=1).squeeze(1)
|
||||
conf_i = np.take_along_axis(pi, pred_i[:, None], axis=1).squeeze(1)
|
||||
conf_m = np.take_along_axis(pm, pred_m[:, None], axis=1).squeeze(1)
|
||||
# how much did fusion shift confidence vs the average of the two towers?
|
||||
conf_delta = conf_f - 0.5 * (conf_i + conf_m)
|
||||
|
||||
f_ok = pred_f == y_true
|
||||
i_ok = pred_i == y_true
|
||||
m_ok = pred_m == y_true
|
||||
|
||||
# 6 non-trivial bridge-effect event types
|
||||
full_correction = f_ok & ~i_ok & ~m_ok # both towers wrong → fused right
|
||||
img_assist = f_ok & ~i_ok & m_ok # img wrong, md right → fused right (md carried it)
|
||||
md_assist = f_ok & i_ok & ~m_ok # md wrong, img right → fused right (img carried it)
|
||||
full_error = ~f_ok & i_ok & m_ok # both towers right → fused wrong
|
||||
img_drag = ~f_ok & ~i_ok & m_ok # img wrong, md right → fused wrong (img dragged it down)
|
||||
md_drag = ~f_ok & i_ok & ~m_ok # md wrong, img right → fused wrong (md dragged it down)
|
||||
concordant_ok = f_ok & i_ok & m_ok
|
||||
concordant_bad = ~f_ok & ~i_ok & ~m_ok
|
||||
|
||||
event_labels = [
|
||||
"full correction\n(both wrong→fused right)",
|
||||
"img assist\n(img wrong, md right→right)",
|
||||
"md assist\n(md wrong, img right→right)",
|
||||
"full error\n(both right→fused wrong)",
|
||||
"img drag\n(img wrong, md right→wrong)",
|
||||
"md drag\n(md wrong, img right→wrong)",
|
||||
]
|
||||
event_masks = [full_correction, img_assist, md_assist, full_error, img_drag, md_drag]
|
||||
event_colors = ["#2ca02c", "#98df8a", "#b5d46e", "#d62728", "#ff9896", "#ffbf9b"]
|
||||
event_keys = ["full_correction", "img_assist", "md_assist",
|
||||
"full_error", "img_drag", "md_drag"]
|
||||
counts = [int(m.sum()) for m in event_masks]
|
||||
|
||||
print(f" N={N}", flush=True)
|
||||
for label, count in zip(event_labels, counts):
|
||||
print(f" {label.replace(chr(10), ' '):55s}: {count}", flush=True)
|
||||
n_corr, n_err = counts[0], counts[3]
|
||||
ratio_str = f"{n_corr}/{n_err}" if n_err > 0 else f"{n_corr}/0"
|
||||
print(f" full correction/error ratio: {ratio_str}", flush=True)
|
||||
print(f" conf_delta mean={conf_delta.mean():+.4f} median={np.median(conf_delta):+.4f}",
|
||||
flush=True)
|
||||
|
||||
# ---- CSV ----
|
||||
import csv
|
||||
event_type = np.where(concordant_ok, "concordant_correct",
|
||||
np.where(concordant_bad, "concordant_wrong", "other")).astype(object)
|
||||
for mask, key in zip(event_masks, event_keys):
|
||||
event_type[mask] = key
|
||||
|
||||
rows = []
|
||||
for idx in range(N):
|
||||
rows.append({
|
||||
"patient_idx": idx,
|
||||
"y_true": int(y_true[idx]),
|
||||
"pred_fused": int(pred_f[idx]),
|
||||
"pred_img": int(pred_i[idx]),
|
||||
"pred_md": int(pred_m[idx]),
|
||||
"conf_fused": float(conf_f[idx]),
|
||||
"conf_img": float(conf_i[idx]),
|
||||
"conf_md": float(conf_m[idx]),
|
||||
"conf_delta": float(conf_delta[idx]),
|
||||
"event_type": event_type[idx],
|
||||
})
|
||||
csv_path = out_dir / "fusion_events.csv"
|
||||
with csv_path.open("w", newline="") as f:
|
||||
writer = csv.DictWriter(f, fieldnames=list(rows[0].keys()))
|
||||
writer.writeheader()
|
||||
writer.writerows(rows)
|
||||
print(f" Saved → {csv_path}", flush=True)
|
||||
|
||||
# ---- plot ----
|
||||
fig, axes = plt.subplots(1, 3, figsize=(15, 4))
|
||||
|
||||
# Panel 1: stacked bar — positive events vs negative events
|
||||
pos_counts = counts[:3]
|
||||
neg_counts = counts[3:]
|
||||
pos_colors = event_colors[:3]
|
||||
neg_colors = event_colors[3:]
|
||||
for bar_x, bar_counts, bar_colors in ((0, pos_counts, pos_colors),
|
||||
(1, neg_counts, neg_colors)):
|
||||
bot = 0
|
||||
for c, col in zip(bar_counts, bar_colors):
|
||||
axes[0].bar(bar_x, c, bottom=bot, color=col, width=0.5)
|
||||
if c > 0:
|
||||
axes[0].text(bar_x, bot + c / 2, str(c), ha="center", va="center",
|
||||
fontsize=9, fontweight="bold")
|
||||
bot += c
|
||||
axes[0].set_xticks([0, 1])
|
||||
axes[0].set_xticklabels(["Positive\nevents", "Negative\nevents"])
|
||||
axes[0].set_ylabel("Count")
|
||||
axes[0].set_title(f"Fusion Events (N={N})")
|
||||
patches = [mpatches.Patch(color=c, label=l.replace("\n", " "))
|
||||
for c, l in zip(event_colors, event_labels)]
|
||||
axes[0].legend(handles=patches, fontsize=6, loc="upper right")
|
||||
|
||||
# Panel 2: conf_delta boxplot per event type (only non-empty)
|
||||
box_data = [conf_delta[m] for m in event_masks if m.sum() > 0]
|
||||
box_labels = [l.split("\n")[0] for m, l in zip(event_masks, event_labels) if m.sum() > 0]
|
||||
box_cols = [c for m, c in zip(event_masks, event_colors) if m.sum() > 0]
|
||||
if box_data:
|
||||
bp = axes[1].boxplot(box_data, patch_artist=True, widths=0.5)
|
||||
for patch, color in zip(bp["boxes"], box_cols):
|
||||
patch.set_facecolor(color)
|
||||
axes[1].set_xticks(range(1, len(box_labels) + 1))
|
||||
axes[1].set_xticklabels(box_labels, rotation=35, ha="right", fontsize=7)
|
||||
axes[1].axhline(0, color="black", linewidth=0.8, linestyle="--")
|
||||
axes[1].set_ylabel("conf_delta\n(fused − avg(img, md))")
|
||||
axes[1].set_title("Confidence delta by event type")
|
||||
|
||||
# Panel 3: img vs md confidence space, coloured by event type
|
||||
for mask, color, label in zip(event_masks, event_colors, event_labels):
|
||||
if mask.sum() > 0:
|
||||
axes[2].scatter(conf_i[mask], conf_m[mask], c=color,
|
||||
label=label.split("\n")[0], alpha=0.85, s=45, edgecolors="none")
|
||||
if concordant_ok.sum() > 0:
|
||||
axes[2].scatter(conf_i[concordant_ok], conf_m[concordant_ok],
|
||||
c="lightgrey", alpha=0.4, s=20, edgecolors="none", label="concordant correct")
|
||||
if concordant_bad.sum() > 0:
|
||||
axes[2].scatter(conf_i[concordant_bad], conf_m[concordant_bad],
|
||||
c="darkgrey", alpha=0.4, s=20, edgecolors="none", label="concordant wrong")
|
||||
axes[2].plot([0, 1], [0, 1], "k--", linewidth=0.5, alpha=0.4)
|
||||
axes[2].set_xlabel("conf_img")
|
||||
axes[2].set_ylabel("conf_md")
|
||||
axes[2].set_title("Tower confidence space\ncoloured by fusion event")
|
||||
axes[2].legend(fontsize=6, loc="lower right")
|
||||
|
||||
fig.tight_layout()
|
||||
fig.savefig(out_dir / "fusion_events.png", dpi=150)
|
||||
plt.close(fig)
|
||||
print(f" Saved → {out_dir / 'fusion_events.png'}", flush=True)
|
||||
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
fold_dir = args.fold_dir.resolve()
|
||||
if not fold_dir.is_dir():
|
||||
sys.exit(f"[ERROR] fold_dir does not exist: {fold_dir}")
|
||||
|
||||
ckpt_path = fold_dir / args.checkpoint
|
||||
if not ckpt_path.exists():
|
||||
sys.exit(
|
||||
f"[ERROR] Checkpoint not found: {ckpt_path}\n"
|
||||
f" Run training with --save-checkpoints (now the default) to produce checkpoints."
|
||||
)
|
||||
|
||||
# ---- read config from summary.json in parent (tower-mode) dir ----
|
||||
summary_path = fold_dir.parent / "summary.json"
|
||||
if not summary_path.exists():
|
||||
sys.exit(f"[ERROR] summary.json not found: {summary_path}")
|
||||
summary = json.loads(summary_path.read_text())
|
||||
backbone = summary["backbone"]
|
||||
eval_mode = summary["eval_mode"]
|
||||
tower_mode = summary.get("tower_mode", "single")
|
||||
fold_idx = int(fold_dir.name.replace("fold", ""))
|
||||
print(
|
||||
f"[explain_fold] fold={fold_idx} backbone={backbone} eval_mode={eval_mode} tower_mode={tower_mode}"
|
||||
)
|
||||
|
||||
if tower_mode not in ("single", "ensemble"):
|
||||
sys.exit(
|
||||
f"[ERROR] explain_fold currently supports single/ensemble tower modes, got: {tower_mode!r}"
|
||||
)
|
||||
|
||||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
print(f"[explain_fold] device={device} checkpoint={args.checkpoint}")
|
||||
|
||||
# ---- build DataBundle ----
|
||||
print("[explain_fold] Loading clinical data ...", flush=True)
|
||||
data = build_papila_data(
|
||||
image_dir=args.image_dir,
|
||||
clinical_dir=args.clinical_dir,
|
||||
label_col=args.label_col,
|
||||
cat_cols=args.cat_cols,
|
||||
n_splits=args.n_splits,
|
||||
random_seed=args.fold_seed,
|
||||
)
|
||||
df_mode = data.df.copy()
|
||||
if eval_mode == "binary":
|
||||
df_mode = df_mode[df_mode[args.label_col].isin([0, 1])].reset_index(drop=True)
|
||||
num_classes = 2 if eval_mode == "binary" else int(df_mode[args.label_col].nunique())
|
||||
|
||||
# ---- reconstruct the exact same split ----
|
||||
print("[explain_fold] Reconstructing split ...", flush=True)
|
||||
splitter = PatientFirstSplitManager(
|
||||
patient_col="Patient ID", label_col=args.label_col
|
||||
)
|
||||
split_args = SimpleNamespace(
|
||||
eval_mode=eval_mode,
|
||||
holdout_per_class=args.holdout_per_class,
|
||||
holdout_seed=args.holdout_seed,
|
||||
n_splits=args.n_splits,
|
||||
fold_seed=args.fold_seed,
|
||||
)
|
||||
clinical_ns = SimpleNamespace(df=df_mode, label_col=args.label_col)
|
||||
plans = splitter.build_plans(clinical=clinical_ns, args=split_args, profile=None)
|
||||
if fold_idx >= len(plans):
|
||||
sys.exit(f"[ERROR] fold_idx={fold_idx} but only {len(plans)} plans built.")
|
||||
split = plans[fold_idx]
|
||||
|
||||
if (
|
||||
args.split == "holdout"
|
||||
and split.holdout is not None
|
||||
and not split.holdout.empty
|
||||
):
|
||||
eval_df = split.holdout
|
||||
split_name = "holdout"
|
||||
else:
|
||||
if args.split == "holdout":
|
||||
print(" [WARN] No holdout set available; falling back to val.", flush=True)
|
||||
eval_df = split.val
|
||||
split_name = "val"
|
||||
print(
|
||||
f" Using {split_name} set: {eval_df['Patient ID'].nunique()} patients",
|
||||
flush=True,
|
||||
)
|
||||
|
||||
# ---- build loader ----
|
||||
profile_patient = build_papila_profile(
|
||||
patient_col="Patient ID", label_col=args.label_col, sample_mode="patient"
|
||||
)
|
||||
samples = filter_bilateral_samples(
|
||||
profile_patient.build_samples(df=eval_df, clinical=data)
|
||||
)
|
||||
if not samples:
|
||||
sys.exit("[ERROR] No bilateral samples found in the eval set.")
|
||||
loader = make_loader(
|
||||
samples,
|
||||
profile_patient.slot_descriptors(),
|
||||
image_transform=build_eval_transform(backbone),
|
||||
image_preprocessor=None,
|
||||
batch_size=args.batch_size,
|
||||
shuffle=False,
|
||||
num_workers=0,
|
||||
)
|
||||
|
||||
# ---- load model ----
|
||||
print(f"[explain_fold] Loading model from {ckpt_path} ...", flush=True)
|
||||
model = SingleEyeHT(
|
||||
backbone=backbone,
|
||||
freeze_ratio=0.0,
|
||||
augment=False,
|
||||
clinical_data=data,
|
||||
num_classes=num_classes,
|
||||
).to(device)
|
||||
state = torch.load(ckpt_path, map_location=device)
|
||||
model.load_state_dict(state)
|
||||
model.eval()
|
||||
|
||||
# ---- output directory ----
|
||||
out_dir = fold_dir / "explainability"
|
||||
out_dir.mkdir(exist_ok=True)
|
||||
print(f"[explain_fold] Output → {out_dir}", flush=True)
|
||||
|
||||
# ---- Phase 1 ----
|
||||
if not args.no_phase1:
|
||||
run_permutation_importance(
|
||||
model=model,
|
||||
loader=loader,
|
||||
data=data,
|
||||
num_classes=num_classes,
|
||||
device=device,
|
||||
n_permutations=args.n_permutations,
|
||||
seed=args.seed,
|
||||
out_dir=out_dir,
|
||||
)
|
||||
|
||||
# ---- Phase 2 ----
|
||||
if not args.no_phase2:
|
||||
run_gradcam(
|
||||
model=model,
|
||||
loader=loader,
|
||||
data=data,
|
||||
eval_df=eval_df,
|
||||
eval_mode=eval_mode,
|
||||
backbone=backbone,
|
||||
device=device,
|
||||
alpha=args.alpha,
|
||||
out_dir=out_dir,
|
||||
)
|
||||
|
||||
# ---- Phase 3 ----
|
||||
if not args.no_phase3:
|
||||
run_fusion_event_analysis(
|
||||
model=model,
|
||||
loader=loader,
|
||||
device=device,
|
||||
out_dir=out_dir,
|
||||
)
|
||||
|
||||
print("\n[explain_fold] Done.", flush=True)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user