""" Evaluation script — Attention U-Net. Loads each test GeoTIFF, normalizes and pads/crops to 128×128 (same as training), predicts, then evaluates only on the VALID (non-padded) pixels so zero-padding does not inflate scores. All outputs go directly into outputs/ — no subdirectories. Outputs ------- outputs/ attn_evaluation_metrics.csv — global threshold sweep 0.30–0.70 attn_per_image_metrics.csv — per-image TP/FP/FN/TN, IoU, F1, etc. attn_confusion_matrix.csv — aggregate CM at best IoU threshold attn_evaluation_summary.txt — human-readable summary attn_pred_*.png — RGB / GT / prediction side-by-side """ import os import glob import numpy as np np.Inf = np.inf import pandas as pd import rasterio import tensorflow as tf from sklearn.metrics import confusion_matrix from data_generator import GLOBAL_MEAN, GLOBAL_STD from models.attention_unet import build_attention_unet # ───────────────────────────────────────────────────────────────────────────── # Config # ───────────────────────────────────────────────────────────────────────────── TEST_FEAT = "dataset/test/features" TEST_MASK = "dataset/test/masks" CKPT_PATH = "outputs/checkpoints/attention_unet_best.weights.h5" OUT_DIR = "outputs" INPUT_SHAPE = (128, 128, 12) TARGET = INPUT_SHAPE[:2] BATCH_SIZE = 8 THRESHOLDS = [0.30, 0.35, 0.40, 0.45, 0.50, 0.55, 0.60, 0.65, 0.70] os.makedirs(OUT_DIR, exist_ok=True) # ───────────────────────────────────────────────────────────────────────────── # Loss / metric stubs (needed for model.compile) # ───────────────────────────────────────────────────────────────────────────── def dice_loss(y_true, y_pred): smooth = 1e-6 yt = tf.cast(y_true, tf.float32) i = tf.reduce_sum(yt * y_pred, axis=[1, 2, 3]) u = tf.reduce_sum(yt, axis=[1, 2, 3]) + tf.reduce_sum(y_pred, axis=[1, 2, 3]) return 1.0 - tf.reduce_mean((2.0 * i + smooth) / (u + smooth)) def focal_loss(alpha=0.75, gamma=2.0): def _l(y_true, y_pred): yt = tf.cast(y_true, tf.float32) yp = tf.clip_by_value(y_pred, tf.keras.backend.epsilon(), 1-tf.keras.backend.epsilon()) at = yt*alpha + (1-yt)*(1-alpha) pt = yt*yp + (1-yt)*(1-yp) return tf.reduce_mean(-at * tf.pow(1-pt, gamma) * tf.math.log(pt)) return _l def combined_loss(y_true, y_pred): return focal_loss()(y_true, y_pred)*2.0 + dice_loss(y_true, y_pred) def dice_coef(y_true, y_pred): yp = tf.cast(y_pred > 0.5, tf.float32); yt = tf.cast(y_true, tf.float32) return (2*tf.reduce_sum(yt*yp)+1e-6)/(tf.reduce_sum(yt)+tf.reduce_sum(yp)+1e-6) def precision_m(y_true, y_pred): yp = tf.cast(y_pred > 0.5, tf.float32); yt = tf.cast(y_true, tf.float32) tp = tf.reduce_sum(yt*yp); fp = tf.reduce_sum(yp*(1-yt)) return (tp+1e-6)/(tp+fp+1e-6) def recall_m(y_true, y_pred): yp = tf.cast(y_pred > 0.5, tf.float32); yt = tf.cast(y_true, tf.float32) tp = tf.reduce_sum(yt*yp); fn = tf.reduce_sum(yt*(1-yp)) return (tp+1e-6)/(tp+fn+1e-6) def f1_m(y_true, y_pred): p = precision_m(y_true, y_pred); r = recall_m(y_true, y_pred) return (2*p*r+1e-6)/(p+r+1e-6) def iou_m(y_true, y_pred): yp = tf.cast(y_pred > 0.5, tf.float32); yt = tf.cast(y_true, tf.float32) i = tf.reduce_sum(yt*yp); u = tf.reduce_sum(yt)+tf.reduce_sum(yp)-i return (i+1e-6)/(u+1e-6) # ───────────────────────────────────────────────────────────────────────────── # Image preprocessing (identical to training — no data leakage) # ───────────────────────────────────────────────────────────────────────────── def load_and_prep(f_path, m_path): with rasterio.open(f_path) as s: feat = np.moveaxis(s.read(), 0, -1).astype(np.float32) with rasterio.open(m_path) as s: mask = np.expand_dims(s.read(1), -1).astype(np.float32) H, W = feat.shape[:2] feat = np.where((feat == -9999) | ~np.isfinite(feat), 0, feat) feat = (feat - GLOBAL_MEAN) / (GLOBAL_STD + 1e-7) th, tw = TARGET pad_h, pad_w = max(0, th - H), max(0, tw - W) if pad_h or pad_w: feat = np.pad(feat, ((0, pad_h), (0, pad_w), (0, 0))) mask = np.pad(mask, ((0, pad_h), (0, pad_w), (0, 0))) h, w = feat.shape[:2] if h > th or w > tw: hs = (h - th) // 2 if h > th else 0 ws = (w - tw) // 2 if w > tw else 0 feat = feat[hs:hs+th, ws:ws+tw, :] mask = mask[hs:hs+th, ws:ws+tw, :] valid_h = min(H, th) valid_w = min(W, tw) padded = pad_h > 0 or pad_w > 0 cropped = H > th or W > tw return feat, mask, H, W, valid_h, valid_w, padded, cropped # ───────────────────────────────────────────────────────────────────────────── # Build + load weights # ───────────────────────────────────────────────────────────────────────────── print("Building Attention U-Net…") model = build_attention_unet(input_shape=INPUT_SHAPE) model.compile(loss=combined_loss, metrics=['accuracy', dice_coef, precision_m, recall_m, f1_m, iou_m]) if not os.path.exists(CKPT_PATH): raise FileNotFoundError( f"Checkpoint not found: {CKPT_PATH}\n" "Run: conda run -n brats python train_attention.py") model.load_weights(CKPT_PATH) print("Weights loaded.") # ───────────────────────────────────────────────────────────────────────────── # Load all test images # ───────────────────────────────────────────────────────────────────────────── feat_files = sorted(glob.glob(os.path.join(TEST_FEAT, 'feature_*.tif'))) mask_files = sorted(glob.glob(os.path.join(TEST_MASK, 'mask_*.tif'))) assert len(feat_files) > 0, f"No test images found in {TEST_FEAT}" assert len(feat_files) == len(mask_files) print(f"\nLoading {len(feat_files)} test images…") X, Y, meta, ids = [], [], [], [] for fp, mp in zip(feat_files, mask_files): feat, mask, H, W, vh, vw, padded, cropped = load_and_prep(fp, mp) X.append(feat); Y.append(mask) meta.append((H, W, vh, vw, padded, cropped)) ids.append(os.path.basename(fp).replace('feature_', '').replace('.tif', '')) X = np.array(X, dtype=np.float32) Y = np.array(Y, dtype=np.float32) print("Predicting…") Y_prob = model.predict(X, batch_size=BATCH_SIZE, verbose=1) # Keras evaluate at t=0.5 print("\nmodel.evaluate (t=0.5):") eval_vals = model.evaluate(X, Y, batch_size=BATCH_SIZE, verbose=1) eval_names = model.metrics_names # ───────────────────────────────────────────────────────────────────────────── # Helper # ───────────────────────────────────────────────────────────────────────────── def cm_metrics(tp, fp, fn, tn): prec = tp / (tp + fp + 1e-9) rec = tp / (tp + fn + 1e-9) f1 = 2 * prec * rec / (prec + rec + 1e-9) iou = tp / (tp + fp + fn + 1e-9) dice = 2 * tp / (2*tp + fp + fn + 1e-9) acc = (tp + tn) / (tp + tn + fp + fn + 1e-9) spec = tn / (tn + fp + 1e-9) return prec, rec, f1, iou, dice, acc, spec # ───────────────────────────────────────────────────────────────────────────── # Per-image metrics @ t=0.5 (valid region only — no padding pixels) # ───────────────────────────────────────────────────────────────────────────── all_true, all_prob = [], [] per_image = [] for i, iid in enumerate(ids): H, W, vh, vw, padded, cropped = meta[i] yt = Y[i, :vh, :vw, 0].flatten().astype(np.int32) yprob = Y_prob[i,:vh, :vw, 0].flatten() all_true.append(yt); all_prob.append(yprob) yp = (yprob >= 0.5).astype(np.int32) tn, fp, fn, tp = confusion_matrix(yt, yp, labels=[0,1]).ravel() prec, rec, f1, iou, dice, acc, spec = cm_metrics(tp, fp, fn, tn) per_image.append({ 'image_id': iid, 'orig_H': H, 'orig_W': W, 'valid_px': int(yt.size), 'was_padded': bool(padded), 'was_cropped': bool(cropped), 'glacier_px_gt': int(yt.sum()), 'glacier_frac_gt': round(float(yt.mean()), 4), 'glacier_px_pred': int(yp.sum()), 'TP': int(tp), 'FP': int(fp), 'FN': int(fn), 'TN': int(tn), 'precision': round(prec, 4), 'recall': round(rec, 4), 'f1_score': round(f1, 4), 'iou': round(iou, 4), 'dice': round(dice, 4), 'accuracy': round(acc, 4), 'specificity':round(spec, 4), }) per_img_df = pd.DataFrame(per_image).sort_values('iou') per_img_df.to_csv(os.path.join(OUT_DIR, 'attn_per_image_metrics.csv'), index=False) # ───────────────────────────────────────────────────────────────────────────── # Global threshold sweep # ───────────────────────────────────────────────────────────────────────────── y_true_all = np.concatenate(all_true) y_prob_all = np.concatenate(all_prob) thresh_rows = [] for t in THRESHOLDS: yp = (y_prob_all >= t).astype(np.int32) tn, fp, fn, tp = confusion_matrix(y_true_all, yp, labels=[0,1]).ravel() prec, rec, f1, iou, dice, acc, spec = cm_metrics(tp, fp, fn, tn) thresh_rows.append({ 'threshold': t, 'TP': int(tp), 'FP': int(fp), 'TN': int(tn), 'FN': int(fn), 'precision': round(prec, 6), 'recall': round(rec, 6), 'specificity': round(spec, 6), 'f1_score': round(f1, 6), 'iou': round(iou, 6), 'dice': round(dice, 6), 'accuracy': round(acc, 6), }) thresh_df = pd.DataFrame(thresh_rows) thresh_df.to_csv(os.path.join(OUT_DIR, 'attn_evaluation_metrics.csv'), index=False) # Best threshold by IoU best_row = thresh_df.loc[thresh_df['iou'].idxmax()] best_t = float(best_row['threshold']) # Aggregate confusion matrix at best threshold cm = confusion_matrix(y_true_all, (y_prob_all >= best_t).astype(int), labels=[0,1]) pd.DataFrame(cm, index=['actual_background', 'actual_glacier'], columns=['pred_background', 'pred_glacier'] ).to_csv(os.path.join(OUT_DIR, 'attn_confusion_matrix.csv')) # ───────────────────────────────────────────────────────────────────────────── # Human-readable summary # ───────────────────────────────────────────────────────────────────────────── lines = [ "=" * 58, " ATTENTION U-NET — EVALUATION SUMMARY", "=" * 58, "", "model.evaluate (threshold = 0.5):", ] for n, v in zip(eval_names, eval_vals): lines.append(f" {n:<22}: {v:.6f}") lines += [ "", f"Best threshold by IoU: {best_t}", "", " {:<10} {:<10} {:<10} {:<10} {:<10} {:<10}".format( "thresh", "precision", "recall", "f1", "iou", "dice"), " " + "-" * 60, ] for _, r in thresh_df.iterrows(): tag = " ◄ best" if r['threshold'] == best_t else "" lines.append(" {:<10} {:<10.4f} {:<10.4f} {:<10.4f} {:<10.4f} {:<10.4f}{}".format( r['threshold'], r['precision'], r['recall'], r['f1_score'], r['iou'], r['dice'], tag)) lines += [ "", "Per-image stats (valid pixels only, t=0.5):", f" Mean IoU : {per_img_df['iou'].mean():.4f}", f" Median IoU: {per_img_df['iou'].median():.4f}", f" Worst 5 :", ] for _, r in per_img_df.head(5).iterrows(): lines.append(f" {r['image_id']} — IoU={r['iou']:.4f} F1={r['f1_score']:.4f}" f" glacier={r['glacier_frac_gt']*100:.1f}% padded={r['was_padded']}") lines += [ "", f"Total test images : {len(feat_files)}", f"Output dir : {OUT_DIR}/", "", "Files written:", " attn_evaluation_metrics.csv — full threshold sweep", " attn_per_image_metrics.csv — per-image metrics (sorted worst IoU first)", " attn_confusion_matrix.csv — aggregate CM at best threshold", " attn_evaluation_summary.txt", " attn_pred_*.png — comparison figures", "", "─" * 58, "Baseline (SwinUNet 128×128, 70 epochs):", " Precision=0.6233 Recall=0.8990 F1=0.7362 IoU=0.5825", "─" * 58, ] summary = "\n".join(lines) print("\n" + summary) with open(os.path.join(OUT_DIR, 'attn_evaluation_summary.txt'), 'w') as f: f.write(summary) # ───────────────────────────────────────────────────────────────────────────── # Save comparison PNGs for top 20 worst + 5 best images # ───────────────────────────────────────────────────────────────────────────── try: import matplotlib matplotlib.use('Agg') import matplotlib.pyplot as plt y_pred_bin = (Y_prob >= best_t).astype(np.float32) save_indices = list(per_img_df.head(10).index) + list(per_img_df.tail(5).index) save_indices = list(dict.fromkeys(save_indices))[:20] for i in save_indices: iid = ids[i] H, W, vh, vw, _, _ = meta[i] rgb = X[i, :, :, [2, 1, 0]] lo, hi = np.percentile(rgb, 2), np.percentile(rgb, 98) rgb = np.clip((rgb - lo) / (hi - lo + 1e-8), 0, 1) fig, axes = plt.subplots(1, 3, figsize=(9, 3)) axes[0].imshow(rgb); axes[0].set_title('RGB'); axes[0].axis('off') axes[1].imshow(Y[i,:,:,0], cmap='gray'); axes[1].set_title('Ground Truth'); axes[1].axis('off') axes[2].imshow(y_pred_bin[i,:,:,0], cmap='gray') iou_val = per_img_df.loc[per_img_df['image_id']==iid, 'iou'].values axes[2].set_title(f'Pred IoU={iou_val[0]:.3f}' if len(iou_val) else 'Pred') axes[2].axis('off') plt.tight_layout() plt.savefig(os.path.join(OUT_DIR, f'attn_pred_{iid}.png'), dpi=80) plt.close() print(f"Saved {len(save_indices)} PNGs to {OUT_DIR}/") except Exception as e: print(f"PNG save skipped: {e}") print("\nEvaluation complete.")