"""Stage 1 v1 final: 4-model ensemble stacking (convnext, swin, mamba, resnet). Both cRT and plain ensembles. 3 combiner methods: A. logreg_logits_balanced — LogReg on 12-dim raw logits (4 backbones × 3 classes) B. logreg_softmax_balanced — LogReg on 12-dim softmax probabilities C. avg_softmax — simple softmax-probability averaging Best method selected by pooled OOF macro-F1, then applied to test. Usage: python ensemble.py --resize letterbox python ensemble.py --resize squeeze """ import argparse import json import pathlib import sys import numpy as np from scipy.special import softmax from sklearn.linear_model import LogisticRegression from sklearn.preprocessing import StandardScaler from sklearn.metrics import (f1_score, accuracy_score, balanced_accuracy_score, confusion_matrix, classification_report) ROOT = pathlib.Path(__file__).parent sys.path.insert(0, str(ROOT)) BACKBONES = ["convnext", "swin", "mamba", "resnet"] NUM_CLASSES = 3 SEED = 17 METHODS = ["logreg_logits_balanced", "logreg_softmax_balanced", "avg_softmax"] def get_paths(resize): return ROOT / "logits" / resize, ROOT / "results" / resize def load_split(logits_dir, kind, prefix=""): """kind = 'oof' or 'test'. prefix = '' for cRT, 'plain_' for plain. Logits are stored in per-backbone subdirectories: logits_dir/{bb}/{prefix}{bb}_{kind}.npz """ out = {} for bb in BACKBONES: f = logits_dir / bb / f"{prefix}{bb}_{kind}.npz" z = np.load(f, allow_pickle=True) out[bb] = {k: z[k] for k in z.files} return out def stack_features(d, backbones=BACKBONES, use_softmax=False): arrs = [] for bb in backbones: lg = d[bb]["logits"].astype(np.float32) arrs.append(softmax(lg, axis=1) if use_softmax else lg) return np.concatenate(arrs, axis=1).astype(np.float32) def compute_per_class_precision(true, pred, names): true = np.asarray(true); pred = np.asarray(pred) result = {} for i, name in enumerate(names): tp = ((pred == i) & (true == i)).sum() fp = ((pred == i) & (true != i)).sum() result[name] = float(tp / max(tp + fp, 1)) return result def metrics_from_pred(pred, true, src=None, names=None): if names is None: names = ["malignant", "non_malignant", "non_roi"] pred = np.asarray(pred); true = np.asarray(true) m = { "macro_f1": float(f1_score(true, pred, average="macro", labels=list(range(NUM_CLASSES)), zero_division=0)), "accuracy": float(accuracy_score(true, pred)), "balanced_acc": float(balanced_accuracy_score(true, pred)), "per_class_recall": {names[i]: float(((true == i) & (pred == i)).sum() / max((true == i).sum(), 1)) for i in range(NUM_CLASSES)}, "per_class_precision": compute_per_class_precision(true, pred, names), "confusion_matrix": confusion_matrix(true, pred, labels=list(range(NUM_CLASSES))).tolist(), "classification_report": classification_report(true, pred, labels=list(range(NUM_CLASSES)), target_names=names, output_dict=True, zero_division=0), "n": len(pred), } if src is not None: src = np.asarray(src) for s in ["B", "C", "D"]: mk = src == s if mk.sum() > 0: m[f"src_{s}_acc"] = float(accuracy_score(true[mk], pred[mk])) m[f"src_{s}_n"] = int(mk.sum()) return m def make_logreg(): return LogisticRegression(C=1.0, max_iter=2000, solver="lbfgs", class_weight="balanced", random_state=SEED) def predict_method(method, X_tr, y_tr, X_va, n_backbones=len(BACKBONES)): if method == "avg_softmax": blocks = [X_va[:, i*NUM_CLASSES:(i+1)*NUM_CLASSES] for i in range(n_backbones)] avg = np.mean(blocks, axis=0) return avg.argmax(1) scaler = StandardScaler() Xtr = scaler.fit_transform(X_tr) Xva = scaler.transform(X_va) clf = make_logreg() clf.fit(Xtr, y_tr) return clf.predict(Xva) def run_cv_method(method, X_oof, y_oof, cv_folds, src_oof): fold_metrics = [] for k in range(5): tr_mask = cv_folds != k va_mask = cv_folds == k pred_va = predict_method(method, X_oof[tr_mask], y_oof[tr_mask], X_oof[va_mask]) m = metrics_from_pred(pred_va, y_oof[va_mask], src_oof[va_mask]) m["fold"] = int(k) fold_metrics.append(m) f1s = [fm["macro_f1"] for fm in fold_metrics] nms = [fm["per_class_recall"]["non_malignant"] for fm in fold_metrics] baccs = [fm["balanced_acc"] for fm in fold_metrics] prec_mal = [fm["per_class_precision"]["malignant"] for fm in fold_metrics] prec_nm = [fm["per_class_precision"]["non_malignant"] for fm in fold_metrics] prec_nroi = [fm["per_class_precision"]["non_roi"] for fm in fold_metrics] rec_mal = [fm["per_class_recall"]["malignant"] for fm in fold_metrics] rec_nm = [fm["per_class_recall"]["non_malignant"] for fm in fold_metrics] rec_nroi = [fm["per_class_recall"]["non_roi"] for fm in fold_metrics] return { "folds": fold_metrics, "cv_macro_f1_mean": float(np.mean(f1s)), "cv_macro_f1_std": float(np.std(f1s)), "cv_non_malignant_recall_mean": float(np.mean(nms)), "cv_non_malignant_recall_std": float(np.std(nms)), "cv_balanced_acc_mean": float(np.mean(baccs)), "cv_balanced_acc_std": float(np.std(baccs)), "cv_malignant_precision_mean": float(np.mean(prec_mal)), "cv_malignant_precision_std": float(np.std(prec_mal)), "cv_non_malignant_precision_mean": float(np.mean(prec_nm)), "cv_non_malignant_precision_std": float(np.std(prec_nm)), "cv_non_roi_precision_mean": float(np.mean(prec_nroi)), "cv_non_roi_precision_std": float(np.std(prec_nroi)), "cv_malignant_recall_mean": float(np.mean(rec_mal)), "cv_malignant_recall_std": float(np.std(rec_mal)), "cv_non_malignant_recall_mean": float(np.mean(nms)), "cv_non_malignant_recall_std": float(np.std(nms)), "cv_non_roi_recall_mean": float(np.mean(rec_nroi)), "cv_non_roi_recall_std": float(np.std(rec_nroi)), } def run_pooled_method(method, X_oof, y_oof, cv_folds, src_oof): oof_pred = np.zeros(len(y_oof), dtype=np.int64) for k in range(5): tr_mask = cv_folds != k va_mask = cv_folds == k oof_pred[va_mask] = predict_method(method, X_oof[tr_mask], y_oof[tr_mask], X_oof[va_mask]) return metrics_from_pred(oof_pred, y_oof, src_oof) def fit_predict_test(method, X_oof, y_oof, X_test, n_backbones=len(BACKBONES)): if method == "avg_softmax": blocks = [X_test[:, i*NUM_CLASSES:(i+1)*NUM_CLASSES] for i in range(n_backbones)] avg = np.mean(blocks, axis=0) return avg.argmax(1), avg scaler = StandardScaler() Xtr = scaler.fit_transform(X_oof) Xte = scaler.transform(X_test) clf = make_logreg() clf.fit(Xtr, y_oof) pred = clf.predict(Xte) proba = clf.predict_proba(Xte) return pred, proba def single_metrics(logits_dir, kind, prefix=""): d = load_split(logits_dir, kind, prefix=prefix) out = {} for bb in BACKBONES: pred = d[bb]["logits"].argmax(1) y = d[bb]["labels"]; src = d[bb].get("source", None) out[bb] = metrics_from_pred(pred, y, src) return out def run_ensemble(label, prefix, logits_dir, results_dir): """label='cRT' or 'plain'. prefix='' for cRT, 'plain_' for plain.""" print(f"\n=== Ensemble stacking ({label}, 4 backbones: {BACKBONES}) ===") oof = load_split(logits_dir, "oof", prefix=prefix) test = load_split(logits_dir, "test", prefix=prefix) for bb in BACKBONES: assert np.array_equal(oof[bb]["labels"], oof["convnext"]["labels"]), f"oof labels mismatch {bb}" assert np.array_equal(test[bb]["labels"], test["convnext"]["labels"]), f"test labels mismatch {bb}" print("alignment OK") y_oof = oof["convnext"]["labels"].astype(np.int64) y_test = test["convnext"]["labels"].astype(np.int64) cv_folds = oof["convnext"]["cv_fold"].astype(np.int64) src_oof = oof["convnext"]["source"] src_test = test["convnext"]["source"] X_oof_raw = stack_features(oof, use_softmax=False) X_oof_sm = stack_features(oof, use_softmax=True) X_test_raw = stack_features(test, use_softmax=False) X_test_sm = stack_features(test, use_softmax=True) print(f"X_oof_raw {X_oof_raw.shape} X_oof_sm {X_oof_sm.shape}") method_features = { "logreg_logits_balanced": (X_oof_raw, X_test_raw), "logreg_softmax_balanced": (X_oof_sm, X_test_sm), "avg_softmax": (X_oof_sm, X_test_sm), } all_cv = {} all_pooled = {} all_test = {} for method in METHODS: Xo, Xt = method_features[method] print(f"\n--- {method} CV per-fold ---") cv_res = run_cv_method(method, Xo, y_oof, cv_folds, src_oof) for fm in cv_res["folds"]: print(f" fold{fm['fold']}: f1={fm['macro_f1']:.4f} nm_rec={fm['per_class_recall']['non_malignant']:.4f}") print(f" CV: f1={cv_res['cv_macro_f1_mean']:.4f}+-{cv_res['cv_macro_f1_std']:.4f}") print(f"--- {method} CV pooled ---") pooled = run_pooled_method(method, Xo, y_oof, cv_folds, src_oof) print(f" pooled: f1={pooled['macro_f1']:.4f} nm_rec={pooled['per_class_recall']['non_malignant']:.4f}") print(f"--- {method} Test ---") pred_te, proba_te = fit_predict_test(method, Xo, y_oof, Xt) tm = metrics_from_pred(pred_te, y_test, src_test) print(f" test: f1={tm['macro_f1']:.4f} nm_rec={tm['per_class_recall']['non_malignant']:.4f} bacc={tm['balanced_acc']:.4f}") all_cv[method] = cv_res all_pooled[method] = pooled all_test[method] = {"test_metrics": tm, "pred": pred_te.tolist(), "proba": proba_te.tolist()} best_method = max(METHODS, key=lambda m: all_pooled[m]["macro_f1"]) print(f"\n=== Best method: {best_method} (pooled OOF macro-F1={all_pooled[best_method]['macro_f1']:.4f}) ===") tag = label.lower() # 'crt' or 'plain' cv_summary = { "label": label, "methods": METHODS, "best_method": best_method, "backbones": BACKBONES, "per_method_cv": {m: all_cv[m] for m in METHODS}, "per_method_pooled": {m: all_pooled[m] for m in METHODS}, "best_cv": all_cv[best_method], "best_pooled": all_pooled[best_method], } (results_dir / f"ensemble_{tag}_cv_summary.json").write_text(json.dumps(cv_summary, indent=2)) best_test = all_test[best_method] test_summary = { "label": label, "best_method": best_method, "test_metrics": best_test["test_metrics"], "all_methods_test": {m: all_test[m]["test_metrics"] for m in METHODS}, } (results_dir / f"ensemble_{tag}_test.json").write_text(json.dumps(test_summary, indent=2)) # Single model metrics single_cv = single_metrics(logits_dir, "oof", prefix=prefix) single_te = single_metrics(logits_dir, "test", prefix=prefix) comp = { "cv_single": {bb: single_cv[bb] for bb in BACKBONES}, "cv_ensemble_best_per_fold": all_cv[best_method], "cv_ensemble_best_pooled": all_pooled[best_method], "test_single": {bb: single_te[bb] for bb in BACKBONES}, "test_ensemble_best": best_test["test_metrics"], "best_method": best_method, "all_methods_test": {m: all_test[m]["test_metrics"] for m in METHODS}, } (results_dir / f"ensemble_{tag}_single_vs_ensemble.json").write_text(json.dumps(comp, indent=2)) print(f"\nBest method: {best_method}") print(f"CV per-fold f1: {all_cv[best_method]['cv_macro_f1_mean']:.4f}+-{all_cv[best_method]['cv_macro_f1_std']:.4f}") print(f"CV pooled f1: {all_pooled[best_method]['macro_f1']:.4f}") print(f"Test f1: {best_test['test_metrics']['macro_f1']:.4f}") return cv_summary, test_summary, comp def main(): ap = argparse.ArgumentParser() ap.add_argument("--resize", choices=["letterbox", "squeeze"], required=True) args = ap.parse_args() # Map "squeeze" arg to "resized" directory name resize_dir = "resized" if args.resize == "squeeze" else args.resize logits_dir, results_dir = get_paths(resize_dir) results_dir.mkdir(parents=True, exist_ok=True) crt_cv, crt_test, crt_comp = run_ensemble("cRT", "", logits_dir, results_dir) plain_cv, plain_test, plain_comp = run_ensemble("plain", "plain_", logits_dir, results_dir) # Consolidated all_results = { "backbones": BACKBONES, "resize": args.resize, "crt": {"cv_summary": crt_cv, "test_summary": crt_test, "comparison": crt_comp}, "plain": {"cv_summary": plain_cv, "test_summary": plain_test, "comparison": plain_comp}, } (results_dir / "ensemble_all.json").write_text(json.dumps(all_results, indent=2)) print("\n=== All 4-model ensembles complete ===") if __name__ == "__main__": main()