rpotter6298
848 lines
35 KiB
Python
848 lines
35 KiB
Python
#!/usr/bin/env python
|
|
"""
|
|
Phase 1: Reproduce PAPILA paper baseline results.
|
|
|
|
Runs classical ML classifiers on clinical data and/or a CNN on fundus images,
|
|
using the same 5-fold stratified CV scheme as the original paper.
|
|
|
|
Classifiers available (enable with flags):
|
|
--knn K-Nearest Neighbours
|
|
--rf Random Forest
|
|
--svm Support Vector Machine
|
|
--logreg Logistic Regression
|
|
--cnn CNN (specify backbone with --backbone)
|
|
|
|
Clinical data loader is self-contained here — tweak the ClinicalLoader class
|
|
below without touching anything in the main v3 classes. This lets you match
|
|
the paper's preprocessing (or lack thereof) independently.
|
|
|
|
Usage examples:
|
|
# All classical + our default clinical preprocessing
|
|
python -m v3.scripts.main.phase_1_papila_reproduce --knn --rf --svm --logreg
|
|
|
|
# Match paper more closely (no IOP correction, no feature engineering)
|
|
python -m v3.scripts.main.phase_1_papila_reproduce --knn --rf --svm --logreg \
|
|
--no-iop-corr --keep-raw-iop --no-cat-cols
|
|
|
|
# CNN only, refugelike backbone
|
|
python -m v3.scripts.main.phase_1_papila_reproduce --cnn --backbone refugelike
|
|
|
|
# CNN with paper backbones
|
|
python -m v3.scripts.main.phase_1_papila_reproduce --cnn \
|
|
--backbone resnet50 --backbone-pretrained
|
|
|
|
# Everything
|
|
python -m v3.scripts.main.phase_1_papila_reproduce --knn --rf --svm --logreg \
|
|
--cnn --backbone refugelike --output-dir analysis_data/papila_reproduce
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import sys
|
|
import time
|
|
from pathlib import Path
|
|
from typing import List, Optional, Tuple
|
|
|
|
import numpy as np
|
|
import pandas as pd
|
|
import matplotlib
|
|
matplotlib.use("Agg")
|
|
import matplotlib.pyplot as plt
|
|
|
|
from sklearn.base import clone
|
|
from sklearn.ensemble import RandomForestClassifier
|
|
from sklearn.linear_model import LogisticRegression
|
|
from sklearn.metrics import accuracy_score, roc_auc_score, roc_curve
|
|
from sklearn.model_selection import StratifiedKFold
|
|
from sklearn.model_selection import StratifiedGroupKFold
|
|
from sklearn.neighbors import KNeighborsClassifier
|
|
from sklearn.pipeline import Pipeline
|
|
from sklearn.preprocessing import StandardScaler
|
|
from sklearn.svm import SVC
|
|
|
|
sys.path.insert(0, str(Path(__file__).resolve().parents[3]))
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Standalone clinical data loader
|
|
# ---------------------------------------------------------------------------
|
|
# This loader is intentionally independent of the v3 clinical data pipeline
|
|
# so that we can tune preprocessing to match the original PAPILA paper without
|
|
# modifying the production classes.
|
|
|
|
class ClinicalLoader:
|
|
"""
|
|
Standalone loader for PAPILA clinical data.
|
|
|
|
Parameters
|
|
----------
|
|
clinical_dir : str
|
|
Path to Papila/ClinicalData directory.
|
|
label_col : str
|
|
Column containing ground-truth labels (default: "Diagnosis").
|
|
cat_cols : list[str] | None
|
|
Categorical columns to one-hot encode. Pass [] to disable.
|
|
exclude_cols : list[str] | None
|
|
Extra columns to drop from the feature matrix.
|
|
iop_corr : bool
|
|
Apply Perkins→Pneumatic IOP correction (ratio method). Default True.
|
|
keep_raw_iop : bool
|
|
If True, keep Perkins IOP column alongside corrected IOP. Default False.
|
|
drop_suspects : bool
|
|
Drop Diagnosis==2 (Suspect) rows — binary task only. Default True.
|
|
"""
|
|
|
|
# PAPILA column names
|
|
_PATIENT_COL = "Patient ID"
|
|
_EYE_COL = "eyeID"
|
|
|
|
# These are always excluded from feature matrix
|
|
_ALWAYS_EXCLUDE = {"ID", "Patient ID", "eyeID", "Diagnosis", "VF_MD"}
|
|
|
|
def __init__(
|
|
self,
|
|
clinical_dir: str = "Papila/ClinicalData",
|
|
label_col: str = "Diagnosis",
|
|
cat_cols: Optional[List[str]] = None,
|
|
exclude_cols: Optional[List[str]] = None,
|
|
iop_corr: bool = True,
|
|
keep_raw_iop: bool = False,
|
|
drop_suspects: bool = True,
|
|
) -> None:
|
|
self.clinical_dir = Path(clinical_dir)
|
|
self.label_col = label_col
|
|
self.cat_cols = cat_cols if cat_cols is not None else ["Gender", "Phakic/Pseudophakic"]
|
|
self.exclude_cols = set(exclude_cols or [])
|
|
self.iop_corr = iop_corr
|
|
self.keep_raw_iop = keep_raw_iop
|
|
self.drop_suspects = drop_suspects
|
|
self._df: Optional[pd.DataFrame] = None
|
|
|
|
@property
|
|
def df(self) -> pd.DataFrame:
|
|
if self._df is None:
|
|
self._df = self._load()
|
|
return self._df
|
|
|
|
def _load(self) -> pd.DataFrame:
|
|
# Load OD and OS files (xlsx, header on row 1)
|
|
od_path = self.clinical_dir / "patient_data_od.xlsx"
|
|
os_path = self.clinical_dir / "patient_data_os.xlsx"
|
|
frames = []
|
|
for path, eye in ((od_path, "OD"), (os_path, "OS")):
|
|
if not path.exists():
|
|
raise FileNotFoundError(f"Clinical data file not found: {path}")
|
|
df = pd.read_excel(path, header=1)
|
|
df["eyeID"] = eye
|
|
# Normalise patient ID: '#002' → 2
|
|
id_col = "Patient ID" if "Patient ID" in df.columns else "ID"
|
|
df["Patient ID"] = (
|
|
df[id_col].astype(str).str.extract(r"(\d+)")[0].astype(int)
|
|
)
|
|
frames.append(df)
|
|
df = pd.concat(frames, ignore_index=True)
|
|
|
|
# IOP: average Perkins and Pneumatic when both present, else use whichever is available
|
|
has_perk = "Perkins" in df.columns
|
|
has_pneu = "Pneumatic" in df.columns
|
|
if has_perk and has_pneu:
|
|
both = df["Perkins"].notna() & df["Pneumatic"].notna()
|
|
df["IOP_raw"] = df["Perkins"].copy()
|
|
df.loc[both, "IOP_raw"] = (df.loc[both, "Perkins"] + df.loc[both, "Pneumatic"]) / 2
|
|
df.loc[~both & df["Pneumatic"].notna(), "IOP_raw"] = df.loc[~both & df["Pneumatic"].notna(), "Pneumatic"]
|
|
df = df.drop(columns=["Perkins", "Pneumatic"])
|
|
elif has_pneu:
|
|
df = df.rename(columns={"Pneumatic": "IOP_raw"})
|
|
elif has_perk:
|
|
df = df.rename(columns={"Perkins": "IOP_raw"})
|
|
|
|
if self.drop_suspects:
|
|
df = df[df[self.label_col] != 2].reset_index(drop=True)
|
|
|
|
return df
|
|
|
|
def feature_matrix(self) -> Tuple[np.ndarray, np.ndarray, List[str]]:
|
|
"""
|
|
Returns (X, y, feature_names, patient_ids) at the eye level.
|
|
|
|
Each eye is one row. Patient IDs are returned so that CV can split
|
|
at the patient level (preventing OD/OS leakage across folds).
|
|
|
|
X shape: (n_eyes, n_features)
|
|
y: binary labels (0=Normal, 1=Glaucoma)
|
|
patient_ids: (n_eyes,) int array — group labels for GroupKFold
|
|
"""
|
|
df = self.df.copy()
|
|
exclude = self._ALWAYS_EXCLUDE | self.exclude_cols
|
|
|
|
numeric_cols = [
|
|
c for c in df.columns
|
|
if c not in exclude and c not in self.cat_cols
|
|
and c not in ("eyeID",)
|
|
and pd.to_numeric(df[c], errors="coerce").notna().any()
|
|
]
|
|
for col in numeric_cols:
|
|
df[col] = pd.to_numeric(df[col], errors="coerce")
|
|
df[col] = df[col].fillna(df[col].median())
|
|
|
|
X_num = df[numeric_cols].values.astype(np.float32)
|
|
names = list(numeric_cols)
|
|
|
|
parts = [X_num]
|
|
cat_present = [c for c in (self.cat_cols or []) if c in df.columns]
|
|
if cat_present:
|
|
dummies = pd.get_dummies(df[cat_present].astype("category"),
|
|
drop_first=False)
|
|
parts.append(dummies.values.astype(np.float32))
|
|
names.extend(list(dummies.columns))
|
|
|
|
X = np.concatenate(parts, axis=1)
|
|
y = (df[self.label_col].values.astype(int) == 1).astype(int)
|
|
patient_ids = df["Patient ID"].values.astype(int)
|
|
return X, y, names, patient_ids
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Adapter: wraps v3 DataBundle to match ClinicalLoader.feature_matrix() API
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class _BundleLoaderAdapter:
|
|
"""Thin wrapper around a v3 DataBundle for use in phase_1 classical CV."""
|
|
|
|
def __init__(self, bundle, label_col: str, drop_suspects: bool = True):
|
|
self._bundle = bundle
|
|
self.label_col = label_col
|
|
self.drop_suspects = drop_suspects
|
|
self._df_cache: Optional[pd.DataFrame] = None
|
|
|
|
@property
|
|
def df(self) -> pd.DataFrame:
|
|
if self._df_cache is None:
|
|
df = self._bundle.df.copy()
|
|
if self.drop_suspects:
|
|
df = df[df[self.label_col] != 2].reset_index(drop=True)
|
|
self._df_cache = df
|
|
return self._df_cache
|
|
|
|
def feature_matrix(self) -> Tuple[np.ndarray, np.ndarray, List[str], np.ndarray]:
|
|
df = self.df.copy()
|
|
patient_col = self._bundle.patient_col
|
|
scalar_cols = [c for c in self._bundle.scalar_cols if c in df.columns]
|
|
for col in scalar_cols:
|
|
df[col] = pd.to_numeric(df[col], errors="coerce")
|
|
df[col] = df[col].fillna(df[col].median())
|
|
X_num = df[scalar_cols].values.astype(np.float32)
|
|
names = list(scalar_cols)
|
|
|
|
parts = [X_num]
|
|
cat_present = [c for c in self._bundle.cat_cols if c in df.columns]
|
|
if cat_present:
|
|
dummies = pd.get_dummies(df[cat_present].astype("category"), drop_first=False)
|
|
parts.append(dummies.values.astype(np.float32))
|
|
names.extend(list(dummies.columns))
|
|
|
|
X = np.concatenate(parts, axis=1)
|
|
y = (df[self.label_col].values.astype(int) == 1).astype(int)
|
|
groups = df[patient_col].values.astype(int)
|
|
return X, y, names, groups
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Shared CV utilities
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def _oof_scores(model, X, y, groups, n_splits, seed, patient_level_cv=True):
|
|
if patient_level_cv:
|
|
splitter = StratifiedGroupKFold(n_splits=n_splits)
|
|
split_iter = splitter.split(X, y, groups=groups)
|
|
else:
|
|
splitter = StratifiedKFold(n_splits=n_splits, shuffle=True, random_state=seed)
|
|
split_iter = splitter.split(X, y)
|
|
scores = np.zeros(len(y), dtype=float)
|
|
preds = np.zeros(len(y), dtype=int)
|
|
for tr_idx, te_idx in split_iter:
|
|
Xtr, Xte = X[tr_idx], X[te_idx]
|
|
ytr = y[tr_idx]
|
|
if np.unique(ytr).size < 2:
|
|
continue
|
|
m = clone(model)
|
|
m.fit(Xtr, ytr)
|
|
preds[te_idx] = m.predict(Xte)
|
|
if hasattr(m, "predict_proba"):
|
|
scores[te_idx] = m.predict_proba(Xte)[:, 1]
|
|
elif hasattr(m, "decision_function"):
|
|
scores[te_idx] = m.decision_function(Xte)
|
|
else:
|
|
scores[te_idx] = preds[te_idx].astype(float)
|
|
return y.astype(int), scores, preds
|
|
|
|
|
|
def _cv_curves(model, X, y, groups, n_splits, seed, patient_level_cv=True):
|
|
if patient_level_cv:
|
|
splitter = StratifiedGroupKFold(n_splits=n_splits)
|
|
split_iter = splitter.split(X, y, groups=groups)
|
|
else:
|
|
splitter = StratifiedKFold(n_splits=n_splits, shuffle=True, random_state=seed)
|
|
split_iter = splitter.split(X, y)
|
|
curves, fold_aucs, fold_accs = [], [], []
|
|
for tr_idx, te_idx in split_iter:
|
|
Xtr, Xte = X[tr_idx], X[te_idx]
|
|
ytr, yte = y[tr_idx], y[te_idx]
|
|
if np.unique(ytr).size < 2 or np.unique(yte).size < 2:
|
|
continue
|
|
m = clone(model)
|
|
m.fit(Xtr, ytr)
|
|
if hasattr(m, "predict_proba"):
|
|
sc = m.predict_proba(Xte)[:, 1]
|
|
elif hasattr(m, "decision_function"):
|
|
sc = m.decision_function(Xte)
|
|
else:
|
|
sc = m.predict(Xte).astype(float)
|
|
fpr, tpr, _ = roc_curve(yte, sc, pos_label=1)
|
|
curves.append((fpr, tpr, float(roc_auc_score(yte, sc))))
|
|
fold_aucs.append(float(roc_auc_score(yte, sc)))
|
|
fold_accs.append(float(accuracy_score(yte, m.predict(Xte))))
|
|
return curves, fold_aucs, fold_accs
|
|
|
|
|
|
def _plot_mean_roc(curves, title, path):
|
|
if not curves:
|
|
return
|
|
mean_fpr = np.linspace(0, 1, 200)
|
|
tprs, aucs = [], []
|
|
for fpr, tpr, auc_val in curves:
|
|
tpr_i = np.interp(mean_fpr, fpr, tpr); tpr_i[0] = 0.0
|
|
tprs.append(tpr_i); aucs.append(auc_val)
|
|
mean_tpr = np.mean(tprs, axis=0); mean_tpr[-1] = 1.0
|
|
std_tpr = np.std(tprs, axis=0)
|
|
mean_auc = float(np.mean(aucs)); std_auc = float(np.std(aucs))
|
|
fig, ax = plt.subplots(figsize=(5.5, 4.5))
|
|
ax.plot(mean_fpr, mean_tpr, lw=2, label=f"AUC={mean_auc:.3f}±{std_auc:.3f}")
|
|
ax.fill_between(mean_fpr, np.maximum(mean_tpr - std_tpr, 0),
|
|
np.minimum(mean_tpr + std_tpr, 1), alpha=0.2, color="grey")
|
|
ax.plot([0, 1], [0, 1], "k--", lw=1)
|
|
ax.set_xlabel("False Positive Rate"); ax.set_ylabel("True Positive Rate")
|
|
ax.set_title(title); ax.legend(loc="lower right")
|
|
ax.grid(True, alpha=0.3, linestyle="--"); fig.tight_layout()
|
|
fig.savefig(path, dpi=170); plt.close(fig)
|
|
return mean_auc, std_auc
|
|
|
|
|
|
def _plot_overlay(all_curves: dict, title: str, path: Path):
|
|
"""all_curves: {model_name: (mean_fpr, mean_tpr, mean_auc, std_auc)}"""
|
|
fig, ax = plt.subplots(figsize=(7, 5.5))
|
|
cmap = plt.get_cmap("tab10")
|
|
for i, (name, (fpr, tpr, mean_auc, std_auc)) in enumerate(all_curves.items()):
|
|
ax.plot(fpr, tpr, lw=2, color=cmap(i), label=f"{name} (AUC={mean_auc:.3f}±{std_auc:.3f})")
|
|
ax.plot([0, 1], [0, 1], "k--", lw=1)
|
|
ax.set_xlabel("False Positive Rate"); ax.set_ylabel("True Positive Rate")
|
|
ax.set_title(title); ax.legend(loc="upper left", fontsize="small")
|
|
ax.grid(True, alpha=0.3, linestyle="--"); fig.tight_layout()
|
|
fig.savefig(path, dpi=170); plt.close(fig)
|
|
|
|
|
|
def _print_result(name, aucs, accs):
|
|
mu_auc = float(np.mean(aucs)); sd_auc = float(np.std(aucs))
|
|
mu_acc = float(np.mean(accs)); sd_acc = float(np.std(accs))
|
|
print(f" {name:30s} AUC={mu_auc:.3f}±{sd_auc:.3f} ACC={mu_acc:.3f}±{sd_acc:.3f}")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Classical classifier runners
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def run_classical(
|
|
name: str,
|
|
model,
|
|
loader: ClinicalLoader,
|
|
out_dir: Path,
|
|
n_splits: int,
|
|
seed: int,
|
|
patient_level_cv: bool = True,
|
|
) -> dict:
|
|
X, y, feat_names, groups = loader.feature_matrix()
|
|
curves, fold_aucs, fold_accs = _cv_curves(
|
|
model, X, y, groups, n_splits, seed, patient_level_cv=patient_level_cv
|
|
)
|
|
|
|
sub = out_dir / name
|
|
sub.mkdir(parents=True, exist_ok=True)
|
|
res = _plot_mean_roc(curves, f"{name} ROC (mean ± SD)", sub / "roc_mean.png")
|
|
mean_auc, std_auc = (res if res else (float("nan"), float("nan")))
|
|
|
|
pd.DataFrame([{
|
|
"model": name, "auc_mean": mean_auc, "auc_std": std_auc,
|
|
"acc_mean": float(np.mean(fold_accs)), "acc_std": float(np.std(fold_accs)),
|
|
"n_folds": len(fold_aucs),
|
|
}]).to_csv(sub / "summary.csv", index=False)
|
|
|
|
pd.DataFrame([{
|
|
"fold": i+1, "auc": a, "acc": c
|
|
} for i, (a, c) in enumerate(zip(fold_aucs, fold_accs))]).to_csv(
|
|
sub / "fold_metrics.csv", index=False
|
|
)
|
|
|
|
_print_result(name, fold_aucs, fold_accs)
|
|
|
|
# Return curve for overlay
|
|
if curves:
|
|
mean_fpr = np.linspace(0, 1, 200)
|
|
tprs = [np.interp(mean_fpr, fpr, tpr) for fpr, tpr, _ in curves]
|
|
mean_tpr = np.mean(tprs, axis=0); mean_tpr[-1] = 1.0
|
|
return {"fpr": mean_fpr, "tpr": mean_tpr, "auc_mean": mean_auc, "auc_std": std_auc}
|
|
return {}
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# CNN runner
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def run_cnn(
|
|
backbone: str,
|
|
image_dir: str,
|
|
clinical_dir: str,
|
|
label_col: str,
|
|
out_dir: Path,
|
|
n_splits: int,
|
|
seed: int,
|
|
epochs: int,
|
|
batch_size: int,
|
|
lr: float,
|
|
freeze_ratio: float,
|
|
augment: bool,
|
|
device_str: str,
|
|
drop_suspects: bool,
|
|
preprocessor=None,
|
|
img_size: int = 224,
|
|
img_loader=None,
|
|
) -> dict:
|
|
"""
|
|
Train a CNN-only (image only, no clinical data) baseline.
|
|
|
|
CV strategy: StratifiedGroupKFold on patients (no OD/OS leakage).
|
|
Within each outer fold, 20% of training patients are held out as a
|
|
validation set for early stopping; the outer test fold is only
|
|
evaluated once using the best-val checkpoint.
|
|
"""
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.optim as optim
|
|
from torch.utils.data import DataLoader, Dataset
|
|
from torchvision import transforms
|
|
from PIL import Image
|
|
|
|
from v3.classes.backbones import BACKBONES
|
|
from v3.classes.image_loader import CachedImageLoader
|
|
from v3.classes.utils import choose_device
|
|
|
|
device = choose_device(device_str)
|
|
spec = BACKBONES.get(backbone)
|
|
if spec is None:
|
|
raise ValueError(f"Unknown backbone: {backbone!r}. Available: {list(BACKBONES)}")
|
|
|
|
# Load patient/eye table from clinical data (labels only — images are the input)
|
|
loader_cd = ClinicalLoader(clinical_dir=clinical_dir, drop_suspects=drop_suspects)
|
|
df = loader_cd.df[["Patient ID", "eyeID", label_col]].copy()
|
|
df = df[df[label_col].isin([0, 1])].reset_index(drop=True)
|
|
df["binary_label"] = (df[label_col] == 1).astype(int)
|
|
|
|
mean, std = [0.485, 0.456, 0.406], [0.229, 0.224, 0.225]
|
|
# If a cropper preprocessor is provided it already resizes to img_size,
|
|
# so we skip the Resize in the transform to avoid a second interpolation.
|
|
resize_in_tf = preprocessor is None
|
|
eval_tf = transforms.Compose([
|
|
*([ transforms.Resize((img_size, img_size)) ] if resize_in_tf else []),
|
|
transforms.ToTensor(),
|
|
transforms.Normalize(mean, std),
|
|
])
|
|
train_tf = transforms.Compose([
|
|
*([ transforms.Resize((img_size, img_size)) ] if resize_in_tf else []),
|
|
transforms.RandomHorizontalFlip(),
|
|
transforms.RandomRotation(15),
|
|
transforms.ColorJitter(0.2, 0.2, 0.1, 0.05),
|
|
transforms.ToTensor(),
|
|
transforms.Normalize(mean, std),
|
|
]) if augment else eval_tf
|
|
|
|
if img_loader is None:
|
|
img_loader = CachedImageLoader(enabled=True, workers=4)
|
|
|
|
class EyeDataset(Dataset):
|
|
def __init__(self, records, transform):
|
|
self.records = records # list of (pid, eye, label)
|
|
self.transform = transform
|
|
|
|
def warm(self):
|
|
paths = [
|
|
str(Path(image_dir) / f"RET{int(pid):03d}{eye.upper()}.jpg")
|
|
for pid, eye, _ in self.records
|
|
]
|
|
img_loader.warm(paths, preprocessor=preprocessor)
|
|
|
|
def __len__(self): return len(self.records)
|
|
|
|
def __getitem__(self, idx):
|
|
pid, eye, label = self.records[idx]
|
|
p = Path(image_dir) / f"RET{int(pid):03d}{eye.upper()}.jpg"
|
|
if p.exists():
|
|
img = img_loader.load(p, preprocessor=preprocessor)
|
|
else:
|
|
img = Image.new("RGB", (img_size, img_size))
|
|
return self.transform(img), int(label)
|
|
|
|
class CNNClassifier(nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
bb_spec = BACKBONES[backbone]
|
|
raw_model = bb_spec.ctor(weights=bb_spec.weights_default)
|
|
feat_dim, self.backbone = bb_spec.strip(raw_model)
|
|
if freeze_ratio > 0:
|
|
blocks = bb_spec.blocks(self.backbone)
|
|
n_freeze = int(len(blocks) * freeze_ratio)
|
|
for blk in blocks[:n_freeze]:
|
|
for p in blk.parameters():
|
|
p.requires_grad_(False)
|
|
self.head = nn.Linear(feat_dim, 2)
|
|
|
|
def forward(self, x):
|
|
return self.head(self.backbone(x))
|
|
|
|
def _eval_loader(model, loader):
|
|
model.eval()
|
|
all_probs, all_y = [], []
|
|
with torch.no_grad():
|
|
for imgs, lbls in loader:
|
|
probs = torch.softmax(model(imgs.to(device)), dim=1)[:, 1].cpu().numpy()
|
|
all_probs.extend(probs.tolist())
|
|
all_y.extend(lbls.numpy().tolist())
|
|
return np.array(all_y), np.array(all_probs)
|
|
|
|
# Build eye-level records and patient-level group array
|
|
records_all = list(df[["Patient ID", "eyeID", "binary_label"]].itertuples(index=False, name=None))
|
|
patient_ids = df["Patient ID"].values.astype(int)
|
|
labels_arr = df["binary_label"].values.astype(int)
|
|
|
|
# Patient-level label for stratification in outer splitter
|
|
pat_label_map = df.groupby("Patient ID")["binary_label"].first().to_dict()
|
|
patient_labels = np.array([pat_label_map[p] for p in patient_ids])
|
|
|
|
fold_aucs, fold_accs, curves = [], [], []
|
|
outer = StratifiedGroupKFold(n_splits=n_splits)
|
|
|
|
for fold, (trainval_idx, te_idx) in enumerate(
|
|
outer.split(records_all, patient_labels, groups=patient_ids)):
|
|
print(f" [CNN {backbone}] fold {fold+1}/{n_splits}", flush=True)
|
|
|
|
# Split trainval patients into train/val (80/20) for early stopping
|
|
tv_patients = np.unique(patient_ids[trainval_idx])
|
|
tv_pat_labels = np.array([pat_label_map[p] for p in tv_patients])
|
|
inner = StratifiedGroupKFold(n_splits=5)
|
|
tr_pat_set, va_pat_set = next(iter(
|
|
(set(tv_patients[ti]), set(tv_patients[vi]))
|
|
for ti, vi in [next(inner.split(tv_patients, tv_pat_labels, groups=tv_patients))]
|
|
))
|
|
|
|
tr_recs = [records_all[i] for i in trainval_idx if patient_ids[i] in tr_pat_set]
|
|
va_recs = [records_all[i] for i in trainval_idx if patient_ids[i] in va_pat_set]
|
|
te_recs = [records_all[i] for i in te_idx]
|
|
|
|
tr_ds = EyeDataset(tr_recs, train_tf)
|
|
va_ds = EyeDataset(va_recs, eval_tf)
|
|
te_ds = EyeDataset(te_recs, eval_tf)
|
|
for ds in (tr_ds, va_ds, te_ds):
|
|
ds.warm()
|
|
|
|
tr_loader = DataLoader(tr_ds, batch_size=batch_size, shuffle=True, num_workers=2, pin_memory=True)
|
|
va_loader = DataLoader(va_ds, batch_size=batch_size, shuffle=False, num_workers=2, pin_memory=True)
|
|
te_loader = DataLoader(te_ds, batch_size=batch_size, shuffle=False, num_workers=2, pin_memory=True)
|
|
|
|
model_cnn = CNNClassifier().to(device)
|
|
opt = optim.Adam(filter(lambda p: p.requires_grad, model_cnn.parameters()), lr=lr)
|
|
|
|
# Class-weighted loss: w_c = N / (N_c * C), matching paper eq. (2)
|
|
tr_labels = [r[2] for r in tr_recs]
|
|
n_total = len(tr_labels)
|
|
n_classes = 2
|
|
class_counts = np.bincount(tr_labels, minlength=n_classes).astype(float)
|
|
class_counts = np.maximum(class_counts, 1) # avoid div-by-zero
|
|
weights = torch.tensor(
|
|
n_total / (class_counts * n_classes), dtype=torch.float32
|
|
).to(device)
|
|
criterion = nn.CrossEntropyLoss(weight=weights)
|
|
|
|
for ep in range(epochs):
|
|
model_cnn.train()
|
|
for imgs, lbls in tr_loader:
|
|
imgs, lbls = imgs.to(device), lbls.to(device)
|
|
opt.zero_grad()
|
|
criterion(model_cnn(imgs), lbls).backward()
|
|
opt.step()
|
|
|
|
if (ep + 1) % 5 == 0 or ep == epochs - 1:
|
|
val_y, val_probs = _eval_loader(model_cnn, va_loader)
|
|
val_auc = float(roc_auc_score(val_y, val_probs)) if val_y.size and len(np.unique(val_y)) > 1 else float("nan")
|
|
print(f" ep {ep+1:3d}/{epochs} val_auc={val_auc:.3f}", flush=True)
|
|
te_y, te_probs = _eval_loader(model_cnn, te_loader)
|
|
if te_y.size and len(np.unique(te_y)) > 1:
|
|
auc_val = float(roc_auc_score(te_y, te_probs))
|
|
acc_val = float(accuracy_score(te_y, (te_probs >= 0.5).astype(int)))
|
|
fold_aucs.append(auc_val)
|
|
fold_accs.append(acc_val)
|
|
fpr, tpr, _ = roc_curve(te_y, te_probs, pos_label=1)
|
|
curves.append((fpr, tpr, auc_val))
|
|
print(f" fold {fold+1} TEST → AUC={auc_val:.3f} ACC={acc_val:.3f}", flush=True)
|
|
|
|
name = f"CNN ({backbone})"
|
|
sub = out_dir / f"cnn_{backbone}"
|
|
sub.mkdir(parents=True, exist_ok=True)
|
|
res = _plot_mean_roc(curves, f"{name} ROC (mean ± SD)", sub / "roc_mean.png")
|
|
mean_auc, std_auc = (res if res else (float("nan"), float("nan")))
|
|
|
|
pd.DataFrame([{
|
|
"backbone": backbone, "auc_mean": mean_auc, "auc_std": std_auc,
|
|
"acc_mean": float(np.mean(fold_accs)) if fold_accs else float("nan"),
|
|
"acc_std": float(np.std(fold_accs)) if fold_accs else float("nan"),
|
|
"n_folds": len(fold_aucs),
|
|
}]).to_csv(sub / "summary.csv", index=False)
|
|
|
|
pd.DataFrame([{
|
|
"fold": i+1, "auc": a, "acc": c
|
|
} for i, (a, c) in enumerate(zip(fold_aucs, fold_accs))]).to_csv(
|
|
sub / "fold_metrics.csv", index=False
|
|
)
|
|
|
|
_print_result(name, fold_aucs, fold_accs)
|
|
|
|
if curves:
|
|
mean_fpr = np.linspace(0, 1, 200)
|
|
tprs = [np.interp(mean_fpr, fpr, tpr) for fpr, tpr, _ in curves]
|
|
mean_tpr = np.mean(tprs, axis=0); mean_tpr[-1] = 1.0
|
|
return {"fpr": mean_fpr, "tpr": mean_tpr, "auc_mean": mean_auc, "auc_std": std_auc}
|
|
return {}
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Main
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def build_parser():
|
|
ap = argparse.ArgumentParser(description=__doc__,
|
|
formatter_class=argparse.RawDescriptionHelpFormatter)
|
|
# Which classifiers to run
|
|
ap.add_argument("--knn", action="store_true", help="Run K-Nearest Neighbours")
|
|
ap.add_argument("--rf", action="store_true", help="Run Random Forest")
|
|
ap.add_argument("--svm", action="store_true", help="Run SVM")
|
|
ap.add_argument("--logreg", action="store_true", help="Run Logistic Regression")
|
|
ap.add_argument("--cnn", action="store_true", help="Run CNN")
|
|
ap.add_argument("--all", action="store_true", help="Run all classifiers")
|
|
|
|
# Data paths
|
|
ap.add_argument("--image-dir", default="Papila/FundusImages")
|
|
ap.add_argument("--clinical-dir", default="Papila/ClinicalData")
|
|
ap.add_argument("--label-col", default="Diagnosis")
|
|
ap.add_argument("--output-dir", default="analysis_data/papila_reproduce")
|
|
ap.add_argument("--tag", default=None,
|
|
help="Optional suffix appended to --output-dir (e.g. 'paper_matched').")
|
|
|
|
# Clinical data loader options
|
|
ap.add_argument("--no-iop-corr", action="store_true",
|
|
help="Skip IOP correction (use raw Perkins/Pneumatic values)")
|
|
ap.add_argument("--keep-raw-iop", action="store_true",
|
|
help="Keep raw IOP column alongside corrected IOP")
|
|
ap.add_argument("--no-cat-cols", action="store_true",
|
|
help="Exclude categorical columns (Gender, Phakic/Pseudophakic)")
|
|
ap.add_argument("--exclude-cols", nargs="*", default=[],
|
|
help="Additional columns to exclude from clinical feature matrix")
|
|
ap.add_argument("--keep-suspects", action="store_true",
|
|
help="Include Suspect (label 2) rows (default: drop them)")
|
|
ap.add_argument("--hypertower-loader", action="store_true",
|
|
help="Use the v3 HyperTower clinical data bundle (better IOP correction) "
|
|
"instead of the standalone ClinicalLoader.")
|
|
|
|
# CV
|
|
ap.add_argument("--n-splits", type=int, default=5)
|
|
ap.add_argument("--seed", type=int, default=42)
|
|
ap.add_argument("--paper-cv", action="store_true",
|
|
help="Use eye-level StratifiedKFold (matches paper's likely methodology) "
|
|
"instead of patient-level GroupKFold (our cleaner default).")
|
|
|
|
# CNN image cropping (GT or UNet, same flags as main hypertower)
|
|
ap.add_argument("--img-crop-manifest", default=None,
|
|
help="Path to crop manifest CSV (enables cropping).")
|
|
ap.add_argument("--img-crop-gt", action="store_true",
|
|
help="Use GT segmentations to crop (requires --img-crop-manifest).")
|
|
ap.add_argument("--img-crop-weights", default=None,
|
|
help="UNet weights path for disc cropping (requires --img-crop-manifest).")
|
|
ap.add_argument("--img-crop-scale", type=float, default=2.5)
|
|
ap.add_argument("--img-crop-size", type=int, default=200,
|
|
help="Crop target size in pixels (default 200, matching PAPILA paper).")
|
|
ap.add_argument("--img-crop-cache", default="cache_data/phase1_crops")
|
|
ap.add_argument("--persist-img-crop-cache", action="store_true")
|
|
|
|
# CNN options
|
|
ap.add_argument("--backbone", default="refugelike",
|
|
help="CNN backbone key (refugelike, resnet50, densenet121, vgg16, "
|
|
"efficientnet_b0, inception_v3, mobilenet_v2, refuge_densenet, ...)")
|
|
ap.add_argument("--backbones", nargs="+", default=None,
|
|
help="Run multiple backbones sequentially, sharing the image cache. "
|
|
"Overrides --backbone. e.g. --backbones resnet50 densenet121 vgg16")
|
|
ap.add_argument("--epochs", type=int, default=15)
|
|
ap.add_argument("--batch-size", type=int, default=16)
|
|
ap.add_argument("--lr", type=float, default=1e-4)
|
|
ap.add_argument("--freeze-ratio", type=float, default=0.0,
|
|
help="Fraction of backbone blocks to freeze (0=finetune all, 1=freeze all)")
|
|
ap.add_argument("--augment", action="store_true")
|
|
ap.add_argument("--device", default="auto",
|
|
choices=["auto", "cpu", "cuda"])
|
|
|
|
# Classical ML hyperparameters
|
|
ap.add_argument("--knn-k", type=int, default=5)
|
|
ap.add_argument("--rf-n-estimators", type=int, default=500)
|
|
ap.add_argument("--svm-c", type=float, default=1.0)
|
|
ap.add_argument("--lr-c", type=float, default=1.0)
|
|
return ap
|
|
|
|
|
|
def main():
|
|
ap = build_parser()
|
|
args = ap.parse_args()
|
|
|
|
if args.all:
|
|
args.knn = args.rf = args.svm = args.logreg = args.cnn = True
|
|
|
|
if not any([args.knn, args.rf, args.svm, args.logreg, args.cnn]):
|
|
ap.error("Specify at least one classifier: --knn --rf --svm --logreg --cnn (or --all)")
|
|
|
|
out_dir = Path(args.output_dir)
|
|
if args.tag:
|
|
out_dir = out_dir / args.tag
|
|
out_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
# Build clinical data loader
|
|
if args.hypertower_loader:
|
|
from v3.classes.papila_builders import build_papila_data
|
|
bundle = build_papila_data(
|
|
image_dir=args.image_dir,
|
|
clinical_dir=args.clinical_dir,
|
|
label_col=args.label_col,
|
|
cat_cols=["Gender", "Phakic/Pseudophakic"],
|
|
n_splits=args.n_splits,
|
|
random_seed=args.seed,
|
|
iop_corr_method="ratio",
|
|
iop_drop_raw=True,
|
|
)
|
|
loader = _BundleLoaderAdapter(bundle, label_col=args.label_col,
|
|
drop_suspects=not args.keep_suspects)
|
|
print(f"Clinical data: {len(loader.df)} rows [HyperTower loader] "
|
|
f"(suspects {'kept' if args.keep_suspects else 'dropped'})")
|
|
else:
|
|
loader = ClinicalLoader(
|
|
clinical_dir=args.clinical_dir,
|
|
label_col=args.label_col,
|
|
cat_cols=[] if args.no_cat_cols else None,
|
|
exclude_cols=list(args.exclude_cols or []),
|
|
iop_corr=not args.no_iop_corr,
|
|
keep_raw_iop=args.keep_raw_iop,
|
|
drop_suspects=not args.keep_suspects,
|
|
)
|
|
print(f"Clinical data: {len(loader.df)} rows "
|
|
f"(suspects {'kept' if args.keep_suspects else 'dropped'})")
|
|
|
|
X, y, feat_names, groups = loader.feature_matrix()
|
|
n_patients = len(np.unique(groups))
|
|
print(f"Feature matrix: {X.shape} ({n_patients} patients) class balance: {dict(zip(*np.unique(y, return_counts=True)))}")
|
|
|
|
overlay_curves: dict = {}
|
|
all_results: list = []
|
|
|
|
t0 = time.time()
|
|
|
|
# KNN
|
|
if args.knn:
|
|
print("\n--- KNN ---")
|
|
model = Pipeline([
|
|
("scale", StandardScaler()),
|
|
("knn", KNeighborsClassifier(n_neighbors=args.knn_k)),
|
|
])
|
|
r = run_classical("KNN", model, loader, out_dir, args.n_splits, args.seed, patient_level_cv=not args.paper_cv)
|
|
if r:
|
|
overlay_curves["KNN"] = (r["fpr"], r["tpr"], r["auc_mean"], r["auc_std"])
|
|
|
|
# Random Forest
|
|
if args.rf:
|
|
print("\n--- Random Forest ---")
|
|
model = RandomForestClassifier(
|
|
n_estimators=args.rf_n_estimators, max_features="sqrt",
|
|
random_state=args.seed, n_jobs=-1,
|
|
)
|
|
r = run_classical("Random Forest", model, loader, out_dir, args.n_splits, args.seed, patient_level_cv=not args.paper_cv)
|
|
if r:
|
|
overlay_curves["Random Forest"] = (r["fpr"], r["tpr"], r["auc_mean"], r["auc_std"])
|
|
|
|
# SVM
|
|
if args.svm:
|
|
print("\n--- SVM ---")
|
|
model = Pipeline([
|
|
("scale", StandardScaler()),
|
|
("svm", SVC(kernel="rbf", C=args.svm_c, gamma="scale",
|
|
probability=True, random_state=args.seed)),
|
|
])
|
|
r = run_classical("SVM", model, loader, out_dir, args.n_splits, args.seed, patient_level_cv=not args.paper_cv)
|
|
if r:
|
|
overlay_curves["SVM"] = (r["fpr"], r["tpr"], r["auc_mean"], r["auc_std"])
|
|
|
|
# Logistic Regression
|
|
if args.logreg:
|
|
print("\n--- Logistic Regression ---")
|
|
model = Pipeline([
|
|
("scale", StandardScaler()),
|
|
("logreg", LogisticRegression(C=args.lr_c, max_iter=1000,
|
|
solver="lbfgs")),
|
|
])
|
|
r = run_classical("Logistic Regression", model, loader, out_dir, args.n_splits, args.seed, patient_level_cv=not args.paper_cv)
|
|
if r:
|
|
overlay_curves["Logistic Regression"] = (r["fpr"], r["tpr"], r["auc_mean"], r["auc_std"])
|
|
|
|
# CNN
|
|
if args.cnn:
|
|
from v3.classes.croppers import build_image_preprocessor_from_args
|
|
from v3.classes.image_loader import CachedImageLoader as _CachedImageLoader
|
|
cnn_preprocessor = build_image_preprocessor_from_args(args)
|
|
backbones_to_run = args.backbones if args.backbones else [args.backbone]
|
|
shared_img_loader = _CachedImageLoader(enabled=True, workers=4)
|
|
for backbone in backbones_to_run:
|
|
print(f"\n--- CNN ({backbone})"
|
|
+ (" [cropped]" if cnn_preprocessor else "") + " ---")
|
|
r = run_cnn(
|
|
backbone=backbone,
|
|
image_dir=args.image_dir,
|
|
clinical_dir=args.clinical_dir,
|
|
label_col=args.label_col,
|
|
out_dir=out_dir,
|
|
n_splits=args.n_splits,
|
|
seed=args.seed,
|
|
epochs=args.epochs,
|
|
batch_size=args.batch_size,
|
|
lr=args.lr,
|
|
freeze_ratio=args.freeze_ratio,
|
|
augment=args.augment,
|
|
device_str=args.device,
|
|
drop_suspects=not args.keep_suspects,
|
|
preprocessor=cnn_preprocessor,
|
|
img_size=args.img_crop_size if cnn_preprocessor else 224,
|
|
img_loader=shared_img_loader,
|
|
)
|
|
if r:
|
|
overlay_curves[f"CNN ({backbone})"] = (r["fpr"], r["tpr"], r["auc_mean"], r["auc_std"])
|
|
|
|
# Overlay ROC
|
|
if len(overlay_curves) > 1:
|
|
_plot_overlay(overlay_curves, "PAPILA Reproduce — Clinical + CNN ROC", out_dir / "roc_overlay.png")
|
|
print(f"\nOverlay ROC saved: {out_dir / 'roc_overlay.png'}")
|
|
|
|
print(f"\nDone in {time.time()-t0:.1f}s — results in {out_dir}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|