#!/usr/bin/env python3 """Validate PCA-50 feature K-means on the same Task06 test patients as siamese.""" import os, sys, json, argparse import numpy as np import matplotlib; matplotlib.use("Agg") import matplotlib.pyplot as plt from PIL import Image as PILImage 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 PatientIdentifier, FeatureExtractor 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__)))) PNG_DIR = os.path.join(ROOT, "features", "task06_pngs", "test") RESULTS_DIR = os.path.join(ROOT, "results", "clustering_validation") PLOTS_DIR = os.path.join(ROOT, "plots") SEED = 42 MODELS = ["VGG16", "DenseNet121", "EfficientNetB1", "MobileNetV2", "ResNet50"] 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 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') 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(ax.images[0], 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}") def main(): ap = argparse.ArgumentParser() ap.add_argument("--model", default="VGG16", choices=MODELS) ap.add_argument("--all-models", action="store_true") ap.add_argument("--tag", default="") ap.add_argument("--full", action="store_true", help="Run on full NIfTI dataset (63 patients, not just test).") 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) 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) test_slices = ds.slices_rgb test_pids = ds.patient_labels test_fnames = ds.filenames y_true_map = {f: pid for f, pid in zip(test_fnames, test_pids)} k = ds.n_patients y_true = np.array([y_true_map[f] for f in test_fnames]) print(f" FULL dataset: {len(test_fnames)} slices, {k} patients") 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}. Run train_siamese.py first.") sys.exit(1) with open(manifest_path) as f: png = json.load(f) test_paths = png["paths"] test_pids = np.array(png["patient_ids"]) 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)) y_true = np.array([y_true_map[f] for f in test_fnames]) test_slices = [PILImage.open(p).convert("L") for p in test_paths] print(f" {len(test_fnames)} slices, {k} patients (same test set as siamese)") models_to_run = MODELS if args.all_models else [args.model] for model_name in models_to_run: print(f"\n{'─'*50}") print(f"Model: {model_name}") print(f"{'─'*50}") print(f" Extracting features ...", flush=True) ext = FeatureExtractor(model_name=model_name) X, _, _ = ext.extract_from_images(test_slices) print(f" Clustering PCA-50 into k={k} ...", flush=True) ident = PatientIdentifier(patient_estimates={"lung": k}, random_state=SEED) groups = ident.identify_from_features(X, np.full(len(X), "lung")) y_pred = np.array([str(g) for g in groups]) 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) print(f"\n RESULTS — PCA-50 K-means ({model_name})") print(f" {'─'*45}") print(f" ARI: {ari:.4f}") print(f" NMI: {nmi:.4f}") print(f" Overall purity: {purity['overall']:.3f} ({purity['overall']*100:.1f}%)") print(f" Mean capture: {capture['mean']:.3f} ({capture['mean']*100:.1f}%)") results = {"method": f"pca50_{model_name}", "model": model_name, "n_patients": int(k), "n_slices": len(test_fnames), "ARI": ari, "NMI": nmi, "cluster_purity": purity, "patient_capture": capture} out_json = os.path.join(RESULTS_DIR, f"validate_pca50_{model_name}{tag}.json") with open(out_json, "w") as f: json.dump(results, f, indent=2) print(f" Metrics → {out_json}") plot_assignment_matrix(y_true, y_pred, os.path.join(PLOTS_DIR, "clustering_validation", "pca50", f"assignment_matrix_{model_name}{tag}.png"), f"PCA-50 K-means {model_name} (ARI={ari:.3f}, purity={purity['overall']:.1%}, capture mean={capture['mean']:.1%})") # Summary table if all models if args.all_models: print(f"\n{'='*60}") print("SUMMARY — All models") print(f"{'='*60}") print(f"{'Model':<18s} {'ARI':>7s} {'Purity':>8s} {'Capture':>8s}") print("-" * 43) for model_name in MODELS: j = os.path.join(RESULTS_DIR, f"validate_pca50_{model_name}{tag}.json") if os.path.exists(j): with open(j) as f: d = json.load(f) print(f"{model_name:<18s} {d['ARI']:>7.3f} {d['cluster_purity']['overall']:>7.1%} {d['patient_capture']['mean']:>7.1%}") print("DONE") if __name__ == "__main__": main()