"""Unified pipeline for the corrected thesis experiments. Given a holdout split manifest (build_holdout_split.py), this script: 1. Runs stratified 5-fold CV on the train+val pool for every selected model. 2. After CV, retrains the model on the full train+val pool. 3. Evaluates on the held-out independent test set and stores per-sample predictions for paired statistical testing. Designed to run unattended on the Azure T4 VM. All artefacts go to ``output_dir`` so the user can inspect them after the long run finishes. """ from __future__ import annotations import argparse import copy import json import random import time from collections import defaultdict from pathlib import Path from typing import Sequence import numpy as np import torch import torch.nn as nn import torch.nn.functional as F from PIL import Image from sklearn.metrics import ( accuracy_score, average_precision_score, confusion_matrix, f1_score, precision_recall_fscore_support, roc_auc_score, ) from sklearn.model_selection import StratifiedKFold from torch.utils.data import DataLoader, Dataset, WeightedRandomSampler from torchvision import models, transforms IMAGENET_MEAN = [0.485, 0.456, 0.406] IMAGENET_STD = [0.229, 0.224, 0.225] CLIP_MEAN = [0.48145466, 0.4578275, 0.40821073] CLIP_STD = [0.26862954, 0.26130258, 0.27577711] INCEPTION_SIZE = 299 class ImageListDataset(Dataset): def __init__(self, root: Path, samples: Sequence[tuple[str, int]], transform): self.root = root self.samples = samples self.transform = transform def __len__(self): return len(self.samples) def __getitem__(self, idx): rel_path, label = self.samples[idx] with Image.open(self.root / rel_path) as img: img = img.convert("RGB") if self.transform is not None: img = self.transform(img) return img, label class OpenCLIPClassifier(nn.Module): def __init__(self, num_classes, model_name="ViT-B-16", pretrained="openai"): super().__init__() import open_clip clip_model, _, _ = open_clip.create_model_and_transforms(model_name, pretrained=pretrained) self.backbone = clip_model with torch.no_grad(): dummy = torch.zeros(1, 3, 224, 224) feat_dim = self.backbone.encode_image(dummy).shape[-1] self.head = nn.Linear(feat_dim, num_classes) def forward(self, x): feats = self.backbone.encode_image(x) return self.head(feats.float()) def build_transforms(image_size, use_clip_norm=False): mean = CLIP_MEAN if use_clip_norm else IMAGENET_MEAN std = CLIP_STD if use_clip_norm else IMAGENET_STD train_tf = transforms.Compose([ transforms.RandomRotation(30), transforms.RandomHorizontalFlip(), transforms.RandomResizedCrop((image_size, image_size)), transforms.ToTensor(), transforms.Normalize(mean, std), ]) eval_tf = transforms.Compose([ transforms.Resize((image_size, image_size)), transforms.ToTensor(), transforms.Normalize(mean, std), ]) return train_tf, eval_tf def build_model(name, num_classes): if name == "vgg19": m = models.vgg19(weights=models.VGG19_Weights.IMAGENET1K_V1) m.classifier[-1] = nn.Linear(m.classifier[-1].in_features, num_classes) return m, 224, False if name == "resnet50": m = models.resnet50(weights=models.ResNet50_Weights.IMAGENET1K_V1) m.fc = nn.Linear(m.fc.in_features, num_classes) return m, 224, False if name == "resnet101": m = models.resnet101(weights=models.ResNet101_Weights.IMAGENET1K_V2) m.fc = nn.Linear(m.fc.in_features, num_classes) return m, 224, False if name == "densenet121": m = models.densenet121(weights=models.DenseNet121_Weights.IMAGENET1K_V1) m.classifier = nn.Linear(m.classifier.in_features, num_classes) return m, 224, False if name == "inception_v3": m = models.inception_v3(weights=models.Inception_V3_Weights.IMAGENET1K_V1, aux_logits=True) m.fc = nn.Linear(m.fc.in_features, num_classes) m.AuxLogits.fc = nn.Linear(m.AuxLogits.fc.in_features, num_classes) return m, INCEPTION_SIZE, False if name == "clip_openai": return OpenCLIPClassifier(num_classes), 224, True raise ValueError(name) def extract_logits_loss(out, labels, criterion): if isinstance(out, tuple): logits = out[0] loss = criterion(logits, labels) + 0.4 * criterion(out[1], labels) return logits, loss return out, criterion(out, labels) def train_epoch(model, loader, criterion, optim, scaler, device): model.train() total, count = 0.0, 0 for x, y in loader: x = x.to(device, non_blocking=True) y = y.to(device, non_blocking=True) optim.zero_grad(set_to_none=True) with torch.amp.autocast("cuda", enabled=scaler.is_enabled()): out = model(x) _, loss = extract_logits_loss(out, y, criterion) scaler.scale(loss).backward() scaler.step(optim) scaler.update() total += loss.item() * x.size(0) count += x.size(0) return total / max(count, 1) @torch.no_grad() def evaluate(model, loader, device, return_preds=False, num_classes=None): model.eval() preds, labels, probs = [], [], [] for x, y in loader: x = x.to(device, non_blocking=True) out = model(x) logits = out[0] if isinstance(out, tuple) else out p = F.softmax(logits, dim=1) preds.extend(logits.argmax(1).cpu().tolist()) probs.extend(p.cpu().numpy().tolist()) labels.extend(y.tolist()) p_m, r_m, f_m, _ = precision_recall_fscore_support(labels, preds, average="macro", zero_division=0) acc = accuracy_score(labels, preds) res = {"accuracy": acc, "precision_macro": p_m, "recall_macro": r_m, "f1_macro": f_m} # ROC-AUC and PR-AUC (one-vs-rest, macro) try: n = num_classes or (max(labels) + 1) y_onehot = np.eye(n)[np.array(labels)] probs_arr = np.array(probs) res["roc_auc_macro"] = float(roc_auc_score(y_onehot, probs_arr, average="macro", multi_class="ovr")) res["pr_auc_macro"] = float(average_precision_score(y_onehot, probs_arr, average="macro")) except Exception as exc: res["roc_auc_macro"] = None res["pr_auc_macro"] = None # Expected Calibration Error (15 bins) res["ece"] = float(expected_calibration_error(np.array(probs), np.array(labels))) if return_preds: res["preds"] = preds res["labels"] = labels res["probs"] = probs return res def expected_calibration_error(probs, labels, n_bins=15): confidences = probs.max(axis=1) predictions = probs.argmax(axis=1) accuracies = (predictions == labels).astype(float) bin_boundaries = np.linspace(0, 1, n_bins + 1) ece = 0.0 for lo, hi in zip(bin_boundaries[:-1], bin_boundaries[1:]): in_bin = (confidences > lo) & (confidences <= hi) if in_bin.sum() > 0: avg_conf = confidences[in_bin].mean() avg_acc = accuracies[in_bin].mean() ece += (in_bin.sum() / len(probs)) * abs(avg_conf - avg_acc) return ece def bootstrap_ci(labels, preds, metric_fn, n_resamples=1000, alpha=0.05, seed=0): rng = np.random.RandomState(seed) labels = np.array(labels) preds = np.array(preds) stats = [] n = len(labels) for _ in range(n_resamples): idx = rng.randint(0, n, size=n) stats.append(metric_fn(labels[idx], preds[idx])) stats = np.array(stats) return {"mean": float(stats.mean()), "lo": float(np.quantile(stats, alpha / 2)), "hi": float(np.quantile(stats, 1 - alpha / 2))} def compute_class_weights(labels, num_classes, smoothing="sqrt"): counts = np.bincount(labels, minlength=num_classes).astype(float) counts[counts == 0] = 1.0 # avoid div0 if smoothing == "sqrt": weights = 1.0 / np.sqrt(counts) else: weights = 1.0 / counts weights = weights * num_classes / weights.sum() return torch.tensor(weights, dtype=torch.float32) def class_names_to_int(samples, classes): cls_to_idx = {c: i for i, c in enumerate(classes)} return [(p, cls_to_idx[c]) for p, c in samples] def train_model(model_name, train_samples, val_samples, classes, root, args, device, log): model_tuple = build_model(model_name, len(classes)) model, image_size, use_clip_norm = model_tuple model = model.to(device) train_tf, eval_tf = build_transforms(image_size, use_clip_norm) train_set = ImageListDataset(root, train_samples, train_tf) val_set = ImageListDataset(root, val_samples, eval_tf) train_loader = DataLoader(train_set, batch_size=args.batch_size, shuffle=True, num_workers=args.workers, pin_memory=device.type == "cuda") val_loader = DataLoader(val_set, batch_size=args.batch_size, shuffle=False, num_workers=args.workers, pin_memory=device.type == "cuda") # Class-weighted cross-entropy to mitigate severe imbalance train_labels = [s[1] for s in train_samples] cls_weights = compute_class_weights(train_labels, len(classes)).to(device) if args.class_weighted else None criterion = nn.CrossEntropyLoss(weight=cls_weights) optim = torch.optim.Adam(model.parameters(), lr=args.lr, weight_decay=1e-4) scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optim, T_max=args.epochs) scaler = torch.amp.GradScaler("cuda", enabled=device.type == "cuda") best_state, best_acc, best_epoch = None, 0.0, 0 bad_epochs = 0 for ep in range(1, args.epochs + 1): tloss = train_epoch(model, train_loader, criterion, optim, scaler, device) scheduler.step() v = evaluate(model, val_loader, device, num_classes=len(classes)) line = (f"{model_name} ep {ep:03d} train_loss={tloss:.4f} " f"val_acc={v['accuracy']:.4f} val_f1={v['f1_macro']:.4f}") log(line) if v["accuracy"] > best_acc: best_acc = v["accuracy"] best_epoch = ep best_state = copy.deepcopy(model.state_dict()) bad_epochs = 0 else: bad_epochs += 1 if bad_epochs >= args.patience: log(f" early stop at ep {ep} (best ep {best_epoch} acc {best_acc:.4f})") break if best_state is not None: model.load_state_dict(best_state) return model, {"best_val_acc": best_acc, "best_epoch": best_epoch, "image_size": image_size, "use_clip_norm": use_clip_norm} def run_kfold(model_name, samples_pool, classes, root, args, device, log): paths = [s[0] for s in samples_pool] labels = [s[1] for s in samples_pool] skf = StratifiedKFold(n_splits=args.folds, shuffle=True, random_state=args.seed) fold_metrics = [] for k, (tr_idx, vl_idx) in enumerate(skf.split(paths, labels), start=1): log(f"=== {model_name} fold {k}/{args.folds} ===") tr = [samples_pool[i] for i in tr_idx] vl = [samples_pool[i] for i in vl_idx] model, info = train_model(model_name, tr, vl, classes, root, args, device, log) # final fold metrics on val _, eval_tf = build_transforms(info["image_size"], info["use_clip_norm"]) v_loader = DataLoader( ImageListDataset(root, vl, eval_tf), batch_size=args.batch_size, shuffle=False, num_workers=args.workers, pin_memory=device.type == "cuda", ) m = evaluate(model, v_loader, device, return_preds=True, num_classes=len(classes)) fold_metrics.append({"fold": k, **{k2: v for k2, v in m.items() if k2 not in ("preds", "labels", "probs")}}) log(f" fold {k} acc={m['accuracy']:.4f} f1={m['f1_macro']:.4f}") del model torch.cuda.empty_cache() accs = np.array([f["accuracy"] for f in fold_metrics]) f1s = np.array([f["f1_macro"] for f in fold_metrics]) summary = { "model": model_name, "folds": args.folds, "accuracy_mean": float(accs.mean()), "accuracy_std": float(accs.std(ddof=1)), "accuracy_ci95": [float(accs.mean() - 1.96 * accs.std(ddof=1) / np.sqrt(len(accs))), float(accs.mean() + 1.96 * accs.std(ddof=1) / np.sqrt(len(accs)))], "f1_macro_mean": float(f1s.mean()), "f1_macro_std": float(f1s.std(ddof=1)), "fold_metrics": fold_metrics, } return summary def run_indep_test(model_name, samples_pool, test_samples, classes, root, args, device, log, weights_dir): log(f"=== {model_name} FINAL (train on pool, eval on indep test) ===") # 90/10 split within the pool to keep an internal val for early stopping pool_paths = [s[0] for s in samples_pool] pool_labels = [s[1] for s in samples_pool] rng = np.random.RandomState(args.seed) idx = np.arange(len(samples_pool)) rng.shuffle(idx) cut = int(0.9 * len(idx)) tr = [samples_pool[i] for i in idx[:cut]] vl = [samples_pool[i] for i in idx[cut:]] model, info = train_model(model_name, tr, vl, classes, root, args, device, log) weights_path = weights_dir / f"{model_name}_final.pth" torch.save(model.state_dict(), weights_path) log(f" saved weights: {weights_path}") _, eval_tf = build_transforms(info["image_size"], info["use_clip_norm"]) test_loader = DataLoader( ImageListDataset(root, test_samples, eval_tf), batch_size=args.batch_size, shuffle=False, num_workers=args.workers, pin_memory=device.type == "cuda", ) m = evaluate(model, test_loader, device, return_preds=True, num_classes=len(classes)) cm = confusion_matrix(m["labels"], m["preds"], labels=list(range(len(classes)))) # Bootstrap 95% CIs on accuracy and macro-F1 boot_acc = bootstrap_ci(m["labels"], m["preds"], lambda y, p: float((y == p).mean()), seed=args.seed) boot_f1 = bootstrap_ci(m["labels"], m["preds"], lambda y, p: f1_score(y, p, average="macro", zero_division=0), seed=args.seed + 1) log(f" test acc={m['accuracy']:.4f} f1={m['f1_macro']:.4f} " f"roc_auc={m.get('roc_auc_macro')} ece={m['ece']:.4f}") del model torch.cuda.empty_cache() return { "model": model_name, "test_accuracy": m["accuracy"], "test_precision_macro": m["precision_macro"], "test_recall_macro": m["recall_macro"], "test_f1_macro": m["f1_macro"], "test_roc_auc_macro": m.get("roc_auc_macro"), "test_pr_auc_macro": m.get("pr_auc_macro"), "test_ece": m["ece"], "bootstrap_accuracy_ci95": boot_acc, "bootstrap_f1_macro_ci95": boot_f1, "preds": m["preds"], "labels": m["labels"], "probs": m["probs"], "confusion_matrix": cm.tolist(), } def mcnemar_pairwise(test_results): """Compute pairwise McNemar p-values between models on the independent test set.""" from scipy.stats import binom out = {} names = sorted(test_results.keys()) for i in range(len(names)): for j in range(i + 1, len(names)): a, b = names[i], names[j] preds_a = np.array(test_results[a]["preds"]) preds_b = np.array(test_results[b]["preds"]) labels = np.array(test_results[a]["labels"]) ca = preds_a == labels cb = preds_b == labels b_count = int(np.sum(ca & ~cb)) # a right, b wrong c_count = int(np.sum(~ca & cb)) # a wrong, b right n = b_count + c_count if n == 0: p = 1.0 else: k = min(b_count, c_count) p = float(2 * binom.cdf(k, n, 0.5)) p = min(p, 1.0) out[f"{a}_vs_{b}"] = {"b": b_count, "c": c_count, "p_value": p} return out def parse_args(): p = argparse.ArgumentParser() p.add_argument("--manifest", default="holdout_split.json") p.add_argument("--models", nargs="+", default=[ "vgg19", "resnet50", "resnet101", "densenet121", "inception_v3", "clip_openai", ]) p.add_argument("--epochs", type=int, default=60) p.add_argument("--batch-size", type=int, default=32) p.add_argument("--workers", type=int, default=4) p.add_argument("--lr", type=float, default=1e-4) p.add_argument("--folds", type=int, default=5) p.add_argument("--patience", type=int, default=10) p.add_argument("--seed", type=int, default=42) p.add_argument("--output-dir", default="final_experiments") p.add_argument("--skip-kfold", action="store_true") p.add_argument("--skip-test", action="store_true") p.add_argument("--class-weighted", action="store_true", default=True, help="Use sqrt-inverse-frequency class-weighted CE loss to mitigate imbalance.") p.add_argument("--no-class-weighted", dest="class_weighted", action="store_false") return p.parse_args() def main(): args = parse_args() manifest = json.loads(Path(args.manifest).read_text()) root = Path(manifest["data_dir"]) classes = manifest["classes"] pool_samples = class_names_to_int( [(p, c) for p, c in zip(manifest["kfold"]["pool_paths"], manifest["kfold"]["pool_labels"])], classes, ) test_samples = class_names_to_int(manifest["splits"]["test"], classes) out_dir = Path(args.output_dir) out_dir.mkdir(parents=True, exist_ok=True) weights_dir = out_dir / "weights" weights_dir.mkdir(exist_ok=True) log_path = out_dir / "run.log" def log(line): msg = f"[{time.strftime('%Y-%m-%d %H:%M:%S')}] {line}" print(msg, flush=True) with log_path.open("a") as fh: fh.write(msg + "\n") random.seed(args.seed) np.random.seed(args.seed) torch.manual_seed(args.seed) torch.cuda.manual_seed_all(args.seed) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") log(f"device={device} classes={len(classes)} pool={len(pool_samples)} test={len(test_samples)}") kfold_summaries = {} test_results = {} for model_name in args.models: log(f"########## {model_name} ##########") try: if not args.skip_kfold: s = run_kfold(model_name, pool_samples, classes, root, args, device, log) kfold_summaries[model_name] = s (out_dir / f"{model_name}_kfold.json").write_text(json.dumps(s, indent=2)) if not args.skip_test: r = run_indep_test(model_name, pool_samples, test_samples, classes, root, args, device, log, weights_dir) test_results[model_name] = r # Save without huge preds/labels/probs arrays inline slim = {k: v for k, v in r.items() if k not in ("preds", "labels", "probs")} (out_dir / f"{model_name}_test.json").write_text(json.dumps(slim, indent=2)) # Save preds + probs separately for stat tests, ROC, calibration analysis (out_dir / f"{model_name}_test_preds.json").write_text(json.dumps({ "preds": r["preds"], "labels": r["labels"], "probs": r["probs"]})) except Exception as exc: log(f"!! {model_name} FAILED: {exc!r}") if test_results: mc = mcnemar_pairwise(test_results) (out_dir / "mcnemar.json").write_text(json.dumps(mc, indent=2)) log(f"McNemar pairwise saved -> {out_dir / 'mcnemar.json'}") if kfold_summaries: summary = {m: {k: v for k, v in s.items() if k != "fold_metrics"} for m, s in kfold_summaries.items()} (out_dir / "kfold_summary.json").write_text(json.dumps(summary, indent=2)) log("ALL DONE") if __name__ == "__main__": main()