reworked iop_corr, added explainability tools and plotting tools, cleanup codebase
This commit is contained in:
rpotter6298
@@ -9,3 +9,5 @@ cache_data/
|
||||
# model artifacts
|
||||
models/refuge/
|
||||
models/v2/refuge/
|
||||
**/.archive/
|
||||
.archive/
|
||||
@@ -1,105 +0,0 @@
|
||||
# bridge.py
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from classes.SE_attention import SEBlock, SEGateLogger
|
||||
|
||||
# class SEBlock(nn.Module):
|
||||
# def __init__(self, dim: int, reduction: int = 16):
|
||||
# super().__init__()
|
||||
# hidden = max(1, dim // max(1, reduction))
|
||||
# self.net = nn.Sequential(
|
||||
# nn.Linear(dim, hidden, bias=True),
|
||||
# nn.ReLU(inplace=True),
|
||||
# nn.Linear(hidden, dim, bias=True),
|
||||
# nn.Sigmoid(),
|
||||
# )
|
||||
|
||||
# def forward(self, x):
|
||||
# return self.net(x)
|
||||
|
||||
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
|
||||
# self.se_reduction = se_reduction
|
||||
# self.se_pre_norm = se_pre_norm
|
||||
|
||||
#project towers to equal width
|
||||
self.W_img = nn.Linear(img_dim, fusion_dim)
|
||||
self.W_md = nn.Linear(meta_dim, fusion_dim)
|
||||
|
||||
#(optional) : set layernorm for se so one tower doesn't dominate the other
|
||||
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,264 +0,0 @@
|
||||
# clinical_data.py
|
||||
from __future__ import annotations
|
||||
from pathlib import Path
|
||||
from typing import Iterable, Optional, Dict, List, Tuple
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
|
||||
class ClinicalData:
|
||||
"""
|
||||
Torch-free container for clinical metadata and file/label bookkeeping.
|
||||
- Holds one or more dataframes (via add_df) and harmonizes columns
|
||||
- Canonical IDs: 'Patient ID' must exist (or be specified and will be renamed)
|
||||
- Canonical eye column: 'eyeID' recoded to 'OS'/'OD' if present; if absent, set to 0
|
||||
- Feature typing (if cat_cols not provided):
|
||||
* Categorical if (a) <= max_unique categorical threshold (default 4), or
|
||||
(b) values cannot be coerced to float; otherwise numeric (scalar)
|
||||
- Scaling/imputation:
|
||||
* Numeric: min–max to [0,1], median imputation; + one missing flag per numeric feature
|
||||
* Categorical: one-hot with '<UNK>' bucket at index 0
|
||||
- Patient-level K-fold indices stored as dict: folds[k] -> {'train_ids': [...], 'test_ids': [...]}
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
image_dir: str,
|
||||
clinical_dir: Optional[str],
|
||||
label_col: str,
|
||||
# typing / detection
|
||||
cat_cols: Optional[Iterable[str]] = None,
|
||||
max_unique_for_cat: int = 4,
|
||||
# splitting
|
||||
n_splits: int = 5,
|
||||
random_seed: int = 42,
|
||||
):
|
||||
self.image_dir = Path(image_dir)
|
||||
self.clinical_dir = Path(clinical_dir) if clinical_dir else None
|
||||
self.label_col = label_col
|
||||
self.max_unique_for_cat = max_unique_for_cat
|
||||
self.n_splits = n_splits
|
||||
|
||||
# Internal state
|
||||
self.frames: List[pd.DataFrame] = [] # raw frames as added
|
||||
self.df: pd.DataFrame = pd.DataFrame() # concatenated
|
||||
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]]] = {} # fold -> {'train_ids': [], 'test_ids': []}
|
||||
self.random_seed = int(random_seed)
|
||||
|
||||
# ------------------- Public API -------------------
|
||||
def add_df(
|
||||
self,
|
||||
df: pd.DataFrame,
|
||||
id_column: Optional[str] = None,
|
||||
eye_column: Optional[str] = None,
|
||||
exclude_cols: Optional[Iterable[str]] = None,
|
||||
) -> None:
|
||||
"""
|
||||
Add a dataframe and re-run harmonization, typing, stats, and K-fold indices.
|
||||
QC rules:
|
||||
- Must have patient ID column; if not provided under that name, specify id_column.
|
||||
- eyeID, if present, must be binary; recoded to 'OS'/'OD'. If absent, create and set to 0.
|
||||
"""
|
||||
df = df.copy()
|
||||
# --- QC: Patient ID ---
|
||||
pid_col = self._ensure_patient_id(df, id_column)
|
||||
# --- QC: eyeID ---
|
||||
self._canonicalize_eye_column(df, eye_column)
|
||||
# --- Normalize label presence ---
|
||||
if self.label_col not in df.columns:
|
||||
raise ValueError(f"label_col '{self.label_col}' not found in added dataframe")
|
||||
|
||||
# append & refresh
|
||||
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['Patient ID'].isin(train_ids)].reset_index(drop=True)
|
||||
test_df = self.df[self.df['Patient ID'].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, filename_template: str = "RET{pid:03d}{eye}.jpg") -> Path:
|
||||
pid = int(row['Patient ID']); eye = row.get('eyeID', 0)
|
||||
if eye in ("OS", "OD"):
|
||||
eye_str = eye
|
||||
else:
|
||||
eye_str = str(eye)
|
||||
return self.image_dir / filename_template.format(pid=pid, eye=eye_str)
|
||||
|
||||
# ------------------- Internal helpers -------------------
|
||||
def _ensure_patient_id(self, df: pd.DataFrame, id_column: Optional[str]) -> str:
|
||||
if 'Patient ID' in df.columns:
|
||||
return 'Patient ID'
|
||||
if id_column and id_column in df.columns:
|
||||
df.rename(columns={id_column: 'Patient ID'}, inplace=True)
|
||||
return 'Patient ID'
|
||||
# try auto-detect common variants
|
||||
candidates = [c for c in df.columns if c.lower().replace(" ", "") in {"patientid","patient","pid"}]
|
||||
if len(candidates) == 1:
|
||||
df.rename(columns={candidates[0]: 'Patient ID'}, inplace=True)
|
||||
return 'Patient ID'
|
||||
raise ValueError("A 'Patient ID' column is required; provide id_column=... if it has a different name.")
|
||||
|
||||
def _canonicalize_eye_column(self, df: pd.DataFrame, eye_column: Optional[str]) -> None:
|
||||
# Find source
|
||||
src = None
|
||||
if 'eyeID' in df.columns: src = 'eyeID'
|
||||
elif eye_column and eye_column in df.columns: src = eye_column
|
||||
else:
|
||||
# try auto detect
|
||||
for c in df.columns:
|
||||
if 'eye' in c.lower():
|
||||
src = c; break
|
||||
if src is None:
|
||||
df['eyeID'] = 0
|
||||
return
|
||||
# Map to OS/OD
|
||||
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"
|
||||
# numbers like 2? fall back by parity
|
||||
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")
|
||||
if src != 'eyeID':
|
||||
# keep original too if you want, but we standardize on 'eyeID'
|
||||
pass
|
||||
|
||||
def _refresh_master_df(self, exclude_cols: Optional[Iterable[str]] = None) -> None:
|
||||
self.df = pd.concat(self.frames, axis=0, ignore_index=True)
|
||||
# drop columns explicitly excluded
|
||||
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, 'Patient ID'}
|
||||
# we keep canonical 'eyeID' as categorical if present
|
||||
feature_candidates = [c for c in self.df.columns if c not in excluded]
|
||||
# If user pre-specified cat_cols in __init__, respect them and infer the rest
|
||||
cats = set(self.cat_cols) if self.cat_cols else set()
|
||||
scalars = set()
|
||||
for c in feature_candidates:
|
||||
if c == 'eyeID':
|
||||
cats.add('eyeID'); continue
|
||||
if c in cats: continue
|
||||
s = self.df[c]
|
||||
# try numeric coercion
|
||||
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:
|
||||
# categorical if few uniques OR non-numeric
|
||||
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:
|
||||
# unique patients and a per-patient label for stratification if possible
|
||||
pats = self.df['Patient ID'].unique().tolist()
|
||||
# Derive a patient label as the mode of their rows (fallback to first valid)
|
||||
labels_by_pat = {}
|
||||
for pid, grp in self.df.groupby('Patient ID'):
|
||||
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 to use StratifiedGroupKFold if available, else fall back to StratifiedKFold on patient labels
|
||||
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,58 +0,0 @@
|
||||
# dataset.py
|
||||
from torch.utils.data import Dataset
|
||||
from PIL import Image
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
|
||||
class ClinicalDataset(Dataset):
|
||||
"""Generic dataset wrapping a ClinicalData 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,
|
||||
):
|
||||
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
|
||||
|
||||
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)
|
||||
orig_img = Image.open(img_path).convert("RGB")
|
||||
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
|
||||
@@ -1,89 +0,0 @@
|
||||
import math, copy, torch
|
||||
|
||||
class EarlyStopper:
|
||||
def __init__(self, monitor: str, mode: str = "auto",
|
||||
patience: int = 5, min_delta: float = 0.0,
|
||||
save_path: str | None = None, restore_best: bool = True):
|
||||
"""
|
||||
monitor: key in your epoch row, e.g. 'eval_loss', 'auc_fused', 'acc_fused'
|
||||
mode: 'max' (higher is better), 'min', or 'auto' (min for '*loss*', else max)
|
||||
patience: epochs without improvement before stopping
|
||||
min_delta: required improvement magnitude
|
||||
save_path: optional .pth file to save best weights each time it improves
|
||||
restore_best: if True, load best weights back at the end
|
||||
"""
|
||||
self.monitor = monitor
|
||||
if mode == "auto":
|
||||
mode = "min" if "loss" in monitor.lower() else "max"
|
||||
self.mode = mode
|
||||
self.patience = int(patience)
|
||||
self.min_delta = float(min_delta)
|
||||
self.save_path = save_path
|
||||
self.restore_best = restore_best
|
||||
|
||||
self.best = -math.inf if mode == "max" else math.inf
|
||||
self.bad_epochs = 0
|
||||
self.best_state = None
|
||||
self.best_epoch = -1
|
||||
self.last_improved = False
|
||||
|
||||
def _is_better(self, val):
|
||||
if val is None or (isinstance(val, float) and math.isnan(val)):
|
||||
return False
|
||||
if self.mode == "max":
|
||||
return val > (self.best + self.min_delta)
|
||||
else:
|
||||
return val < (self.best - self.min_delta)
|
||||
|
||||
def step(self, metrics: dict, trainer, epoch: int) -> bool:
|
||||
val = metrics.get(self.monitor, None)
|
||||
improved = self._is_better(val)
|
||||
self.last_improved = improved
|
||||
|
||||
if improved:
|
||||
self.best = val
|
||||
self.best_epoch = epoch
|
||||
self.bad_epochs = 0
|
||||
# snapshot + optional save
|
||||
state = {
|
||||
"img_tower": trainer.img_tower.state_dict(),
|
||||
"md_tower": trainer.md_tower.state_dict(),
|
||||
"optimizer": trainer.optimizer.state_dict(),
|
||||
}
|
||||
if hasattr(trainer, "bridge"): state["bridge"] = trainer.bridge.state_dict()
|
||||
if hasattr(trainer, "head_img"): state["head_img"] = trainer.head_img.state_dict()
|
||||
if hasattr(trainer, "head_md"): state["head_md"] = trainer.head_md.state_dict()
|
||||
# keep an in-memory copy for restore(); file save is optional
|
||||
self.best_state = copy.deepcopy(state)
|
||||
if self.save_path: torch.save(state, self.save_path)
|
||||
print(f"[early] ↑ new best {self.monitor}={val:.5f} at epoch {epoch+1}")
|
||||
else:
|
||||
self.bad_epochs += 1
|
||||
|
||||
stop = self.bad_epochs >= self.patience
|
||||
if stop:
|
||||
print(f"[early] stopping: no improvement in {self.patience} epochs "
|
||||
f"(best {self.monitor}={self.best:.5f} @ epoch {self.best_epoch+1})")
|
||||
return stop
|
||||
|
||||
def restore(self, trainer):
|
||||
if not self.restore_best:
|
||||
return
|
||||
# Prefer in-memory best state; otherwise try loading from save_path
|
||||
st = self.best_state
|
||||
if st is None and self.save_path:
|
||||
try:
|
||||
st = torch.load(self.save_path, map_location="cpu")
|
||||
except Exception:
|
||||
st = None
|
||||
if st is None:
|
||||
return
|
||||
trainer.img_tower.load_state_dict(st["img_tower"])
|
||||
trainer.md_tower.load_state_dict(st["md_tower"])
|
||||
if "bridge" in st and hasattr(trainer, "bridge"):
|
||||
trainer.bridge.load_state_dict(st["bridge"])
|
||||
if "head_img" in st and hasattr(trainer, "head_img"):
|
||||
trainer.head_img.load_state_dict(st["head_img"])
|
||||
if "head_md" in st and hasattr(trainer, "head_md"):
|
||||
trainer.head_md.load_state_dict(st["head_md"])
|
||||
trainer.optimizer.load_state_dict(st["optimizer"])
|
||||
-1402
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -1,125 +0,0 @@
|
||||
# classes/image_tower.py
|
||||
from __future__ import annotations
|
||||
import math
|
||||
from typing import Optional
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
from torchvision import transforms
|
||||
from classes.backbones import BACKBONES, list_names, load_backbone_weights
|
||||
from classes.SE_attention import SEBlock
|
||||
|
||||
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.
|
||||
ratio in [0,1]."""
|
||||
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
|
||||
@@ -1,54 +0,0 @@
|
||||
# md_tower.py
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from classes import ClinicalData
|
||||
from classes.SE_attention import SEBlock
|
||||
|
||||
class MDTower(nn.Module):
|
||||
"""MLP over ClinicalData.vectorize_row outputs (convert to torch inside tower)."""
|
||||
def __init__(self, clinical_data: ClinicalData, 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.
|
||||
With two blocks, ratio≥0.5 freezes block0; ratio≥1.0 freezes both."""
|
||||
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,99 +0,0 @@
|
||||
# papila_builders.py
|
||||
from typing import List, Dict
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
|
||||
from classes import ClinicalData # adjust import path if needed
|
||||
|
||||
# ---- 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 _pick_iop(row: pd.Series) -> float:
|
||||
"""Prefer Pneumatic, else Perkins; may return NaN."""
|
||||
raw = row["Pneumatic"] if not pd.isna(row.get("Pneumatic", np.nan)) else row.get("Perkins", np.nan)
|
||||
return float(raw) if not pd.isna(raw) else np.nan
|
||||
|
||||
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) -> pd.DataFrame:
|
||||
"""Add IOP_raw/IOP_corr and drop VF_MD if present (in-place safe)."""
|
||||
# IOP_raw
|
||||
df["IOP_raw"] = df.apply(_pick_iop, axis=1)
|
||||
|
||||
# IOP_corr
|
||||
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 VF_MD if present
|
||||
if "VF_MD" in df.columns:
|
||||
df.drop(columns=["VF_MD"], inplace=True)
|
||||
return df
|
||||
|
||||
|
||||
def build_papila_clinical(
|
||||
image_dir: str,
|
||||
clinical_dir: str,
|
||||
label_col: str,
|
||||
cat_cols: List[str],
|
||||
n_splits: int = 5,
|
||||
random_seed: int = 42,
|
||||
) -> ClinicalData:
|
||||
"""
|
||||
Build ClinicalData exactly like the user's original build_clinical:
|
||||
- add_df(OD), set eyeID='OD'
|
||||
- add_df(OS), set eyeID='OS'
|
||||
- normalize 'Patient ID' on frames
|
||||
THEN:
|
||||
- compute IOP_raw / IOP_corr on each frame
|
||||
- drop VF_MD
|
||||
- refresh master df + kfold indices
|
||||
"""
|
||||
clinical = ClinicalData(
|
||||
image_dir=image_dir,
|
||||
clinical_dir=clinical_dir,
|
||||
label_col=label_col,
|
||||
cat_cols=cat_cols,
|
||||
n_splits=n_splits,
|
||||
random_seed=random_seed,
|
||||
)
|
||||
|
||||
# --- Load exactly like original build_clinical ---
|
||||
clinical.add_df(pd.read_excel(f"{clinical_dir}/patient_data_od.xlsx", header=1), id_column="ID")
|
||||
clinical.frames[0]["eyeID"] = "OD"
|
||||
|
||||
clinical.add_df(pd.read_excel(f"{clinical_dir}/patient_data_os.xlsx", header=1), id_column="ID")
|
||||
clinical.frames[1]["eyeID"] = "OS"
|
||||
|
||||
# Normalize 'Patient ID' on the per-eye frames (string → int)
|
||||
for frame in clinical.frames:
|
||||
frame["Patient ID"] = frame["Patient ID"].astype(str).str.extract(r"(\d+)")[0].astype(int)
|
||||
|
||||
# Build initial master as in original
|
||||
clinical._refresh_master_df()
|
||||
|
||||
# --- Post-processing ON THE FRAMES (so everything stays consistent) ---
|
||||
for i in range(len(clinical.frames)):
|
||||
clinical.frames[i] = _apply_iop_and_drop_md(clinical.frames[i])
|
||||
|
||||
# Refresh master again so IOP_raw/IOP_corr & MD removal propagate
|
||||
clinical._refresh_master_df()
|
||||
clinical._build_kfold_indices()
|
||||
|
||||
return clinical
|
||||
@@ -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,149 +0,0 @@
|
||||
# tower_watcher.py
|
||||
import matplotlib.pyplot as plt
|
||||
|
||||
# tower_watcher.py
|
||||
import matplotlib.pyplot as plt
|
||||
|
||||
|
||||
class TowerWatcher:
|
||||
"""
|
||||
Live monitor:
|
||||
- Cumulative batch-level: loss & accuracy per batch across all epochs.
|
||||
- Epoch batch-level: loss & accuracy per batch within the current epoch (resets each epoch).
|
||||
- TP/FP/TN/FN bar charts per tower, one chart each, new group each epoch.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
plt.ion()
|
||||
# 2 line plots (cum loss, cum acc), 2 line plots (epoch loss, epoch acc), 3 bar plots
|
||||
self.fig, self.axs = plt.subplots(7, 1, figsize=(10, 28))
|
||||
self.reset()
|
||||
|
||||
def reset(self):
|
||||
# Cumulative batch-level
|
||||
self.global_batches = []
|
||||
self.loss_cum = {"fusion": [], "image": [], "meta": []}
|
||||
self.acc_cum = {"fusion": [], "image": [], "meta": []}
|
||||
# Epoch batch-level
|
||||
self.epoch_batches = []
|
||||
self.loss_epoch_batch = {"fusion": [], "image": [], "meta": []}
|
||||
self.acc_epoch_batch = {"fusion": [], "image": [], "meta": []}
|
||||
# Epoch markers for cum plots
|
||||
self.epoch_markers = []
|
||||
# Stats per epoch for bars
|
||||
self.epoch_stats = {"fusion": [], "image": [], "meta": []}
|
||||
# Track current epoch
|
||||
self.current_epoch = -1
|
||||
|
||||
def on_epoch_start(self, epoch):
|
||||
# mark epoch boundary in cumulative
|
||||
x = self.global_batches[-1] + 1 if self.global_batches else 0
|
||||
self.epoch_markers.append(x)
|
||||
# reset epoch batch-level data
|
||||
self.epoch_batches = []
|
||||
for d in [self.loss_epoch_batch, self.acc_epoch_batch]:
|
||||
for k in d:
|
||||
d[k].clear()
|
||||
self.current_epoch = epoch
|
||||
|
||||
def on_batch_end(self, idx, stats: dict):
|
||||
# Cumulative
|
||||
self.global_batches.append(len(self.global_batches) + 1)
|
||||
for key, lk, ak in [
|
||||
("fusion", "loss_f", "acc_f"),
|
||||
("image", "loss_i", "acc_i"),
|
||||
("meta", "loss_m", "acc_m"),
|
||||
]:
|
||||
self.loss_cum[key].append(stats.get(lk, 0))
|
||||
self.acc_cum[key].append(stats.get(ak, 0))
|
||||
# Epoch-level
|
||||
self.epoch_batches.append(len(self.epoch_batches) + 1)
|
||||
for key, lk, ak in [
|
||||
("fusion", "loss_f", "acc_f"),
|
||||
("image", "loss_i", "acc_i"),
|
||||
("meta", "loss_m", "acc_m"),
|
||||
]:
|
||||
self.loss_epoch_batch[key].append(stats.get(lk, 0))
|
||||
self.acc_epoch_batch[key].append(stats.get(ak, 0))
|
||||
# redraw
|
||||
self._draw_batch_plots()
|
||||
|
||||
def on_epoch_end(self, epoch, stats: dict):
|
||||
# record per-epoch TP/FP/TN/FN
|
||||
for key in ["fusion", "image", "meta"]:
|
||||
self.epoch_stats[key].append(
|
||||
{
|
||||
"tp": stats.get("tp", 0),
|
||||
"fp": stats.get("fp", 0),
|
||||
"tn": stats.get("tn", 0),
|
||||
"fn": stats.get("fn", 0),
|
||||
}
|
||||
)
|
||||
self._draw_epoch_bars()
|
||||
|
||||
def _draw_batch_plots(self):
|
||||
# Cumulative Loss
|
||||
ax = self.axs[0]
|
||||
ax.clear()
|
||||
ax.plot(self.global_batches, self.loss_cum["fusion"], label="Fusion")
|
||||
ax.plot(self.global_batches, self.loss_cum["image"], label="Image Tower")
|
||||
ax.plot(self.global_batches, self.loss_cum["meta"], label="MD Tower")
|
||||
for x in self.epoch_markers:
|
||||
ax.axvline(x=x, color="gray", linestyle="--")
|
||||
ax.set_ylabel("Cumulative Loss")
|
||||
ax.legend()
|
||||
|
||||
# Epoch Loss
|
||||
ax = self.axs[1]
|
||||
ax.clear()
|
||||
ax.plot(self.epoch_batches, self.loss_epoch_batch["fusion"], label="Fusion")
|
||||
ax.plot(self.epoch_batches, self.loss_epoch_batch["image"], label="Image Tower")
|
||||
ax.plot(self.epoch_batches, self.loss_epoch_batch["meta"], label="MD Tower")
|
||||
ax.set_ylabel(f"Epoch {self.current_epoch+1} Loss")
|
||||
ax.set_xlabel("Batch (Epoch)")
|
||||
ax.legend()
|
||||
|
||||
# Cumulative Accuracy
|
||||
ax = self.axs[2]
|
||||
ax.clear()
|
||||
ax.plot(self.global_batches, self.acc_cum["fusion"], label="Fusion")
|
||||
ax.plot(self.global_batches, self.acc_cum["image"], label="Image Tower")
|
||||
ax.plot(self.global_batches, self.acc_cum["meta"], label="MD Tower")
|
||||
for x in self.epoch_markers:
|
||||
ax.axvline(x=x, color="gray", linestyle="--")
|
||||
ax.set_ylabel("Cumulative Accuracy")
|
||||
ax.legend()
|
||||
|
||||
# Epoch Accuracy
|
||||
ax = self.axs[3]
|
||||
ax.clear()
|
||||
ax.plot(self.epoch_batches, self.acc_epoch_batch["fusion"], label="Fusion")
|
||||
ax.plot(self.epoch_batches, self.acc_epoch_batch["image"], label="Image Tower")
|
||||
ax.plot(self.epoch_batches, self.acc_epoch_batch["meta"], label="MD Tower")
|
||||
ax.set_ylabel(f"Epoch {self.current_epoch+1} Accuracy")
|
||||
ax.set_xlabel("Batch (Epoch)")
|
||||
ax.legend()
|
||||
|
||||
plt.pause(0.01)
|
||||
|
||||
def _draw_epoch_bars(self):
|
||||
# Bar charts per tower
|
||||
for i, key in enumerate(["fusion", "image", "meta"]):
|
||||
ax = self.axs[4 + i]
|
||||
ax.clear()
|
||||
data = self.epoch_stats[key]
|
||||
epochs = list(range(1, len(data) + 1))
|
||||
tp = [d["tp"] for d in data]
|
||||
fp = [d["fp"] for d in data]
|
||||
tn = [d["tn"] for d in data]
|
||||
fn = [d["fn"] for d in data]
|
||||
width = 0.2
|
||||
ax.bar([e - width for e in epochs], tp, width, label="TP")
|
||||
ax.bar(epochs, fp, width, label="FP")
|
||||
ax.bar([e + width for e in epochs], tn, width, label="TN")
|
||||
ax.bar([e + 2 * width for e in epochs], fn, width, label="FN")
|
||||
ax.set_title(f"{key.title()} Tower Stats")
|
||||
ax.set_xlabel("Epoch")
|
||||
ax.set_ylabel("Count")
|
||||
ax.legend()
|
||||
plt.pause(0.01)
|
||||
@@ -3,7 +3,7 @@ from __future__ import annotations
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from classes.SE_attention import SEBlock, SEGateLogger
|
||||
from classes.v2.SE_attention import SEBlock, SEGateLogger
|
||||
|
||||
|
||||
class Bridge(nn.Module):
|
||||
|
||||
+24
-3
@@ -10,9 +10,30 @@ import torch
|
||||
from PIL import Image, ImageDraw
|
||||
from torchvision import transforms
|
||||
|
||||
from classes.geometry_features import compute_geometry_features, disc_cup_from_mask_image
|
||||
from classes.refuge_classification import _geometry_from_mask
|
||||
from classes.unet_segmenter import UNetSegmenter
|
||||
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:
|
||||
|
||||
@@ -219,6 +219,15 @@ def _set_single_phase(model: SingleEyeHT, phase: str) -> None:
|
||||
# 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")
|
||||
@@ -282,12 +291,27 @@ def train_single_epoch(
|
||||
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)
|
||||
|
||||
|
||||
@@ -33,10 +33,21 @@ def _nearest_pachy_key(x: float) -> int:
|
||||
return int(_PACHY_KEYS[idx])
|
||||
|
||||
|
||||
# Ratio derived from patients with both Pneumatic and Perkins readings (n=41, OD+OS combined).
|
||||
# Pneumatic / Perkins mean ratio = 1.158; applied to Perkins-only rows to put them on the
|
||||
# Pneumatic scale before IOP_corr is computed.
|
||||
_PERKINS_TO_PNEUMATIC_RATIO: float = 1.158
|
||||
|
||||
|
||||
def _pick_iop(row: pd.Series) -> float:
|
||||
"""Prefer Pneumatic, else Perkins; may return NaN."""
|
||||
raw = row["Pneumatic"] if not pd.isna(row.get("Pneumatic", np.nan)) else row.get("Perkins", np.nan)
|
||||
return float(raw) if not pd.isna(raw) else np.nan
|
||||
"""Prefer Pneumatic; scale 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 not pd.isna(perkins):
|
||||
return float(perkins) * _PERKINS_TO_PNEUMATIC_RATIO
|
||||
return np.nan
|
||||
|
||||
|
||||
def _correct_iop(raw_iop: float, pachy: float) -> float:
|
||||
|
||||
@@ -7,8 +7,8 @@ import torch
|
||||
from torch import nn
|
||||
from torchvision import transforms
|
||||
|
||||
from classes.backbones import BACKBONES, list_names, load_backbone_weights
|
||||
from classes.SE_attention import SEBlock
|
||||
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
|
||||
|
||||
|
||||
|
||||
@@ -7,7 +7,7 @@ from PIL import Image
|
||||
|
||||
from torchvision import transforms
|
||||
|
||||
from classes.backbones import BACKBONES
|
||||
from classes.v2.backbones import BACKBONES
|
||||
|
||||
|
||||
IMAGENET_MEAN: Tuple[float, float, float] = (0.485, 0.456, 0.406)
|
||||
|
||||
+60
-14
@@ -206,6 +206,9 @@ class V2HyperTower:
|
||||
help="Single-eye model tower warmup (overrides --warmup-tower-epochs).")
|
||||
ap.add_argument("--single-warmup-fused-epochs", type=int, default=None,
|
||||
help="Single-eye model fused warmup (overrides --warmup-fused-epochs).")
|
||||
ap.add_argument("--warmup-md-epochs", type=int, default=0,
|
||||
help="MD-only warmup epochs before tower warmup. Trains only md_tower + "
|
||||
"classifier_md (no CNN forward pass, so 50-100 epochs is cheap).")
|
||||
ap.add_argument("--bilat-warmup-tower-epochs", type=int, default=None,
|
||||
help="Bilateral model tower warmup (overrides --warmup-tower-epochs).")
|
||||
ap.add_argument("--bilat-warmup-fused-epochs", type=int, default=None,
|
||||
@@ -424,7 +427,8 @@ class V2HyperTower:
|
||||
if getattr(args, "single_warmup_fused_epochs", None) is not None
|
||||
else int(_global_warmup_fused) if _global_warmup_fused is not None else 2
|
||||
)
|
||||
_total_epochs = _warmup_tower + _warmup_fused + int(args.epochs) + fusion_epochs
|
||||
_warmup_md = int(getattr(args, "warmup_md_epochs", 0))
|
||||
_total_epochs = _warmup_md + _warmup_tower + _warmup_fused + int(args.epochs) + fusion_epochs
|
||||
|
||||
# sample IDs depend on mode: single uses eye IDs, others use patient IDs
|
||||
if tower_mode in ("single", "classic"):
|
||||
@@ -572,6 +576,7 @@ class V2HyperTower:
|
||||
"run_id": run_name,
|
||||
"backbone": args.backbone,
|
||||
"epochs": args.epochs,
|
||||
"warmup_md_epochs": getattr(args, "warmup_md_epochs", 0),
|
||||
"warmup_tower_epochs": args.warmup_tower_epochs,
|
||||
"warmup_fused_epochs": args.warmup_fused_epochs,
|
||||
"single_warmup_tower_epochs": args.single_warmup_tower_epochs,
|
||||
@@ -659,6 +664,7 @@ class V2HyperTower:
|
||||
if getattr(args, "bilat_warmup_fused_epochs", None) is not None
|
||||
else int(global_warmup_fused) if global_warmup_fused is not None else 3
|
||||
)
|
||||
single_warmup_md = int(getattr(args, "warmup_md_epochs", 0)) if run_single else 0
|
||||
if not run_single:
|
||||
single_warmup_tower = 0
|
||||
single_warmup_fused = 0
|
||||
@@ -666,7 +672,7 @@ class V2HyperTower:
|
||||
bilat_warmup_tower = 0
|
||||
bilat_warmup_fused = 0
|
||||
main_epochs = int(args.epochs)
|
||||
total_single_epochs = (single_warmup_tower + single_warmup_fused + main_epochs) if run_single else 0
|
||||
total_single_epochs = (single_warmup_md + single_warmup_tower + single_warmup_fused + main_epochs) if run_single else 0
|
||||
total_bilat_epochs = (bilat_warmup_tower + bilat_warmup_fused + main_epochs) if run_bilat else 0
|
||||
total_epochs = max(total_single_epochs, total_bilat_epochs)
|
||||
|
||||
@@ -744,6 +750,7 @@ class V2HyperTower:
|
||||
train_single_loader = None
|
||||
train_eval_loader = None # non-shuffled, no sampler — for per-epoch train logging
|
||||
train_bilat_loader = None
|
||||
md_only_loader = None # image-free loader for md_warmup phase
|
||||
if run_single:
|
||||
single_sampler = build_balanced_sampler(eye_train) if use_balanced else None
|
||||
train_single_loader = make_loader(
|
||||
@@ -761,6 +768,20 @@ class V2HyperTower:
|
||||
shuffle=False,
|
||||
**loader_kw,
|
||||
)
|
||||
if single_warmup_md > 0:
|
||||
# MD-only loader: drop image_1 so PIL never opens files during md_warmup.
|
||||
# Always use balanced sampling for md_warmup — MD features alone are weaker
|
||||
# than images and collapse to majority class without class balancing.
|
||||
slots_md_only = {k: v for k, v in slots_eye.items() if k != "image_1"}
|
||||
md_warmup_sampler = single_sampler if single_sampler is not None else build_balanced_sampler(eye_train)
|
||||
md_only_loader = make_loader(
|
||||
eye_train, slots_md_only,
|
||||
image_transform=None,
|
||||
image_preprocessor=None,
|
||||
shuffle=True,
|
||||
sampler=md_warmup_sampler,
|
||||
**loader_kw,
|
||||
)
|
||||
if run_bilat:
|
||||
bilat_sampler = build_balanced_sampler(bilat_train) if use_balanced else None
|
||||
train_bilat_loader = make_loader(
|
||||
@@ -907,7 +928,7 @@ class V2HyperTower:
|
||||
print(
|
||||
f" [fold {fold+1}] single_train_n={len(eye_train)} (eye-level) "
|
||||
f"val_n={len(bilat_val)} "
|
||||
f"single_warmup={single_warmup_tower}+{single_warmup_fused} total={total_single_epochs}",
|
||||
f"single_warmup=md{single_warmup_md}+twr{single_warmup_tower}+fus{single_warmup_fused} total={total_single_epochs}",
|
||||
flush=True,
|
||||
)
|
||||
else:
|
||||
@@ -919,17 +940,20 @@ class V2HyperTower:
|
||||
)
|
||||
|
||||
# ---- epoch loop ------------------------------------------------
|
||||
_prev_phase_single = "inactive" # used to detect md_warmup → next phase transition
|
||||
for epoch in range(total_epochs):
|
||||
if not run_single:
|
||||
phase_single, main_epoch_single, single_active = "inactive", 0, False
|
||||
elif epoch < single_warmup_tower:
|
||||
elif epoch < single_warmup_md:
|
||||
phase_single, main_epoch_single, single_active = "md_warmup", 0, True
|
||||
elif epoch < (single_warmup_md + single_warmup_tower):
|
||||
phase_single, main_epoch_single, single_active = "tower_warmup", 0, True
|
||||
elif epoch < (single_warmup_tower + single_warmup_fused):
|
||||
elif epoch < (single_warmup_md + single_warmup_tower + single_warmup_fused):
|
||||
phase_single, main_epoch_single, single_active = "fused_warmup", 0, True
|
||||
elif epoch < total_single_epochs:
|
||||
phase_single, main_epoch_single, single_active = (
|
||||
"main",
|
||||
epoch - single_warmup_tower - single_warmup_fused + 1,
|
||||
epoch - single_warmup_md - single_warmup_tower - single_warmup_fused + 1,
|
||||
True,
|
||||
)
|
||||
else:
|
||||
@@ -951,8 +975,9 @@ class V2HyperTower:
|
||||
phase_bilat, main_epoch_bilat, bilat_active = "done", main_epochs, False
|
||||
|
||||
if run_single and single_active:
|
||||
_active_loader = md_only_loader if phase_single == "md_warmup" else train_single_loader
|
||||
sl_loss, sl_acc = train_single_epoch(
|
||||
single, train_single_loader, opt_single, device,
|
||||
single, _active_loader, opt_single, device,
|
||||
phase=phase_single, bcd_prob=float(args.bcd_prob),
|
||||
)
|
||||
else:
|
||||
@@ -966,7 +991,9 @@ class V2HyperTower:
|
||||
else:
|
||||
bl_loss, bl_acc = nan, nan
|
||||
|
||||
if run_single and tower_mode == "single":
|
||||
_skip_val_eval = (phase_single == "md_warmup")
|
||||
|
||||
if run_single and tower_mode == "single" and not _skip_val_eval:
|
||||
y_cl, p_cl, p_cl_img, p_cl_md = collect_probs_single_components(
|
||||
single, val_loader, device, aggregate_patient=False
|
||||
)
|
||||
@@ -980,7 +1007,7 @@ class V2HyperTower:
|
||||
en_acc = en_auc = nan
|
||||
en_n = 0
|
||||
en_acc_img = en_acc_md = en_auc_img = en_auc_md = nan
|
||||
elif run_single and tower_mode == "ensemble":
|
||||
elif run_single and tower_mode == "ensemble" and not _skip_val_eval:
|
||||
(y_en,
|
||||
_p_en_f_od, _p_en_i_od, _p_en_m_od,
|
||||
_p_en_f_os, _p_en_i_os, _p_en_m_os,
|
||||
@@ -1010,7 +1037,7 @@ class V2HyperTower:
|
||||
cl_acc_img = cl_acc_md = en_acc_img = en_acc_md = nan
|
||||
cl_auc_img = cl_auc_md = en_auc_img = en_auc_md = nan
|
||||
|
||||
if run_bilat:
|
||||
if run_bilat and not _skip_val_eval:
|
||||
y_bi, p_bi, p_bi_img, p_bi_md = collect_probs_bilateral_components(
|
||||
bilateral, val_loader, device
|
||||
)
|
||||
@@ -1034,7 +1061,7 @@ class V2HyperTower:
|
||||
p_cl_h = p_cl_h_img = p_cl_h_md = _z2
|
||||
p_en_h = p_en_h_img = p_en_h_md = _z2
|
||||
|
||||
if holdout_loader is not None:
|
||||
if holdout_loader is not None and not _skip_val_eval:
|
||||
if run_single and tower_mode == "single":
|
||||
y_cl_h, p_cl_h, p_cl_h_img, p_cl_h_md = collect_probs_single_components(
|
||||
single, holdout_loader, device, aggregate_patient=False
|
||||
@@ -1092,7 +1119,7 @@ class V2HyperTower:
|
||||
tr_fe_corr = tr_fe_err = tr_n = 0
|
||||
y_tr = np.array([], dtype=np.int64)
|
||||
p_tr_f = p_tr_i = p_tr_m = np.zeros((0, num_classes), dtype=np.float32)
|
||||
if run_single and train_eval_loader is not None:
|
||||
if run_single and train_eval_loader is not None and not _skip_val_eval:
|
||||
y_tr, p_tr_f, p_tr_i, p_tr_m, tr_ids = collect_probs_eye_level(
|
||||
single, train_eval_loader, device, return_ids=True
|
||||
)
|
||||
@@ -1284,16 +1311,33 @@ class V2HyperTower:
|
||||
**cm_row,
|
||||
}, optional_cols=epoch_fields)
|
||||
|
||||
# ---- md_warmup progress bar (replaces per-epoch print) --------
|
||||
if phase_single == "md_warmup":
|
||||
_bar_w = 30
|
||||
_filled = int(_bar_w * (epoch + 1) / single_warmup_md)
|
||||
_bar = "#" * _filled + "-" * (_bar_w - _filled)
|
||||
_bar_msg = (
|
||||
f" [fold {fold+1}] md_warmup [{_bar}] "
|
||||
f"{epoch + 1}/{single_warmup_md} loss={sl_loss:.4f}"
|
||||
)
|
||||
print(f"\r{_bar_msg}", end="", flush=True)
|
||||
fold_logger.info(_bar_msg)
|
||||
_prev_phase_single = phase_single
|
||||
continue # skip normal log block entirely
|
||||
|
||||
if _prev_phase_single == "md_warmup":
|
||||
print() # seal the progress bar line
|
||||
|
||||
if args.log_every > 0 and (epoch + 1) % args.log_every == 0:
|
||||
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 ""
|
||||
|
||||
# Human-readable phase progress for console logs.
|
||||
if phase_single == "tower_warmup":
|
||||
single_phase_epoch = epoch + 1
|
||||
single_phase_epoch = epoch - single_warmup_md + 1
|
||||
single_phase_total = single_warmup_tower
|
||||
elif phase_single == "fused_warmup":
|
||||
single_phase_epoch = epoch - single_warmup_tower + 1
|
||||
single_phase_epoch = epoch - single_warmup_md - single_warmup_tower + 1
|
||||
single_phase_total = single_warmup_fused
|
||||
else:
|
||||
single_phase_epoch = main_epoch_single
|
||||
@@ -1343,6 +1387,8 @@ class V2HyperTower:
|
||||
print(msg, flush=True)
|
||||
fold_logger.info(msg)
|
||||
|
||||
_prev_phase_single = phase_single
|
||||
|
||||
fold_logger.close()
|
||||
|
||||
if args.save_checkpoints:
|
||||
|
||||
File diff suppressed because one or more lines are too long
@@ -45,6 +45,13 @@ def main():
|
||||
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
|
||||
|
||||
|
||||
+1318
-35
File diff suppressed because it is too large
Load Diff
@@ -13,7 +13,7 @@ REPO_ROOT = Path(__file__).resolve().parents[3]
|
||||
if str(REPO_ROOT) not in sys.path:
|
||||
sys.path.insert(0, str(REPO_ROOT))
|
||||
|
||||
from classes.unet_segmenter import UNetSegmenter
|
||||
from classes.v2.unet_segmenter import UNetSegmenter
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
|
||||
Executable
+68
@@ -0,0 +1,68 @@
|
||||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
|
||||
# Image-only runs v2.21 (6 total):
|
||||
# No crop: binary | multiclass
|
||||
# GT crop: binary | multiclass
|
||||
# UNet crop: binary | multiclass
|
||||
#
|
||||
# Purpose: isolate the effect of ROI cropping at the single-CNN level,
|
||||
# without any MD tower contribution (bridge-mode=image_only).
|
||||
|
||||
ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/../../.." && pwd)"
|
||||
cd "$ROOT_DIR"
|
||||
|
||||
MANIFEST="manifest.csv"
|
||||
UNET_WEIGHTS="models/v2/refuge/segmentation/per_image/best.pt"
|
||||
|
||||
COMMON=(
|
||||
--epochs 40
|
||||
--n-splits 5
|
||||
--batch-size 8
|
||||
--backbone refugelike
|
||||
--tower-mode single
|
||||
--bridge-mode image_only
|
||||
--single-warmup-tower-epochs 4
|
||||
--single-warmup-fused-epochs 0
|
||||
--img-crop-manifest "$MANIFEST"
|
||||
)
|
||||
|
||||
# ── No crop ──────────────────────────────────────────────────────────────────
|
||||
|
||||
echo "[1/6] No crop — binary, image-only..."
|
||||
python3 scripts/main/v2/run_multifold_v2.py \
|
||||
"${COMMON[@]}" --eval-mode binary \
|
||||
--run-name v2.21_imgonly_binary_nocrop_40ep_5fold
|
||||
|
||||
echo "[2/6] No crop — multiclass, image-only..."
|
||||
python3 scripts/main/v2/run_multifold_v2.py \
|
||||
"${COMMON[@]}" --eval-mode multiclass \
|
||||
--run-name v2.21_imgonly_multiclass_nocrop_40ep_5fold
|
||||
|
||||
# ── GT crop ──────────────────────────────────────────────────────────────────
|
||||
|
||||
echo "[3/6] GT crop — binary, image-only..."
|
||||
python3 scripts/main/v2/run_multifold_v2.py \
|
||||
"${COMMON[@]}" --eval-mode binary --img-crop-gt \
|
||||
--run-name v2.21_imgonly_binary_gt_40ep_5fold
|
||||
|
||||
echo "[4/6] GT crop — multiclass, image-only..."
|
||||
python3 scripts/main/v2/run_multifold_v2.py \
|
||||
"${COMMON[@]}" --eval-mode multiclass --img-crop-gt \
|
||||
--run-name v2.21_imgonly_multiclass_gt_40ep_5fold
|
||||
|
||||
# ── UNet crop ────────────────────────────────────────────────────────────────
|
||||
|
||||
echo "[5/6] UNet crop — binary, image-only..."
|
||||
python3 scripts/main/v2/run_multifold_v2.py \
|
||||
"${COMMON[@]}" --eval-mode binary \
|
||||
--img-crop-weights "$UNET_WEIGHTS" \
|
||||
--run-name v2.21_imgonly_binary_unet_40ep_5fold
|
||||
|
||||
echo "[6/6] UNet crop — multiclass, image-only..."
|
||||
python3 scripts/main/v2/run_multifold_v2.py \
|
||||
"${COMMON[@]}" --eval-mode multiclass \
|
||||
--img-crop-weights "$UNET_WEIGHTS" \
|
||||
--run-name v2.21_imgonly_multiclass_unet_40ep_5fold
|
||||
|
||||
echo "Image-only v2.21 runs complete."
|
||||
Executable
+20
@@ -0,0 +1,20 @@
|
||||
#!/usr/bin/env bash
|
||||
set -e
|
||||
|
||||
echo "=== v2.3 ensemble binary nocrop ==="
|
||||
python scripts/basic_analysis/compare_hypertower_modes.py \
|
||||
--tower-modes ensemble --eval-modes binary \
|
||||
--epochs 40 --n-splits 5 \
|
||||
--backbone refugelike \
|
||||
--img-crop-manifest analysis_data/unet_manifest.csv \
|
||||
--run-name v2.3_ensemble_binary_nocrop
|
||||
|
||||
echo "=== v2.3 ensemble multiclass nocrop ==="
|
||||
python scripts/basic_analysis/compare_hypertower_modes.py \
|
||||
--tower-modes ensemble --eval-modes multiclass \
|
||||
--epochs 40 --n-splits 5 \
|
||||
--backbone refugelike \
|
||||
--img-crop-manifest analysis_data/unet_manifest.csv \
|
||||
--run-name v2.3_ensemble_multiclass_nocrop
|
||||
|
||||
echo "=== done ==="
|
||||
Executable
+24
@@ -0,0 +1,24 @@
|
||||
#!/usr/bin/env bash
|
||||
set -e
|
||||
|
||||
echo "=== v2.3 fused binary nocrop ==="
|
||||
python scripts/basic_analysis/compare_hypertower_modes.py \
|
||||
--tower-modes ensemble --eval-modes binary \
|
||||
--epochs 40 --n-splits 5 \
|
||||
--backbone refugelike \
|
||||
--img-crop-manifest analysis_data/unet_manifest.csv \
|
||||
--warmup-md-epochs 50 \
|
||||
--fused-head \
|
||||
--run-name v2.3_fused_binary_nocrop
|
||||
|
||||
echo "=== v2.3 fused multiclass nocrop ==="
|
||||
python scripts/basic_analysis/compare_hypertower_modes.py \
|
||||
--tower-modes ensemble --eval-modes multiclass \
|
||||
--epochs 40 --n-splits 5 \
|
||||
--backbone refugelike \
|
||||
--img-crop-manifest analysis_data/unet_manifest.csv \
|
||||
--warmup-md-epochs 50 \
|
||||
--fused-head \
|
||||
--run-name v2.3_fused_multiclass_nocrop
|
||||
|
||||
echo "=== done ==="
|
||||
@@ -611,84 +611,70 @@ def run_fusion_event_analysis(
|
||||
pred_i = pi.argmax(axis=1)
|
||||
pred_m = pm.argmax(axis=1)
|
||||
|
||||
corrections = (pred_f == y_true) & (pred_i != y_true) & (pred_m != y_true)
|
||||
errors = (pred_f != y_true) & (pred_i == y_true) & (pred_m == y_true)
|
||||
n_corr = corrections.sum()
|
||||
n_err = errors.sum()
|
||||
both_wrong = ((pred_i != y_true) & (pred_m != y_true)).sum()
|
||||
both_correct = ((pred_i == y_true) & (pred_m == y_true)).sum()
|
||||
# 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)
|
||||
|
||||
print(f" N={N} corrections={n_corr} errors={n_err} ratio={n_corr}/{n_err}", flush=True)
|
||||
print(f" correction rate: {n_corr}/{both_wrong} = {n_corr/max(both_wrong,1):.2%} of both-wrong cases", flush=True)
|
||||
print(f" error rate: {n_err}/{both_correct} = {n_err/max(both_correct,1):.2%} of both-correct cases", flush=True)
|
||||
f_ok = pred_f == y_true
|
||||
i_ok = pred_i == y_true
|
||||
m_ok = pred_m == y_true
|
||||
|
||||
# ---- cache intermediate hm/hi vectors for all patients ----
|
||||
model.eval()
|
||||
hm1_list, hm2_list, hi1_list, hi2_list = [], [], [], []
|
||||
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")
|
||||
if not (torch.is_tensor(x1) and torch.is_tensor(m1)):
|
||||
continue
|
||||
hi1 = model.bridge.ln_img(model.bridge.W_img(model.img_tower(x1.to(device))))
|
||||
hi2 = model.bridge.ln_img(model.bridge.W_img(model.img_tower(x2.to(device))))
|
||||
hm1 = model.bridge.ln_md(model.bridge.W_md(model.md_tower(m1.to(device))))
|
||||
hm2 = model.bridge.ln_md(model.bridge.W_md(model.md_tower(m2.to(device))))
|
||||
hi1_list.append(hi1.cpu()); hi2_list.append(hi2.cpu())
|
||||
hm1_list.append(hm1.cpu()); hm2_list.append(hm2.cpu())
|
||||
# 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
|
||||
|
||||
hi1 = torch.cat(hi1_list) # [N, fusion_dim]
|
||||
hi2 = torch.cat(hi2_list)
|
||||
hm1 = torch.cat(hm1_list) # [N, fusion_dim]
|
||||
hm2 = torch.cat(hm2_list)
|
||||
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]
|
||||
|
||||
hm1_mean = hm1.mean(dim=0, keepdim=True)
|
||||
hm2_mean = hm2.mean(dim=0, keepdim=True)
|
||||
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)
|
||||
|
||||
# ---- for each patient: compare logit[true_class] with real hm vs mean hm ----
|
||||
gains = []
|
||||
with torch.no_grad():
|
||||
for idx in range(N):
|
||||
true_cls = int(y_true[idx])
|
||||
# patient-level average of OD/OS fused vectors (SE skipped: hard to replicate outside forward)
|
||||
fused_real = (hi1[idx:idx+1] * hm1[idx:idx+1] + hi2[idx:idx+1] * hm2[idx:idx+1]) * 0.5
|
||||
fused_mean = (hi1[idx:idx+1] * hm1_mean + hi2[idx:idx+1] * hm2_mean) * 0.5
|
||||
logit_real = model.bridge.classifier_fused(fused_real.to(device))
|
||||
logit_mean = model.bridge.classifier_fused(fused_mean.to(device))
|
||||
gain = (logit_real[0, true_cls] - logit_mean[0, true_cls]).item()
|
||||
gains.append(gain)
|
||||
|
||||
gains = np.array(gains)
|
||||
|
||||
if n_corr > 0:
|
||||
corr_gains = gains[corrections]
|
||||
helped = (corr_gains > 0).sum()
|
||||
print(f"\n Fusion corrections — MD gate gain vs mean gate:", flush=True)
|
||||
print(f" mean gain = {corr_gains.mean():+.4f} median = {np.median(corr_gains):+.4f}", flush=True)
|
||||
print(f" real MD helped {helped}/{n_corr} correction patients ({helped/n_corr:.0%})", flush=True)
|
||||
|
||||
if n_err > 0:
|
||||
err_gains = gains[errors]
|
||||
print(f"\n Fusion errors — MD gate gain vs mean gate:", flush=True)
|
||||
print(f" mean gain = {err_gains.mean():+.4f} median = {np.median(err_gains):+.4f}", flush=True)
|
||||
|
||||
# ---- save CSV ----
|
||||
# ---- 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(pf[idx].max()),
|
||||
"conf_img": float(pi[idx].max()),
|
||||
"conf_md": float(pm[idx].max()),
|
||||
"is_correction": bool(corrections[idx]),
|
||||
"is_error": bool(errors[idx]),
|
||||
"md_gate_gain": float(gains[idx]),
|
||||
"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:
|
||||
@@ -697,25 +683,61 @@ def run_fusion_event_analysis(
|
||||
writer.writerows(rows)
|
||||
print(f" Saved → {csv_path}", flush=True)
|
||||
|
||||
# ---- chart ----
|
||||
fig, axes = plt.subplots(1, 2, figsize=(11, 4))
|
||||
# ---- plot ----
|
||||
fig, axes = plt.subplots(1, 3, figsize=(15, 4))
|
||||
|
||||
categories = ["corrections\n(both wrong→fused right)", "errors\n(both right→fused wrong)"]
|
||||
counts = [int(n_corr), int(n_err)]
|
||||
axes[0].bar(categories, counts, color=["#e05c5c", "#5c9ee0"], width=0.5)
|
||||
# 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})")
|
||||
for i, v in enumerate(counts):
|
||||
axes[0].text(i, v + 0.1, str(v), ha="center", fontsize=11)
|
||||
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")
|
||||
|
||||
if n_corr > 0:
|
||||
axes[1].hist(gains[corrections], bins=10, alpha=0.7, color="#e05c5c", label=f"corrections (n={n_corr})")
|
||||
if n_err > 0:
|
||||
axes[1].hist(gains[errors], bins=10, alpha=0.7, color="#5c9ee0", label=f"errors (n={n_err})")
|
||||
axes[1].axvline(0, color="black", linewidth=0.8)
|
||||
axes[1].set_xlabel("MD gate gain vs mean gate\n(logit[true class]: real − mean)")
|
||||
axes[1].set_title("Does real MD help the fused prediction?")
|
||||
axes[1].legend(fontsize=9)
|
||||
# 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)
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,652 @@
|
||||
"""
|
||||
Plot predicted-probability strip charts for V2 runs.
|
||||
|
||||
Three styles available via --style:
|
||||
|
||||
strips (default)
|
||||
X = true class, Y = P(Glaucoma), color = true class × fold shade.
|
||||
Works for binary and multiclass.
|
||||
|
||||
confidence
|
||||
X = predicted class (major) subdivided by true class (minor sub-column).
|
||||
Y = model confidence in its own prediction (P of the predicted class).
|
||||
Color = true class × fold shade.
|
||||
Makes high-confidence mistakes immediately visible.
|
||||
|
||||
triangle (multiclass only)
|
||||
Ternary / simplex plot. Each corner = 100% probability for one class.
|
||||
Every sample is a dot placed at its softmax probability vector
|
||||
(p_H, p_G, p_S) using barycentric coordinates. The centroid is maximum
|
||||
uncertainty (1/3, 1/3, 1/3). Correctly classified samples cluster near
|
||||
their true-class corner; mistakes drift toward the wrong corner.
|
||||
|
||||
Colour families (light → dark = fold 0 → fold N-1):
|
||||
Blues = Healthy / Normal eyes
|
||||
Reds = Glaucoma eyes
|
||||
Greens = Suspect eyes
|
||||
|
||||
Usage
|
||||
-----
|
||||
python scripts/output_analysis/visualizations/plot_prob_strips.py \
|
||||
--run-dir analysis_data/v2.3_single_multiclass_nocrop/multiclass/single \
|
||||
--style triangle --head fused
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import re
|
||||
from pathlib import Path
|
||||
|
||||
import matplotlib.pyplot as plt
|
||||
import matplotlib.ticker as ticker
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
|
||||
|
||||
# ── colour families ────────────────────────────────────────────────────────────
|
||||
_CLASS_FAMILIES = {
|
||||
0: ("#aac4ff", "#0a3d91"), # blues (healthy / normal)
|
||||
1: ("#ffaaaa", "#8b0000"), # reds (glaucoma)
|
||||
2: ("#aaffcc", "#1a6b3c"), # greens (suspect)
|
||||
}
|
||||
|
||||
_CLASS_LABELS = {
|
||||
0: "Healthy",
|
||||
1: "Glaucoma",
|
||||
2: "Suspect",
|
||||
}
|
||||
|
||||
_CLASS_LABELS_SHORT = {
|
||||
0: "H",
|
||||
1: "G",
|
||||
2: "S",
|
||||
}
|
||||
|
||||
|
||||
def _lerp_hex(c1: str, c2: str, t: float) -> tuple:
|
||||
def h(c): return tuple(int(c.lstrip("#")[i*2:i*2+2], 16) / 255 for i in range(3))
|
||||
r1, g1, b1 = h(c1)
|
||||
r2, g2, b2 = h(c2)
|
||||
return (r1 + t*(r2-r1), g1 + t*(g2-g1), b1 + t*(b2-b1))
|
||||
|
||||
|
||||
def _fold_colours(n_folds: int) -> dict[int, dict[int, tuple]]:
|
||||
out: dict[int, dict[int, tuple]] = {}
|
||||
for cls, (light, dark) in _CLASS_FAMILIES.items():
|
||||
out[cls] = {}
|
||||
for f in range(n_folds):
|
||||
t = f / max(n_folds - 1, 1)
|
||||
out[cls][f] = _lerp_hex(light, dark, t)
|
||||
return out
|
||||
|
||||
|
||||
def _load_folds(run_dir: Path, head: str) -> list[pd.DataFrame]:
|
||||
fold_dirs = sorted(
|
||||
[p for p in run_dir.glob("fold*") if p.is_dir() and re.search(r"\d+", p.name)],
|
||||
key=lambda p: int(re.search(r"\d+", p.name).group()),
|
||||
)
|
||||
if not fold_dirs:
|
||||
raise FileNotFoundError(f"No fold* directories found under {run_dir}")
|
||||
frames = []
|
||||
for fd in fold_dirs:
|
||||
csv = fd / "predictions_classic.csv"
|
||||
if not csv.exists():
|
||||
print(f" [warn] {csv} not found, skipping")
|
||||
continue
|
||||
df = pd.read_csv(csv)
|
||||
df["_fold"] = int(re.search(r"\d+", fd.name).group())
|
||||
frames.append(df)
|
||||
return frames
|
||||
|
||||
|
||||
def _detect_num_classes(df: pd.DataFrame, head: str) -> int:
|
||||
return len([c for c in df.columns if c.startswith(f"prob_{head}_c")])
|
||||
|
||||
|
||||
def _draw_grid_legend(ax_leg: plt.Axes, n_folds: int, classes: list[int],
|
||||
colours: dict, title: str = "True class") -> None:
|
||||
"""Rows = folds, columns = classes grid of coloured dots."""
|
||||
from matplotlib.patches import FancyBboxPatch
|
||||
ax_leg.axis("off")
|
||||
|
||||
col_xs = np.linspace(0.55, 0.88, len(classes)) if len(classes) > 1 else [0.72]
|
||||
row_ys = np.linspace(0.88, 0.05, n_folds + 1)
|
||||
header_y, dot_ys = row_ys[0], row_ys[1:]
|
||||
|
||||
for ci, cls in enumerate(classes):
|
||||
ax_leg.text(col_xs[ci], header_y, _CLASS_LABELS_SHORT.get(cls, f"C{cls}"),
|
||||
ha="center", va="bottom", fontsize=8, fontweight="bold",
|
||||
transform=ax_leg.transAxes)
|
||||
|
||||
for fi in range(n_folds):
|
||||
y = dot_ys[fi]
|
||||
ax_leg.text(0.05, y, f"Fold {fi}", ha="left", va="center", fontsize=9,
|
||||
transform=ax_leg.transAxes)
|
||||
for ci, cls in enumerate(classes):
|
||||
ax_leg.scatter([col_xs[ci]], [y], color=colours[cls][fi], s=70, zorder=3,
|
||||
transform=ax_leg.transAxes, clip_on=False)
|
||||
|
||||
ax_leg.add_patch(FancyBboxPatch((0, 0), 1, 1, boxstyle="round,pad=0.02",
|
||||
linewidth=0.8, edgecolor="#aaaaaa",
|
||||
facecolor="#f9f9f9", zorder=0,
|
||||
transform=ax_leg.transAxes))
|
||||
ax_leg.set_title(title, fontsize=9, pad=4)
|
||||
|
||||
|
||||
# ── strips style ───────────────────────────────────────────────────────────────
|
||||
|
||||
def plot_strip(run_dir: Path, head: str = "fused", out_dir: Path | None = None,
|
||||
jitter_strength: float = 0.08) -> Path:
|
||||
"""X = true class, Y = P(Glaucoma)."""
|
||||
frames = _load_folds(run_dir, head)
|
||||
all_true = sorted({int(v) for df in frames for v in df["y_true"]})
|
||||
n_folds = len(frames)
|
||||
colours = _fold_colours(n_folds)
|
||||
rng = np.random.default_rng(seed=0)
|
||||
prob_col = f"prob_{head}_c1"
|
||||
x_pos = {cls: i for i, cls in enumerate(all_true)}
|
||||
|
||||
leg_width = 0.8 + 0.55 * len(all_true)
|
||||
fig = plt.figure(figsize=(max(6, 2.2 * len(all_true) + 1), 7))
|
||||
gs = fig.add_gridspec(1, 2, width_ratios=[max(6, 2.2 * len(all_true)), leg_width],
|
||||
wspace=0.08)
|
||||
ax = fig.add_subplot(gs[0])
|
||||
ax_leg = fig.add_subplot(gs[1])
|
||||
|
||||
fig.suptitle(f"{run_dir.parent.parent.name} — P(Glaucoma) [{head}]",
|
||||
fontsize=11, y=1.01)
|
||||
|
||||
for fold_idx, df in enumerate(frames):
|
||||
for cls in all_true:
|
||||
sub = df[df["y_true"] == cls]
|
||||
if sub.empty:
|
||||
continue
|
||||
probs = sub[prob_col].values
|
||||
jitter = rng.uniform(-jitter_strength, jitter_strength, size=len(probs))
|
||||
ax.scatter(x_pos[cls] + jitter, probs, color=colours[cls][fold_idx],
|
||||
s=45, alpha=0.88, linewidths=0, zorder=3)
|
||||
|
||||
ax.set_xticks(list(x_pos.values()))
|
||||
ax.set_xticklabels([_CLASS_LABELS.get(c, f"C{c}") for c in all_true], fontsize=12)
|
||||
ax.set_xlim(-0.5, len(all_true) - 0.5)
|
||||
ax.set_ylim(-0.05, 1.05)
|
||||
ax.set_ylabel("Predicted P(Glaucoma)", fontsize=12)
|
||||
ax.set_xlabel("True class", fontsize=12)
|
||||
ax.axhline(0.5, color="grey", linestyle="--", linewidth=0.8, alpha=0.5)
|
||||
ax.grid(axis="y", linestyle=":", alpha=0.4)
|
||||
|
||||
_draw_grid_legend(ax_leg, n_folds, all_true, colours, title="Legend")
|
||||
|
||||
fig.tight_layout()
|
||||
return _save(fig, out_dir or run_dir / "plots", f"prob_strips_{head}.png")
|
||||
|
||||
|
||||
# ── confidence style ───────────────────────────────────────────────────────────
|
||||
|
||||
def plot_confidence(run_dir: Path, head: str = "fused", out_dir: Path | None = None,
|
||||
jitter_strength: float = 0.06) -> Path:
|
||||
"""
|
||||
X = predicted class (major) × true class (minor sub-column).
|
||||
Y = model confidence = P(predicted class).
|
||||
Color = true class × fold shade.
|
||||
"""
|
||||
frames = _load_folds(run_dir, head)
|
||||
num_cls = _detect_num_classes(frames[0], head)
|
||||
all_cls = list(range(num_cls))
|
||||
all_true = sorted({int(v) for df in frames for v in df["y_true"]})
|
||||
n_folds = len(frames)
|
||||
colours = _fold_colours(n_folds)
|
||||
rng = np.random.default_rng(seed=0)
|
||||
|
||||
# Sub-column spacing within each predicted-class group
|
||||
# E.g. for 3 classes: sub-offsets at -0.25, 0, +0.25
|
||||
n_sub = len(all_true)
|
||||
sub_spacing = 0.22
|
||||
sub_offsets = np.linspace(-(n_sub - 1) * sub_spacing / 2,
|
||||
(n_sub - 1) * sub_spacing / 2,
|
||||
n_sub)
|
||||
sub_off = {cls: sub_offsets[i] for i, cls in enumerate(all_true)}
|
||||
|
||||
# Major x positions for each predicted class, spaced so sub-columns don't bleed
|
||||
group_gap = sub_spacing * n_sub + 0.35
|
||||
major_x = {pc: i * group_gap for i, pc in enumerate(all_cls)}
|
||||
|
||||
leg_width = 0.8 + 0.55 * n_sub
|
||||
fig_w = max(7, group_gap * len(all_cls) * 1.8 + 1)
|
||||
fig = plt.figure(figsize=(fig_w, 7))
|
||||
gs = fig.add_gridspec(1, 2, width_ratios=[fig_w - leg_width, leg_width],
|
||||
wspace=0.08)
|
||||
ax = fig.add_subplot(gs[0])
|
||||
ax_leg = fig.add_subplot(gs[1])
|
||||
|
||||
fig.suptitle(f"{run_dir.parent.parent.name} — Prediction confidence [{head}]",
|
||||
fontsize=11, y=1.01)
|
||||
|
||||
for fold_idx, df in enumerate(frames):
|
||||
prob_cols = [f"prob_{head}_c{c}" for c in all_cls]
|
||||
pred_col = f"pred_{head}"
|
||||
for true_cls in all_true:
|
||||
sub = df[df["y_true"] == true_cls].copy()
|
||||
if sub.empty:
|
||||
continue
|
||||
for pred_cls in all_cls:
|
||||
rows = sub[sub[pred_col] == pred_cls]
|
||||
if rows.empty:
|
||||
continue
|
||||
# confidence = probability assigned to the predicted class
|
||||
conf = rows[f"prob_{head}_c{pred_cls}"].values
|
||||
x_base = major_x[pred_cls] + sub_off[true_cls]
|
||||
jitter = rng.uniform(-jitter_strength * 0.5,
|
||||
jitter_strength * 0.5, size=len(conf))
|
||||
ax.scatter(x_base + jitter, conf, color=colours[true_cls][fold_idx],
|
||||
s=45, alpha=0.88, linewidths=0, zorder=3)
|
||||
|
||||
# X-axis: major ticks with predicted-class labels, minor sub-column markers
|
||||
ax.set_xlim(-group_gap * 0.5, group_gap * len(all_cls) - group_gap * 0.5)
|
||||
ax.set_xticks([major_x[pc] for pc in all_cls])
|
||||
ax.set_xticklabels([_CLASS_LABELS.get(pc, f"C{pc}") for pc in all_cls], fontsize=12)
|
||||
|
||||
# Light vertical separators between predicted-class groups
|
||||
for i in range(1, len(all_cls)):
|
||||
sep_x = (major_x[all_cls[i-1]] + major_x[all_cls[i]]) / 2
|
||||
ax.axvline(sep_x, color="#cccccc", linewidth=1.0, zorder=1)
|
||||
|
||||
# Sub-column labels (H/G/S) just below the x-axis
|
||||
for pred_cls in all_cls:
|
||||
for true_cls in all_true:
|
||||
lbl = _CLASS_LABELS_SHORT.get(true_cls, f"C{true_cls}")
|
||||
ax.text(major_x[pred_cls] + sub_off[true_cls], -0.085, lbl,
|
||||
ha="center", va="top", fontsize=7, color="#555555",
|
||||
transform=ax.get_xaxis_transform())
|
||||
|
||||
ax.set_ylim(-0.05, 1.05)
|
||||
ax.set_ylabel("Confidence P(predicted class)", fontsize=12)
|
||||
ax.set_xlabel("Predicted class", fontsize=12, labelpad=18)
|
||||
ax.axhline(0.5, color="grey", linestyle="--", linewidth=0.8, alpha=0.5)
|
||||
ax.grid(axis="y", linestyle=":", alpha=0.4)
|
||||
|
||||
_draw_grid_legend(ax_leg, n_folds, all_true, colours, title="True class")
|
||||
|
||||
fig.tight_layout()
|
||||
return _save(fig, out_dir or run_dir / "plots", f"prob_confidence_{head}.png")
|
||||
|
||||
|
||||
# ── triangle / ternary style ───────────────────────────────────────────────────
|
||||
|
||||
# Equilateral triangle vertices in Cartesian space:
|
||||
# H (Healthy) = bottom-left (0, 0)
|
||||
# G (Glaucoma) = bottom-right (1, 0)
|
||||
# S (Suspect) = apex (0.5, sqrt(3)/2)
|
||||
_TRI_VERTICES = np.array([
|
||||
[0.0, 0.0], # class 0 – Healthy
|
||||
[1.0, 0.0], # class 1 – Glaucoma
|
||||
[0.5, np.sqrt(3) / 2], # class 2 – Suspect
|
||||
])
|
||||
|
||||
|
||||
def _bary_to_cart(probs: np.ndarray) -> np.ndarray:
|
||||
"""
|
||||
Convert Nx3 barycentric coordinates (softmax probs) to Nx2 Cartesian.
|
||||
probs rows must sum to 1.
|
||||
"""
|
||||
return probs @ _TRI_VERTICES
|
||||
|
||||
|
||||
def _draw_triangle_grid(ax: plt.Axes, levels: tuple = (0.25, 0.5, 0.75)) -> None:
|
||||
"""Draw the triangle border and iso-probability grid lines."""
|
||||
from matplotlib.patches import Polygon
|
||||
from matplotlib.lines import Line2D
|
||||
|
||||
# Outer triangle
|
||||
tri = Polygon(_TRI_VERTICES, fill=False, edgecolor="#333333", linewidth=1.5, zorder=2)
|
||||
ax.add_patch(tri)
|
||||
|
||||
# Grid lines: for each class, lines where p_class = level,
|
||||
# parallel to the opposite edge.
|
||||
for level in levels:
|
||||
for cls in range(3):
|
||||
# Points on the two edges adjacent to this vertex at distance `level`
|
||||
v0 = _TRI_VERTICES[cls]
|
||||
v1 = _TRI_VERTICES[(cls + 1) % 3]
|
||||
v2 = _TRI_VERTICES[(cls + 2) % 3]
|
||||
# A line at p_cls = level divides the triangle:
|
||||
# p1 = level * v0 + (1-level) * v1
|
||||
# p2 = level * v0 + (1-level) * v2
|
||||
p1 = level * v0 + (1 - level) * v1
|
||||
p2 = level * v0 + (1 - level) * v2
|
||||
ax.plot([p1[0], p2[0]], [p1[1], p2[1]],
|
||||
color="#cccccc", linewidth=0.6, zorder=1, linestyle="--")
|
||||
|
||||
|
||||
def plot_triangle(run_dir: Path, head: str = "fused", out_dir: Path | None = None) -> Path:
|
||||
"""
|
||||
Ternary simplex plot of softmax probabilities for a 3-class run.
|
||||
Each dot = one eye; position = (p_H, p_G, p_S) in barycentric coords.
|
||||
Color = true class × fold shade.
|
||||
"""
|
||||
frames = _load_folds(run_dir, head)
|
||||
num_cls = _detect_num_classes(frames[0], head)
|
||||
if num_cls != 3:
|
||||
raise ValueError(f"Triangle plot requires 3 classes; got {num_cls}. "
|
||||
"Use --style strips for binary.")
|
||||
|
||||
all_true = sorted({int(v) for df in frames for v in df["y_true"]})
|
||||
n_folds = len(frames)
|
||||
colours = _fold_colours(n_folds)
|
||||
|
||||
leg_width = 0.8 + 0.55 * len(all_true)
|
||||
fig = plt.figure(figsize=(8, 7))
|
||||
gs = fig.add_gridspec(1, 2, width_ratios=[8 - leg_width, leg_width], wspace=0.05)
|
||||
ax = fig.add_subplot(gs[0], aspect="equal")
|
||||
ax_leg = fig.add_subplot(gs[1])
|
||||
|
||||
fig.suptitle(f"{run_dir.parent.parent.name} — Simplex [{head}]",
|
||||
fontsize=11, y=1.01)
|
||||
|
||||
_draw_triangle_grid(ax)
|
||||
|
||||
# Vertex labels — placed just outside each corner
|
||||
offsets = [(-0.07, -0.06), (0.07, -0.06), (0.0, 0.06)]
|
||||
for cls_idx in range(3):
|
||||
lbl = _CLASS_LABELS.get(cls_idx, f"C{cls_idx}")
|
||||
vx, vy = _TRI_VERTICES[cls_idx]
|
||||
dx, dy = offsets[cls_idx]
|
||||
ax.text(vx + dx, vy + dy, lbl, ha="center", va="center",
|
||||
fontsize=12, fontweight="bold")
|
||||
|
||||
# Centroid marker
|
||||
cx, cy = _TRI_VERTICES.mean(axis=0)
|
||||
ax.scatter([cx], [cy], color="#aaaaaa", s=30, marker="+", zorder=2, linewidths=1)
|
||||
ax.text(cx + 0.02, cy - 0.04, "1/3 each", fontsize=7, color="#999999", ha="left")
|
||||
|
||||
# Data dots
|
||||
rng = np.random.default_rng(seed=0)
|
||||
for fold_idx, df in enumerate(frames):
|
||||
prob_cols = [f"prob_{head}_c{c}" for c in range(3)]
|
||||
probs_all = df[prob_cols].values # Nx3
|
||||
y_true = df["y_true"].values.astype(int)
|
||||
for true_cls in all_true:
|
||||
mask = y_true == true_cls
|
||||
if not mask.any():
|
||||
continue
|
||||
probs = probs_all[mask] # Kx3
|
||||
xy = _bary_to_cart(probs) # Kx2
|
||||
noise = rng.normal(0, 0.004, xy.shape)
|
||||
ax.scatter(xy[:, 0] + noise[:, 0],
|
||||
xy[:, 1] + noise[:, 1],
|
||||
color=colours[true_cls][fold_idx],
|
||||
s=30, alpha=0.80, linewidths=0, zorder=3)
|
||||
|
||||
ax.set_xlim(-0.18, 1.18)
|
||||
ax.set_ylim(-0.12, 1.02)
|
||||
ax.axis("off")
|
||||
|
||||
_draw_grid_legend(ax_leg, n_folds, all_true, colours, title="True class")
|
||||
|
||||
fig.tight_layout()
|
||||
return _save(fig, out_dir or run_dir / "plots", f"prob_triangle_{head}.png")
|
||||
|
||||
|
||||
# ── triangle3d style ──────────────────────────────────────────────────────────
|
||||
|
||||
# Triangle vertices for the 3D "bread-slice" view.
|
||||
# Triangles stand upright in the x-z plane; y is the depth (layer) axis.
|
||||
# H (Healthy) = bottom-left (0, 0) — back corner of the base
|
||||
# G (Glaucoma) = top-centre (0.5, √3/2) — apex (visually "at the top")
|
||||
# S (Suspect) = bottom-right (1, 0) — front corner toward the viewer
|
||||
_TRI3D_VERTS = np.array([
|
||||
[0.0, 0.0], # class 0 – Healthy (bottom-left / back)
|
||||
[0.5, np.sqrt(3) / 2], # class 1 – Glaucoma (top)
|
||||
[1.0, 0.0], # class 2 – Suspect (bottom-right / front)
|
||||
])
|
||||
|
||||
|
||||
def _bary_to_cart_3d(probs: np.ndarray) -> np.ndarray:
|
||||
"""Convert Nx3 softmax probs to Nx2 Cartesian using the 3D vertex layout."""
|
||||
return probs @ _TRI3D_VERTS
|
||||
|
||||
|
||||
def plot_triangle_3d(
|
||||
run_dir: Path,
|
||||
head: str = "fused",
|
||||
out_dir: Path | None = None,
|
||||
elev: float = 20,
|
||||
azim: float = -45,
|
||||
y_spacing: float = 0.4,
|
||||
) -> Path:
|
||||
"""
|
||||
3-D ternary "bread-slice" plot.
|
||||
|
||||
Each true class gets its own vertical triangle slice standing in the x-z
|
||||
plane, stacked along the y (depth) axis (camera at ~4:30 / SE position):
|
||||
• Healthy layer is at y=0 (front-left from 4:30 view)
|
||||
• Glaucoma layer is at y=1
|
||||
• Suspect layer is at y=2 (back-right from 4:30 view)
|
||||
|
||||
Within every slice the simplex corners are:
|
||||
• G (Glaucoma) at the top apex
|
||||
• H (Healthy) at the bottom-left (back corner of the base)
|
||||
• S (Suspect) at the bottom-right (front corner toward viewer)
|
||||
|
||||
View with elev/azim to see the slices as near-vertical planes with slight
|
||||
perspective depth.
|
||||
"""
|
||||
from mpl_toolkits.mplot3d import Axes3D # noqa: F401
|
||||
from mpl_toolkits.mplot3d.art3d import Poly3DCollection
|
||||
from matplotlib.patches import Patch
|
||||
|
||||
frames = _load_folds(run_dir, head)
|
||||
num_cls = _detect_num_classes(frames[0], head)
|
||||
if num_cls != 3:
|
||||
raise ValueError("triangle3d requires 3 classes.")
|
||||
|
||||
all_true = sorted({int(v) for df in frames for v in df["y_true"]})
|
||||
n_folds = len(frames)
|
||||
colours = _fold_colours(n_folds)
|
||||
rng = np.random.default_rng(seed=0)
|
||||
|
||||
# Layer y positions: H front-left (y=0), S back-right (y=2) from 4:30 view
|
||||
y_pos = {cls: i * y_spacing for i, cls in enumerate(all_true)}
|
||||
|
||||
# Triangle wireframe in (x, z) — loop closed
|
||||
wire_x = np.append(_TRI3D_VERTS[:, 0], _TRI3D_VERTS[0, 0])
|
||||
wire_z = np.append(_TRI3D_VERTS[:, 1], _TRI3D_VERTS[0, 1])
|
||||
|
||||
fig = plt.figure(figsize=(9, 7))
|
||||
ax = fig.add_subplot(111, projection="3d")
|
||||
ax.view_init(elev=elev, azim=azim)
|
||||
|
||||
fig.suptitle(run_dir.parent.parent.name, fontsize=11)
|
||||
|
||||
# ── pre-compute all xz points per class (needed for KDE) ─────────────────
|
||||
from scipy.stats import gaussian_kde
|
||||
import matplotlib.colors as mcolors
|
||||
|
||||
_xz_by_cls: dict[int, np.ndarray] = {}
|
||||
for df in frames:
|
||||
prob_cols_all = [f"prob_{head}_c{c}" for c in range(3)]
|
||||
p_all = df[prob_cols_all].values
|
||||
yt_all = df["y_true"].values.astype(int)
|
||||
for cls in all_true:
|
||||
m = yt_all == cls
|
||||
if m.any():
|
||||
xz_pts = _bary_to_cart_3d(p_all[m])
|
||||
_xz_by_cls[cls] = (
|
||||
np.vstack([_xz_by_cls[cls], xz_pts])
|
||||
if cls in _xz_by_cls else xz_pts
|
||||
)
|
||||
|
||||
# KDE grid setup — evaluate on a 40×40 grid masked to the triangle interior
|
||||
_G = 40
|
||||
_x_lin = np.linspace(0.0, 1.0, _G)
|
||||
_z_lin = np.linspace(0.0, np.sqrt(3) / 2, _G)
|
||||
_dx = _x_lin[1] - _x_lin[0]
|
||||
_dz = _z_lin[1] - _z_lin[0]
|
||||
_XX, _ZZ = np.meshgrid(_x_lin, _z_lin) # (_G, _G)
|
||||
|
||||
# Vectorised inside-triangle test (barycentric)
|
||||
vH, vG, vS = _TRI3D_VERTS
|
||||
_denom = (vG[1] - vS[1]) * (vH[0] - vS[0]) + (vS[0] - vG[0]) * (vH[1] - vS[1])
|
||||
def _inside(px, pz):
|
||||
la = ((vG[1] - vS[1]) * (px - vS[0]) + (vS[0] - vG[0]) * (pz - vS[1])) / _denom
|
||||
lb = ((vS[1] - vH[1]) * (px - vS[0]) + (vH[0] - vS[0]) * (pz - vS[1])) / _denom
|
||||
return (la >= 0) & (lb >= 0) & ((1 - la - lb) >= 0)
|
||||
|
||||
_inside_mask = _inside(_XX.ravel(), _ZZ.ravel()).reshape(_G, _G)
|
||||
_grid_pts = np.vstack([_XX.ravel(), _ZZ.ravel()]) # 2×(_G²)
|
||||
|
||||
# ── draw triangle wireframe + density heatmap at each y layer ─────────────
|
||||
for cls in all_true:
|
||||
y = y_pos[cls]
|
||||
base_col = colours[cls][n_folds // 2]
|
||||
rgba_base = np.array(mcolors.to_rgba(base_col))
|
||||
|
||||
# Wireframe
|
||||
ax.plot(wire_x, np.full_like(wire_x, y), wire_z,
|
||||
color=base_col, linewidth=1.2, alpha=0.65, zorder=1)
|
||||
|
||||
# ── density heatmap ────────────────────────────────────────────────
|
||||
xz_pts = _xz_by_cls.get(cls)
|
||||
if xz_pts is not None and len(xz_pts) >= 2:
|
||||
kde = gaussian_kde(xz_pts.T, bw_method="silverman")
|
||||
density = kde(_grid_pts).reshape(_G, _G)
|
||||
density[~_inside_mask] = 0.0
|
||||
inside_vals = density[_inside_mask]
|
||||
d_max = inside_vals.max()
|
||||
if d_max > 0:
|
||||
# Normalise by 95th-percentile density so a tight peak at one
|
||||
# corner doesn't wash out the rest of the triangle. Values
|
||||
# above the cap are clipped to the max alpha.
|
||||
d_ref = np.percentile(inside_vals[inside_vals > 0], 95)
|
||||
if d_ref == 0:
|
||||
d_ref = d_max
|
||||
alpha_grid = np.clip(density / d_ref, 0, 1) * 0.50
|
||||
# Build one quad per inside cell, coloured by density alpha
|
||||
quads, face_cols = [], []
|
||||
for i in range(_G):
|
||||
for j in range(_G):
|
||||
if not _inside_mask[i, j]:
|
||||
continue
|
||||
xi, zj = _x_lin[j], _z_lin[i]
|
||||
x0, x1 = xi - _dx / 2, xi + _dx / 2
|
||||
z0, z1 = zj - _dz / 2, zj + _dz / 2
|
||||
quads.append([(x0, y, z0), (x1, y, z0),
|
||||
(x1, y, z1), (x0, y, z1)])
|
||||
fc = rgba_base.copy()
|
||||
fc[3] = float(alpha_grid[i, j])
|
||||
face_cols.append(fc)
|
||||
heat = Poly3DCollection(quads, facecolors=face_cols,
|
||||
edgecolors="none", zorder=0)
|
||||
ax.add_collection3d(heat)
|
||||
|
||||
# Iso-prob grid lines
|
||||
for level in (0.25, 0.5, 0.75):
|
||||
for vi in range(3):
|
||||
v0 = _TRI3D_VERTS[vi]
|
||||
v1 = _TRI3D_VERTS[(vi + 1) % 3]
|
||||
v2 = _TRI3D_VERTS[(vi + 2) % 3]
|
||||
p1 = level * v0 + (1 - level) * v1
|
||||
p2 = level * v0 + (1 - level) * v2
|
||||
ax.plot([p1[0], p2[0]], [y, y], [p1[1], p2[1]],
|
||||
color="#cccccc", linewidth=0.4, linestyle="--",
|
||||
alpha=0.45, zorder=1)
|
||||
|
||||
# ── vertex labels just outside the front face of the merged volume ────────
|
||||
# Place labels at the H-layer plane (y=0, front-left) but pushed slightly
|
||||
# in front of the y-axis so they don't collide with the wireframe.
|
||||
lbl_info = [
|
||||
(0, -0.13, -0.08), # H: bottom-left
|
||||
(1, 0.00, 0.10), # G: top
|
||||
(2, 0.13, -0.08), # S: bottom-right
|
||||
]
|
||||
front_y = min(y_pos.values()) - 0.08
|
||||
for ci, dx, dz in lbl_info:
|
||||
vx, vz = _TRI3D_VERTS[ci]
|
||||
ax.text(vx + dx, front_y, vz + dz,
|
||||
_CLASS_LABELS.get(ci, f"C{ci}"),
|
||||
fontsize=10, fontweight="bold", ha="center", va="center",
|
||||
zorder=10)
|
||||
|
||||
# ── data dots ─────────────────────────────────────────────────────────────
|
||||
for fold_idx, df in enumerate(frames):
|
||||
prob_cols = [f"prob_{head}_c{c}" for c in range(3)]
|
||||
probs_all = df[prob_cols].values
|
||||
y_true = df["y_true"].values.astype(int)
|
||||
for true_cls in all_true:
|
||||
mask = y_true == true_cls
|
||||
if not mask.any():
|
||||
continue
|
||||
probs = probs_all[mask]
|
||||
xz = _bary_to_cart_3d(probs) # Nx2 (x, z)
|
||||
noise = rng.normal(0, 0.004, xz.shape)
|
||||
y_vals = np.full(mask.sum(), y_pos[true_cls])
|
||||
ax.scatter(xz[:, 0] + noise[:, 0],
|
||||
y_vals,
|
||||
xz[:, 1] + noise[:, 1],
|
||||
color=colours[true_cls][fold_idx],
|
||||
s=22, alpha=0.82, linewidths=0,
|
||||
depthshade=False, zorder=5)
|
||||
|
||||
ax.set_xlim(-0.15, 1.15)
|
||||
ax.set_ylim(-0.3, max(y_pos.values()) + 0.3)
|
||||
ax.set_zlim(-0.1, np.sqrt(3) / 2 + 0.1)
|
||||
ax.set_xlabel("")
|
||||
ax.set_ylabel("")
|
||||
ax.set_zlabel("")
|
||||
ax.set_xticks([])
|
||||
ax.set_yticks([])
|
||||
ax.set_zticks([])
|
||||
ax.xaxis.pane.fill = False
|
||||
ax.yaxis.pane.fill = False
|
||||
ax.zaxis.pane.fill = False
|
||||
ax.grid(False)
|
||||
|
||||
# ── legend: same grid style as 2D plots, placed as an inset axes ─────────
|
||||
ax_leg = fig.add_axes([0.68, 0.65, 0.28, 0.30])
|
||||
_draw_grid_legend(ax_leg, n_folds, all_true, colours, title="True class")
|
||||
|
||||
fig.tight_layout()
|
||||
return _save(fig, out_dir or run_dir / "plots", f"prob_triangle3d_{head}.png")
|
||||
|
||||
|
||||
# ── shared save helper ─────────────────────────────────────────────────────────
|
||||
|
||||
def _save(fig: plt.Figure, out_dir: Path, filename: str) -> Path:
|
||||
out_dir = Path(out_dir)
|
||||
out_dir.mkdir(parents=True, exist_ok=True)
|
||||
out_path = out_dir / filename
|
||||
fig.savefig(out_path, dpi=150, bbox_inches="tight")
|
||||
plt.close(fig)
|
||||
print(f"Saved: {out_path}")
|
||||
return out_path
|
||||
|
||||
|
||||
# ── CLI ────────────────────────────────────────────────────────────────────────
|
||||
|
||||
def main():
|
||||
ap = argparse.ArgumentParser(description=__doc__,
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter)
|
||||
ap.add_argument("--run-dir", required=True)
|
||||
ap.add_argument("--head", default="fused", choices=["fused", "img", "md"])
|
||||
ap.add_argument("--style", default="strips",
|
||||
choices=["strips", "confidence", "triangle", "triangle3d"],
|
||||
help="strips: X=true class, Y=P(Glaucoma). "
|
||||
"confidence: X=predicted class × true sub-column, Y=confidence. "
|
||||
"triangle: ternary simplex plot (multiclass only).")
|
||||
ap.add_argument("--out", default=None)
|
||||
args = ap.parse_args()
|
||||
|
||||
rd = Path(args.run_dir)
|
||||
od = Path(args.out) if args.out else None
|
||||
if args.style == "confidence":
|
||||
plot_confidence(rd, head=args.head, out_dir=od)
|
||||
elif args.style == "triangle":
|
||||
plot_triangle(rd, head=args.head, out_dir=od)
|
||||
elif args.style == "triangle3d":
|
||||
plot_triangle_3d(rd, head=args.head, out_dir=od)
|
||||
else:
|
||||
plot_strip(rd, head=args.head, out_dir=od)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user