rpotter6298
328 lines
11 KiB
Python
Executable File
328 lines
11 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""Train a CNN (resnet50 backbone), extract logits, and train RF on logits+metadata with 5-fold CV."""
|
|
from __future__ import annotations
|
|
|
|
import random
|
|
from pathlib import Path
|
|
import sys
|
|
from typing import Dict, List, Tuple
|
|
|
|
import numpy as np
|
|
import pandas as pd
|
|
from PIL import Image
|
|
import torch
|
|
from torch import nn
|
|
from torch.utils.data import DataLoader, Dataset
|
|
from torchvision import transforms
|
|
from sklearn.ensemble import RandomForestClassifier
|
|
from sklearn.metrics import accuracy_score, roc_auc_score
|
|
|
|
REPO_ROOT = Path(__file__).resolve().parents[2]
|
|
if str(REPO_ROOT) not in sys.path:
|
|
sys.path.insert(0, str(REPO_ROOT))
|
|
|
|
from classes import build_papila_clinical
|
|
from classes.backbones import BACKBONES, load_backbone_weights
|
|
|
|
|
|
# ---------------------------
|
|
# Config (edit in IDE)
|
|
# ---------------------------
|
|
IMAGE_DIR = "Papila/FundusImages"
|
|
CLINICAL_DIR = "Papila/ClinicalData"
|
|
LABEL_COL = "Diagnosis"
|
|
CAT_COLS = ["Gender", "Phakic/Pseudophakic"]
|
|
EVAL_MODE = "binary" # "binary" or "multiclass"
|
|
N_SPLITS = 5
|
|
FOLD_SEED = 42
|
|
HOLDOUT_SEED = 123
|
|
HOLDOUT_PATIENTS_PER_CLASS = 6
|
|
|
|
BACKBONE_NAME = "resnet50"
|
|
BATCH_SIZE = 8
|
|
EPOCHS = 40
|
|
LR = 1e-4
|
|
WEIGHT_DECAY = 1e-5
|
|
|
|
RF_TREES = 500
|
|
RF_MAX_DEPTH = None
|
|
RF_MIN_SAMPLES_LEAF = 1
|
|
|
|
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
|
|
OUTPUT_DIR = Path("analysis_data/basic_analysis/cnn_logits_rf_cv")
|
|
PRINT_EPOCH_REPORT = True
|
|
EPOCH_REPORT_EVERY = 1
|
|
|
|
|
|
class PapilaImageDataset(Dataset):
|
|
def __init__(
|
|
self, clinical, df: pd.DataFrame, label_col: str, img_transform
|
|
) -> None:
|
|
self.clinical = clinical
|
|
self.df = df.reset_index(drop=True)
|
|
self.label_col = label_col
|
|
self.img_transform = img_transform
|
|
|
|
def __len__(self) -> int:
|
|
return len(self.df)
|
|
|
|
def __getitem__(self, idx: int):
|
|
row = self.df.iloc[idx]
|
|
img_path = self.clinical.get_image_path(row)
|
|
image = Image.open(img_path).convert("RGB")
|
|
x_img = self.img_transform(image)
|
|
y = int(row[self.label_col])
|
|
x_md = self.clinical.vectorize_row(row).astype(np.float32)
|
|
return x_img, y, x_md
|
|
|
|
|
|
class CNNHead(nn.Module):
|
|
def __init__(self, backbone_name: str, num_classes: int) -> None:
|
|
super().__init__()
|
|
spec = BACKBONES[backbone_name]
|
|
backbone = spec.ctor(weights=spec.weights_default)
|
|
if backbone_name.startswith("refuge"):
|
|
load_backbone_weights(backbone_name, backbone)
|
|
out_dim, backbone = spec.strip(backbone)
|
|
self.backbone = backbone
|
|
self.head = nn.Linear(out_dim, num_classes)
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
feats = self.backbone(x)
|
|
return self.head(feats)
|
|
|
|
|
|
def _set_seed(seed: int) -> None:
|
|
random.seed(seed)
|
|
np.random.seed(seed)
|
|
torch.manual_seed(seed)
|
|
if torch.cuda.is_available():
|
|
torch.cuda.manual_seed_all(seed)
|
|
|
|
|
|
def _auc_score(y_true: np.ndarray, probs: np.ndarray, num_classes: int) -> float:
|
|
try:
|
|
if num_classes == 2:
|
|
return float(roc_auc_score(y_true, probs[:, 1]))
|
|
return float(roc_auc_score(y_true, probs, multi_class="ovr", average="macro"))
|
|
except Exception:
|
|
return float("nan")
|
|
|
|
|
|
def _prepare_clinical() -> Tuple[object, pd.DataFrame]:
|
|
clinical = build_papila_clinical(
|
|
image_dir=IMAGE_DIR,
|
|
clinical_dir=CLINICAL_DIR,
|
|
label_col=LABEL_COL,
|
|
cat_cols=CAT_COLS,
|
|
n_splits=N_SPLITS,
|
|
random_seed=FOLD_SEED,
|
|
)
|
|
df = clinical.df.copy()
|
|
if EVAL_MODE == "binary":
|
|
df = df[df[LABEL_COL].isin([0, 1])].reset_index(drop=True)
|
|
return clinical, df
|
|
|
|
|
|
def _split_holdout_by_patient(df: pd.DataFrame) -> Tuple[pd.DataFrame, pd.DataFrame]:
|
|
rng = np.random.default_rng(HOLDOUT_SEED)
|
|
patient_label = (
|
|
df.groupby("Patient ID")[LABEL_COL]
|
|
.agg(lambda s: int(s.mode().iloc[0]))
|
|
.reset_index()
|
|
)
|
|
holdout_patients = []
|
|
for lbl, grp in patient_label.groupby(LABEL_COL):
|
|
candidates = grp["Patient ID"].to_numpy()
|
|
n = min(HOLDOUT_PATIENTS_PER_CLASS, len(candidates))
|
|
if n <= 0:
|
|
continue
|
|
selected = rng.choice(candidates, size=n, replace=False)
|
|
holdout_patients.extend(selected.tolist())
|
|
holdout_patients = sorted(set(holdout_patients))
|
|
holdout_df = df[df["Patient ID"].isin(holdout_patients)].reset_index(drop=True)
|
|
train_df = df[~df["Patient ID"].isin(holdout_patients)].reset_index(drop=True)
|
|
return train_df, holdout_df
|
|
|
|
|
|
def _rebuild_clinical_from_df(clinical, df: pd.DataFrame) -> object:
|
|
clinical.frames = [df.copy()]
|
|
clinical.df = df.copy()
|
|
clinical._infer_or_validate_feature_types()
|
|
clinical._compute_numeric_stats()
|
|
clinical._build_cat_maps()
|
|
clinical._compute_feature_dim()
|
|
clinical._build_kfold_indices()
|
|
return clinical
|
|
|
|
|
|
def _train_cnn(
|
|
model: nn.Module, loader: DataLoader, num_classes: int, fold: int
|
|
) -> None:
|
|
model.train()
|
|
optimizer = torch.optim.Adam(model.parameters(), lr=LR, weight_decay=WEIGHT_DECAY)
|
|
criterion = nn.CrossEntropyLoss()
|
|
for epoch in range(EPOCHS):
|
|
running_loss = 0.0
|
|
correct = 0
|
|
total = 0
|
|
for x_img, y, _x_md in loader:
|
|
x_img = x_img.to(DEVICE)
|
|
y = y.to(DEVICE)
|
|
optimizer.zero_grad()
|
|
logits = model(x_img)
|
|
loss = criterion(logits, y)
|
|
loss.backward()
|
|
optimizer.step()
|
|
running_loss += float(loss.item()) * int(y.size(0))
|
|
pred = torch.argmax(logits, dim=1)
|
|
correct += int((pred == y).sum().item())
|
|
total += int(y.size(0))
|
|
|
|
if PRINT_EPOCH_REPORT and ((epoch + 1) % EPOCH_REPORT_EVERY == 0):
|
|
avg_loss = running_loss / max(total, 1)
|
|
train_acc = correct / max(total, 1)
|
|
print(
|
|
f"[fold {fold + 1}/{N_SPLITS}] epoch {epoch + 1}/{EPOCHS} "
|
|
f"train_loss={avg_loss:.4f} train_acc={train_acc:.4f}"
|
|
)
|
|
|
|
|
|
def _infer_logits(
|
|
model: nn.Module, loader: DataLoader
|
|
) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
|
|
model.eval()
|
|
logits_all, probs_all, y_all, md_all = [], [], [], []
|
|
with torch.no_grad():
|
|
for x_img, y, x_md in loader:
|
|
x_img = x_img.to(DEVICE)
|
|
logits = model(x_img).cpu().numpy()
|
|
probs = torch.softmax(torch.from_numpy(logits), dim=1).numpy()
|
|
logits_all.append(logits)
|
|
probs_all.append(probs)
|
|
y_all.append(y.numpy())
|
|
md_all.append(x_md.numpy())
|
|
return (
|
|
np.concatenate(y_all, axis=0),
|
|
np.concatenate(logits_all, axis=0),
|
|
np.concatenate(md_all, axis=0),
|
|
)
|
|
|
|
|
|
def main() -> None:
|
|
_set_seed(FOLD_SEED)
|
|
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
|
|
|
|
num_classes = 2 if EVAL_MODE == "binary" else 3
|
|
train_tf = transforms.Compose(
|
|
[
|
|
transforms.Resize((224, 224)),
|
|
transforms.RandomHorizontalFlip(),
|
|
transforms.ToTensor(),
|
|
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
|
|
]
|
|
)
|
|
eval_tf = transforms.Compose(
|
|
[
|
|
transforms.Resize((224, 224)),
|
|
transforms.ToTensor(),
|
|
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
|
|
]
|
|
)
|
|
|
|
clinical, df = _prepare_clinical()
|
|
train_df, holdout_df = _split_holdout_by_patient(df)
|
|
clinical = _rebuild_clinical_from_df(clinical, train_df)
|
|
holdout_df.to_csv(OUTPUT_DIR / "holdout_patients.csv", index=False)
|
|
|
|
rows: List[Dict[str, object]] = []
|
|
holdout_rows: List[Dict[str, object]] = []
|
|
|
|
for fold in range(N_SPLITS):
|
|
print(f"\n[info] Starting fold {fold + 1}/{N_SPLITS}")
|
|
fold_train_df, fold_val_df = clinical.get_split_dfs(fold)
|
|
ds_train = PapilaImageDataset(clinical, fold_train_df, LABEL_COL, train_tf)
|
|
ds_val = PapilaImageDataset(clinical, fold_val_df, LABEL_COL, eval_tf)
|
|
ds_holdout = PapilaImageDataset(clinical, holdout_df, LABEL_COL, eval_tf)
|
|
|
|
dl_train = DataLoader(
|
|
ds_train, batch_size=BATCH_SIZE, shuffle=True, num_workers=0
|
|
)
|
|
dl_val = DataLoader(ds_val, batch_size=BATCH_SIZE, shuffle=False, num_workers=0)
|
|
dl_holdout = DataLoader(
|
|
ds_holdout, batch_size=BATCH_SIZE, shuffle=False, num_workers=0
|
|
)
|
|
|
|
model = CNNHead(BACKBONE_NAME, num_classes=num_classes).to(DEVICE)
|
|
_train_cnn(model, dl_train, num_classes=num_classes, fold=fold)
|
|
|
|
y_tr, log_tr, md_tr = _infer_logits(
|
|
model,
|
|
DataLoader(ds_train, batch_size=BATCH_SIZE, shuffle=False, num_workers=0),
|
|
)
|
|
y_va, log_va, md_va = _infer_logits(model, dl_val)
|
|
y_ho, log_ho, md_ho = _infer_logits(model, dl_holdout)
|
|
|
|
np.save(OUTPUT_DIR / f"fold{fold}_train_logits.npy", log_tr)
|
|
np.save(OUTPUT_DIR / f"fold{fold}_val_logits.npy", log_va)
|
|
np.save(OUTPUT_DIR / f"fold{fold}_holdout_logits.npy", log_ho)
|
|
|
|
X_tr = np.concatenate([log_tr, md_tr], axis=1)
|
|
X_va = np.concatenate([log_va, md_va], axis=1)
|
|
X_ho = np.concatenate([log_ho, md_ho], axis=1)
|
|
|
|
rf = RandomForestClassifier(
|
|
n_estimators=RF_TREES,
|
|
max_depth=RF_MAX_DEPTH,
|
|
min_samples_leaf=RF_MIN_SAMPLES_LEAF,
|
|
class_weight="balanced",
|
|
random_state=FOLD_SEED + fold,
|
|
n_jobs=-1,
|
|
)
|
|
rf.fit(X_tr, y_tr)
|
|
|
|
p_va = rf.predict_proba(X_va)
|
|
p_ho = rf.predict_proba(X_ho)
|
|
pred_va = np.argmax(p_va, axis=1)
|
|
pred_ho = np.argmax(p_ho, axis=1)
|
|
|
|
rows.append(
|
|
{
|
|
"fold": fold,
|
|
"val_acc": float(accuracy_score(y_va, pred_va)),
|
|
"val_auc": _auc_score(y_va, p_va, num_classes),
|
|
"n_val": int(len(y_va)),
|
|
}
|
|
)
|
|
holdout_rows.append(
|
|
{
|
|
"fold": fold,
|
|
"holdout_acc": float(accuracy_score(y_ho, pred_ho)),
|
|
"holdout_auc": _auc_score(y_ho, p_ho, num_classes),
|
|
"n_holdout": int(len(y_ho)),
|
|
}
|
|
)
|
|
print(
|
|
f"[info] Fold {fold + 1} RF: val_acc={rows[-1]['val_acc']:.4f} val_auc={rows[-1]['val_auc']:.4f} "
|
|
f"| holdout_acc={holdout_rows[-1]['holdout_acc']:.4f} holdout_auc={holdout_rows[-1]['holdout_auc']:.4f}"
|
|
)
|
|
|
|
fold_df = pd.DataFrame(rows)
|
|
holdout_df = pd.DataFrame(holdout_rows)
|
|
fold_df.to_csv(OUTPUT_DIR / "rf_val_metrics.csv", index=False)
|
|
holdout_df.to_csv(OUTPUT_DIR / "rf_holdout_metrics.csv", index=False)
|
|
|
|
print("\nRF validation metrics:")
|
|
print(fold_df.to_string(index=False, float_format=lambda x: f"{x:.4f}"))
|
|
print("\nRF holdout metrics:")
|
|
print(holdout_df.to_string(index=False, float_format=lambda x: f"{x:.4f}"))
|
|
print(
|
|
f"\nMeans: val_acc={fold_df['val_acc'].mean():.4f}, val_auc={fold_df['val_auc'].mean():.4f}, "
|
|
f"holdout_acc={holdout_df['holdout_acc'].mean():.4f}, holdout_auc={holdout_df['holdout_auc'].mean():.4f}"
|
|
)
|
|
print(f"\nSaved outputs to: {OUTPUT_DIR}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|