#!/usr/bin/env python3 """ validate_siamese.py — Validate the trained siamese network against Task06 held-out test patients. Runs both connected-components and edge-ranking clustering (known K) on the test set and reports metrics + assignment heatmaps. Usage: conda activate fundus_imaging python scripts/clustering_validation/validate_siamese.py python scripts/clustering_validation/validate_siamese.py --threshold 0.95 --tag v2 """ import os, sys, json, argparse import numpy as np import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt os.environ["OMP_NUM_THREADS"] = "1" os.environ["OPENBLAS_NUM_THREADS"] = "1" os.environ["MKL_NUM_THREADS"] = "1" sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname( os.path.abspath(__file__))))) from classes import SiamesePatientMatcher from sklearn.metrics import adjusted_rand_score, normalized_mutual_info_score ROOT = os.path.dirname(os.path.dirname(os.path.dirname( os.path.abspath(__file__)))) MODELS_DIR = os.path.join(ROOT, "models") RESULTS_DIR = os.path.join(ROOT, "results", "clustering_validation") PLOTS_DIR = os.path.join(ROOT, "plots") PNG_DIR = os.path.join(ROOT, "features", "task06_pngs", "test") VOLUME_DIR = os.path.join(os.path.dirname(ROOT), "Task06_Lung", "imagesTr") SEED = 42 # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- def cluster_purity_stats(y_true, y_pred): purities = [] overall_correct = 0 for c in np.unique(y_pred): mask = y_pred == c _, counts = np.unique(y_true[mask], return_counts=True) purities.append(counts.max() / mask.sum()) overall_correct += counts.max() purities = np.array(purities) return { "overall": float(overall_correct / len(y_true)), "median": float(np.median(purities)), "mean": float(np.mean(purities)), "frac_gt_70": float((purities > 0.7).mean()), "frac_gt_90": float((purities > 0.9).mean()), } def patient_capture_stats(y_true, y_pred): captures = [] for p in np.unique(y_true): mask = y_true == p p_clusters = y_pred[mask] _, counts = np.unique(p_clusters, return_counts=True) captures.append(counts.max() / mask.sum()) captures = np.array(captures) return { "median": float(np.median(captures)), "mean": float(np.mean(captures)), "frac_gt_50": float((captures > 0.5).mean()), } def score_manifest(manifest, y_true_map): """Score a manifest against ground-truth patient IDs.""" preds, truths = [], [] for pred_pid, fnames in manifest.items(): for f in fnames: preds.append(pred_pid) truths.append(y_true_map.get(f, f"unknown_{f}")) y_true = np.array(truths) y_pred = np.array(preds) purity = cluster_purity_stats(y_true, y_pred) capture = patient_capture_stats(y_true, y_pred) ari = adjusted_rand_score(y_true, y_pred) nmi = normalized_mutual_info_score(y_true, y_pred) return {"ARI": ari, "NMI": nmi, "cluster_purity": purity, "patient_capture": capture, "y_true": y_true, "y_pred": y_pred} def plot_assignment_matrix(y_true, y_pred, out_path, title): true_patients = sorted(np.unique(y_true)) pred_clusters = sorted(np.unique(y_pred)) matrix = np.zeros((len(true_patients), len(pred_clusters))) for i, p in enumerate(true_patients): for j, c in enumerate(pred_clusters): matrix[i, j] = ((y_true == p) & (y_pred == c)).sum() matrix_norm = matrix / (matrix.sum(axis=1, keepdims=True) + 1e-8) fig, ax = plt.subplots(figsize=(max(14, len(pred_clusters) * 0.22), max(10, len(true_patients) * 0.18))) cmap = plt.cm.YlOrRd.copy() cmap.set_under('white') im = ax.imshow(matrix_norm, aspect="auto", cmap=cmap, vmin=1e-6, vmax=1) ax.set_xticks(range(len(pred_clusters))) ax.set_xticklabels([f"c{c}" for c in pred_clusters], fontsize=6, rotation=90) ax.set_yticks(range(len(true_patients))) ax.set_yticklabels(true_patients, fontsize=7) ax.set_xlabel("Predicted cluster"); ax.set_ylabel("True patient") ax.set_title(title, fontsize=11) plt.colorbar(im, ax=ax, label="Fraction of patient's slices") plt.tight_layout() os.makedirs(os.path.dirname(out_path), exist_ok=True) plt.savefig(out_path, dpi=150) plt.close() print(f" Heatmap → {out_path}") # --------------------------------------------------------------------------- # Main # --------------------------------------------------------------------------- def main(): ap = argparse.ArgumentParser() ap.add_argument("--model", default=os.path.join(MODELS_DIR, "siamese_resnet18.pt")) ap.add_argument("--backbone", default="resnet18") ap.add_argument("--threshold", type=float, default=0.9) ap.add_argument("--tag", default="", help="Append tag to output filenames") ap.add_argument("--full", action="store_true", help="Run on full NIfTI dataset (leaked training data).") args = ap.parse_args() tag = f"_{args.tag}" if args.tag else "" os.makedirs(RESULTS_DIR, exist_ok=True) os.makedirs(PLOTS_DIR, exist_ok=True) # ---- Load test PNGs ---- if args.full: tag = (tag or "") + "_full" from classes import NiftiSliceDataset VOL = os.path.join(os.path.dirname(ROOT), "Task06_Lung", "imagesTr") ds = NiftiSliceDataset(VOL, random_slices=False, seed=SEED, rotate_deg=90) ds.load_all_slices(stride=1) # Export full PNGs to temp dir import tempfile tmpdir = tempfile.mkdtemp(prefix="task06_full_") full_paths, full_pids, _ = ds.export_pngs(tmpdir) test_paths = full_paths test_pids = np.array(full_pids) print(f" FULL dataset: {len(test_paths)} slices, {len(np.unique(test_pids))} patients") print(f" (includes training data — leakage expected for siamese)") else: manifest_path = os.path.join(PNG_DIR, "manifest.json") if not os.path.exists(manifest_path): print(f"Test PNGs not found at {PNG_DIR}") print("Run train_siamese.py first to generate the test set.") sys.exit(1) with open(manifest_path) as f: png_manifest = json.load(f) test_paths = png_manifest["paths"] test_pids = np.array(png_manifest["patient_ids"]) # Short filenames for matching test_fnames = [p.split("/")[-1].replace(".png", "") for p in test_paths] y_true_map = {f: pid for f, pid in zip(test_fnames, test_pids)} k = len(np.unique(test_pids)) print("=" * 60) print(f"Siamese validation — held-out test set") print("=" * 60) print(f" {len(test_paths)} slices, {k} patients") print(f" Model: {args.model}") # ---- Load siamese model ---- print(f"\nLoading siamese model ...") matcher = SiamesePatientMatcher( args.model, backbone=args.backbone, input_size=224) # ---- Connected components (no known K) ---- print(f"\n{'─'*50}") print("Method 1: Connected components (threshold-based)") print(f"{'─'*50}") manifest_cc = matcher.identify_patients( test_paths, filenames=test_fnames, threshold=args.threshold, top_k=20, k=None) # ---- Edge-ranking clustering (known K) ---- print(f"\n{'─'*50}") print(f"Method 2: Edge-ranking clustering (k={k})") print(f"{'─'*50}") manifest_sc = matcher.identify_patients( test_paths, filenames=test_fnames, threshold=args.threshold, top_k=20, k=k) # ---- Score both methods ---- results_cc = score_manifest(manifest_cc, y_true_map) results_sc = score_manifest(manifest_sc, y_true_map) print(f"\n{'='*60}") print("RESULTS — Siamese patient identification") print(f"{'='*60}") print(f"\n{'Method':<30s} {'ARI':>7s} {'NMI':>7s} {' Purity Capture':>8s}") print(f"{'':30s} {'':>7s} {'':>7s} {' (overall) (mean) ':>8s}") print("-" * 60) for name, r in [("Connected components", results_cc), ("Edge ranking (k=" + str(k) + ")", results_sc)]: p = r["cluster_purity"]["overall"] c = r["patient_capture"]["mean"] print(f"{name:<30s} {r['ARI']:>7.3f} {r['NMI']:>7.3f} " f"{p:>7.1%} {c:>7.1%}") # ---- Save metrics ---- for suffix, r in [("cc", results_cc), ("sc", results_sc)]: out = {k: v for k, v in r.items() if k not in ("y_true", "y_pred")} out["method"] = suffix out["threshold"] = args.threshold out["n_patients"] = int(k) out["n_slices"] = len(test_paths) out_json = os.path.join(RESULTS_DIR, f"validate_siamese_{suffix}{tag}.json") with open(out_json, "w") as f: json.dump(out, f, indent=2) print(f" Metrics → {out_json}") # ---- Plot siamese methods ---- for name, r in [("connected_components", results_cc), ("edge_rank", results_sc)]: p = r["cluster_purity"]["overall"] c = r["patient_capture"]["mean"] title = (f"Siamese Patient-Cluster Assignment ({name})\n" f"(ARI={r['ARI']:.3f}, purity={p:.1%}, " f"capture mean={c:.1%})") plot_assignment_matrix( r["y_true"], r["y_pred"], os.path.join(PLOTS_DIR, "clustering_validation", "siamese", f"assignment_matrix_{name}{tag}.png"), title) print("DONE") if __name__ == "__main__": main()