added better memory caching, multithreaded processing, cleanup scripts dir
This commit is contained in:
rpotter6298
@@ -1,19 +0,0 @@
|
||||
from .clinical_data import ClinicalData
|
||||
from .dataset import ClinicalDataset
|
||||
from .image_tower import ImageTower
|
||||
from .md_tower import MDTower
|
||||
from .bridge import Bridge, VoteBridge
|
||||
# from .hypertower import HyperTower
|
||||
from .backbones import list_names, BackboneSpec, BACKBONES
|
||||
from .papila_builders import build_papila_clinical
|
||||
from .SE_attention import SEBlock, SEGateLogger
|
||||
from .early_stop import EarlyStopper
|
||||
__all__ = [
|
||||
"ClinicalData",
|
||||
"ClinicalDataset",
|
||||
"ImageTower",
|
||||
"MDTower",
|
||||
"Bridge",
|
||||
"VoteBridge",
|
||||
# "HyperTower",
|
||||
]
|
||||
|
||||
Executable
+849
@@ -0,0 +1,849 @@
|
||||
"""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
|
||||
Executable
+306
@@ -0,0 +1,306 @@
|
||||
"""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)
|
||||
Executable
+383
@@ -0,0 +1,383 @@
|
||||
"""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
|
||||
@@ -206,6 +206,7 @@ def make_loader(
|
||||
*,
|
||||
image_transform,
|
||||
image_preprocessor=None,
|
||||
image_cache=None,
|
||||
batch_size: int,
|
||||
shuffle: bool,
|
||||
num_workers: int,
|
||||
@@ -216,6 +217,7 @@ def make_loader(
|
||||
slots,
|
||||
image_transform=image_transform,
|
||||
image_preprocessor=image_preprocessor,
|
||||
image_cache=image_cache,
|
||||
)
|
||||
return DataLoader(
|
||||
ds,
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||
from typing import Any, Callable, Optional
|
||||
|
||||
from pathlib import Path
|
||||
@@ -45,12 +46,14 @@ class SlotDataset(Dataset):
|
||||
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)
|
||||
@@ -74,14 +77,72 @@ class SlotDataset(Dataset):
|
||||
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:
|
||||
|
||||
@@ -222,7 +222,13 @@ class V2HyperTower:
|
||||
ap.add_argument("--augment", action="store_true")
|
||||
ap.add_argument("--balanced-sampling", action="store_true",
|
||||
help="Use WeightedRandomSampler during training to equalise class frequency (default: off).")
|
||||
ap.add_argument("--num-workers", type=int, default=0)
|
||||
ap.add_argument("--num-workers", type=int, default=4)
|
||||
ap.add_argument("--in-memory-cache", action="store_true", default=True,
|
||||
help="Cache preprocessed images in RAM (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 in-memory image cache (default: 4).")
|
||||
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)
|
||||
@@ -452,6 +458,8 @@ class V2HyperTower:
|
||||
n_classes=num_classes,
|
||||
)
|
||||
|
||||
image_cache: dict | None = {} if getattr(args, "in_memory_cache", False) else None
|
||||
|
||||
for fold in range(n_folds):
|
||||
seed_everything(args.seed + fold * 100)
|
||||
fold_dir = tm_dir / f"fold{fold}"
|
||||
@@ -469,6 +477,7 @@ class V2HyperTower:
|
||||
fold_dir=fold_dir,
|
||||
tower_mode=tower_mode,
|
||||
pred_store=pred_store,
|
||||
image_cache=image_cache,
|
||||
)
|
||||
fold_results.append(result)
|
||||
if artifacts.y_true_ensemble is not None:
|
||||
@@ -632,6 +641,7 @@ class V2HyperTower:
|
||||
fold_dir: Path,
|
||||
tower_mode: str,
|
||||
pred_store: "PredictionStore | None" = None,
|
||||
image_cache: "dict | None" = None,
|
||||
) -> tuple[FoldResult, FoldArtifacts]:
|
||||
args = self.args
|
||||
device = self.device
|
||||
@@ -743,7 +753,8 @@ class V2HyperTower:
|
||||
|
||||
slots_eye = profile_eye.slot_descriptors()
|
||||
slots_patient = profile_patient.slot_descriptors()
|
||||
loader_kw = dict(batch_size=args.batch_size, num_workers=args.num_workers)
|
||||
loader_kw = dict(batch_size=args.batch_size, num_workers=args.num_workers,
|
||||
image_cache=image_cache)
|
||||
|
||||
# ---- loaders ---------------------------------------------------
|
||||
use_balanced = bool(getattr(args, "balanced_sampling", False))
|
||||
@@ -831,6 +842,16 @@ class V2HyperTower:
|
||||
if pred_store is not None:
|
||||
pred_store.set_split(fold, [str(s["id_1"]) for s in holdout_bilat], "holdout")
|
||||
|
||||
# ---- prebuild in-memory image cache (fold 0 only; shared dict fills for later folds) ----
|
||||
if image_cache is not None:
|
||||
cache_workers = int(getattr(args, "cache_workers", 4))
|
||||
_loaders_to_warm = [
|
||||
train_single_loader, train_bilat_loader, val_loader, holdout_loader,
|
||||
]
|
||||
for _ldr in _loaders_to_warm:
|
||||
if _ldr is not None:
|
||||
_ldr.dataset.prebuild_image_cache(cache_workers=cache_workers)
|
||||
|
||||
opt_single = torch.optim.Adam(single.parameters(), lr=args.lr) if run_single else None
|
||||
opt_bilateral = torch.optim.Adam(bilateral.parameters(), lr=args.lr) if run_bilat else None
|
||||
|
||||
@@ -942,6 +963,7 @@ class V2HyperTower:
|
||||
# ---- epoch loop ------------------------------------------------
|
||||
_prev_phase_single = "inactive" # used to detect md_warmup → next phase transition
|
||||
for epoch in range(total_epochs):
|
||||
_epoch_t0 = time.time()
|
||||
if not run_single:
|
||||
phase_single, main_epoch_single, single_active = "inactive", 0, False
|
||||
elif epoch < single_warmup_md:
|
||||
@@ -1329,6 +1351,7 @@ class V2HyperTower:
|
||||
print() # seal the progress bar line
|
||||
|
||||
if args.log_every > 0 and (epoch + 1) % args.log_every == 0:
|
||||
_epoch_secs = time.time() - _epoch_t0
|
||||
hld_auc = target_holdout_single_auc if run_single else bi_auc_h
|
||||
hld_suffix = f" hld_auc={hld_auc:.4f}" if holdout_loader is not None else ""
|
||||
|
||||
@@ -1356,7 +1379,7 @@ class V2HyperTower:
|
||||
if run_single:
|
||||
if tower_mode == "single":
|
||||
msg = (
|
||||
f" ep {epoch+1:>3}/{total_epochs} "
|
||||
f" ep {epoch+1:>3}/{total_epochs} ({_epoch_secs:.1f}s) "
|
||||
f"[single:{phase_single} {single_phase_epoch}/{single_phase_total}] "
|
||||
f"fused(acc={cl_acc:.4f},auc={cl_auc:.4f}) "
|
||||
f"img(acc={cl_acc_img:.4f},auc={cl_auc_img:.4f}) "
|
||||
@@ -1366,7 +1389,7 @@ class V2HyperTower:
|
||||
)
|
||||
else:
|
||||
msg = (
|
||||
f" ep {epoch+1:>3}/{total_epochs} "
|
||||
f" ep {epoch+1:>3}/{total_epochs} ({_epoch_secs:.1f}s) "
|
||||
f"[single:{phase_single} {single_phase_epoch}/{single_phase_total}] "
|
||||
f"fused(acc={en_acc:.4f},auc={en_auc:.4f}) "
|
||||
f"img(acc={en_acc_img:.4f},auc={en_auc_img:.4f}) "
|
||||
@@ -1376,7 +1399,7 @@ class V2HyperTower:
|
||||
)
|
||||
else:
|
||||
msg = (
|
||||
f" ep {epoch+1:>3}/{total_epochs} "
|
||||
f" ep {epoch+1:>3}/{total_epochs} ({_epoch_secs:.1f}s) "
|
||||
f"[bilat:{phase_bilat} {bilat_phase_epoch}/{bilat_phase_total}] "
|
||||
f"fused(acc={bi_acc:.4f},auc={bi_auc:.4f}) "
|
||||
f"img(acc={bi_acc_img:.4f},auc={bi_auc_img:.4f}) "
|
||||
|
||||
Reference in New Issue
Block a user