rockglacierNET / evaluate.py
kabiriitb's picture
Upload folder using huggingface_hub
820fab0 verified
Raw History Blame Contribute Delete
22 kB
"""
Evaluation script for RockGlacierNet β€” baseline 128x128 full-image model.
Each test GeoTIFF is normalized and padded/cropped to 128x128 exactly like
training, then predicted. Metrics are computed on the VALID (non-padded) region
of each image so zero-padding does not inflate the scores.
Outputs (outputs/predictions/baseline_model_output/)
-----------------------------------------------------
per_image_metrics.csv β€” ONE ROW PER TEST IMAGE, sorted worst IoU first
(this is the file to inspect for problem images)
evaluation_metrics.csv β€” global threshold sweep 0.30–0.70
confusion_matrix.csv β€” aggregate CM at best IoU threshold
evaluation_summary.txt β€” human-readable summary
pred_*.png β€” RGB / GT / prediction comparison figures
"""
import os
import glob
import json
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.swin_unet import build_model
from models.attention_unet import build_attention_unet
from models.unet import build_unet
from models.segformer import build_segformer
# ─────────────────────────────────────────────────────────────────────────────
# Config
# ─────────────────────────────────────────────────────────────────────────────
import glob as _glob
DATA_DIR = os.environ.get("DATA_DIR", "all_data/patches")
TEST_FEAT = os.environ.get("TEST_FEAT", f"{DATA_DIR}/test/features")
TEST_MASK = os.environ.get("TEST_MASK", f"{DATA_DIR}/test/masks")
# Resolve the run name -> checkpoint/result dirs (mirrors train.py).
RUN_NAME = os.environ.get("RUN_NAME")
if RUN_NAME is None:
_sel = os.environ.get("BANDS", "all")
RUN_NAME = "all" if _sel == "all" else _sel + "band"
CKPT_DIR = os.environ.get("CKPT_DIR",
"all_data/checkpoint" if RUN_NAME == "all" else f"all_data/checkpoint_{RUN_NAME}")
CKPT_PATH = os.environ.get("CKPT_PATH", f"{CKPT_DIR}/best_model.weights.h5")
OUT_DIR = os.environ.get("OUT_DIR",
"all_data/result" if RUN_NAME == "all" else f"all_data/result_{RUN_NAME}")
# Bands MUST match training. Prefer run_config.json written by train.py;
# else BAND_LIST env; else all bands. This makes a mismatch impossible.
_band_names = json.load(open(f"{DATA_DIR}/metadata/norm_stats.json"))["band_names"]
_rc = os.path.join(CKPT_DIR, "run_config.json")
if os.environ.get("BAND_LIST"):
_bands = [b.strip() for b in os.environ["BAND_LIST"].split(",") if b.strip()]
elif os.path.exists(_rc):
_bands = json.load(open(_rc)).get("bands")
else:
_bands = None
if _bands and list(_bands) != list(_band_names):
BAND_INDICES = [_band_names.index(b) for b in _bands]
else:
BAND_INDICES = None
# Patch size auto-detected from a sample test patch.
_ph, _pw = np.load(sorted(_glob.glob(f"{TEST_FEAT}/*.npy"))[0]).shape[:2]
INPUT_SHAPE = (_ph, _pw, len(BAND_INDICES) if BAND_INDICES else len(_band_names))
TARGET = INPUT_SHAPE[:2]
print(f"[eval] run={RUN_NAME} bands={INPUT_SHAPE[2]} patch={_ph}x{_pw} ckpt={CKPT_PATH}")
BATCH_SIZE = 8
THRESHOLDS = [0.30, 0.35, 0.40, 0.45, 0.50, 0.55, 0.60, 0.65, 0.70]
USE_TTA = True # 8-way test-time augmentation (4 rotations x 2 flips)
os.makedirs(OUT_DIR, exist_ok=True)
# ─────────────────────────────────────────────────────────────────────────────
# Normalization: load the exact stats the checkpoint was trained with
# (checkpoint_dataset/norm_stats.json, written by train.py). Falls back to the
# built-in GLOBAL_MEAN/STD when absent so OLD checkpoints stay consistent.
# ─────────────────────────────────────────────────────────────────────────────
def _load_eval_norm():
cand = os.path.join(os.path.dirname(CKPT_PATH), "norm_stats.json")
if os.path.exists(cand):
st = json.load(open(cand))
m = np.array(st["mean"], np.float32).reshape(1, 1, -1)
s = np.array(st["std"], np.float32).reshape(1, 1, -1)
print(f"[eval] normalization: loaded {cand} (matches training)")
return m, s
print("[eval] normalization: built-in GLOBAL_MEAN/STD (no saved stats in checkpoint dir)")
return GLOBAL_MEAN, GLOBAL_STD
NORM_MEAN, NORM_STD = _load_eval_norm()
# ─────────────────────────────────────────────────────────────────────────────
# Loss / metric helpers (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 combined_loss(y_true, y_pred):
bce = tf.reduce_mean(tf.keras.losses.binary_crossentropy(y_true, y_pred)) * 2.0
return bce + 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)
# ─────────────────────────────────────────────────────────────────────────────
# Preprocess one image exactly like the (non-augment) data generator.
# Returns the padded/cropped 128x128 feature + mask AND the valid region size.
# ─────────────────────────────────────────────────────────────────────────────
def load_and_prep(f_path, m_path):
with rasterio.open(f_path) as s:
feat = np.moveaxis(s.read(), 0, -1) # (H, W, 12)
with rasterio.open(m_path) as s:
mask = np.expand_dims(s.read(1), axis=-1) # (H, W, 1)
# Force mask to match feature shape (legacy .tif misalignment fix)
H_f, W_f = feat.shape[:2]
mask = mask[:H_f, :W_f, :]
pad_mh, pad_mw = max(0, H_f - mask.shape[0]), max(0, W_f - mask.shape[1])
if pad_mh > 0 or pad_mw > 0:
mask = np.pad(mask, ((0, pad_mh), (0, pad_mw), (0, 0)), mode='constant')
H, W = feat.shape[:2]
feat = np.where((feat == -9999) | ~np.isfinite(feat), 0, feat).astype(np.float32)
feat = (feat - NORM_MEAN) / (NORM_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)), mode='constant')
mask = np.pad(mask, ((0, pad_h), (0, pad_w), (0, 0)), mode='constant')
h, w = feat.shape[:2]
if h > th or w > tw: # center-crop large images
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 (real, non-padded) region inside the 128x128 canvas
valid_h, valid_w = min(H, th), min(W, tw)
return (feat.astype(np.float32), mask.astype(np.float32),
H, W, valid_h, valid_w, pad_h > 0 or pad_w > 0, H > th or W > tw)
# ─────────────────────────────────────────────────────────────────────────────
# Build model + load weights
# ─────────────────────────────────────────────────────────────────────────────
print("Building model…")
_MODEL_TYPE = os.environ.get("MODEL_NAME", "swin_unet")
if _MODEL_TYPE == "attention_unet":
model = build_attention_unet(input_shape=INPUT_SHAPE, num_classes=1)
elif _MODEL_TYPE == "unet":
model = build_unet(input_shape=INPUT_SHAPE, num_classes=1)
elif _MODEL_TYPE == "segformer":
model = build_segformer(input_shape=INPUT_SHAPE, num_classes=1)
else:
model = build_model(input_shape=INPUT_SHAPE, num_classes=1, num_heads=8)
print(f"[eval] model type: {_MODEL_TYPE}")
model.compile(loss=combined_loss, metrics=['accuracy', dice_coef])
if not os.path.exists(CKPT_PATH):
raise FileNotFoundError(f"Checkpoint not found: {CKPT_PATH}. Train first (python train.py).")
model.load_weights(CKPT_PATH)
print("Weights loaded.")
# ─────────────────────────────────────────────────────────────────────────────
# Load + preprocess every test image
# ─────────────────────────────────────────────────────────────────────────────
# Detect .npy patches (new pipeline) or .tif files (legacy)
npy_files = sorted(glob.glob(os.path.join(TEST_FEAT, '*.npy')))
is_npy = len(npy_files) > 0
if is_npy:
feat_files = npy_files
mask_files = [os.path.join(TEST_MASK, os.path.basename(f)) for f in feat_files]
for mf in mask_files:
assert os.path.exists(mf), f"Missing mask: {mf}"
print(f"\nLoading {len(feat_files)} test patches (.npy)…")
else:
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) == len(mask_files) and len(feat_files) > 0, "Test feature/mask mismatch."
print(f"\nLoading {len(feat_files)} test images (.tif)…")
X, Y, meta, ids = [], [], [], []
for fp, mp in zip(feat_files, mask_files):
if is_npy:
feat = np.load(fp).astype(np.float32) # (H, W, num_bands)
mask = np.load(mp).astype(np.float32) # (H, W, 1)
H, W = feat.shape[:2]
feat = np.where((feat == -9999) | ~np.isfinite(feat), 0, feat)
if BAND_INDICES is not None:
feat = feat[..., BAND_INDICES]
feat = (feat - NORM_MEAN) / (NORM_STD + 1e-7)
meta.append((H, W, H, W, False, False))
ids.append(os.path.basename(fp).replace('.npy', ''))
else:
feat, mask, H, W, vh, vw, padded, cropped = load_and_prep(fp, mp)
meta.append((H, W, vh, vw, padded, cropped))
ids.append(os.path.basename(fp).replace('feature_', '').replace('.tif', ''))
X.append(feat)
Y.append(mask)
X = np.array(X, dtype=np.float32) # (N, 128, 128, 12)
Y = np.array(Y, dtype=np.float32) # (N, 128, 128, 1)
def predict_tta(model, X, batch_size):
"""Average sigmoid probs over the 8 dihedral (D4) transforms."""
acc = np.zeros((X.shape[0], X.shape[1], X.shape[2], 1), np.float32)
n = 0
for flip in (False, True):
for k in (0, 1, 2, 3):
Xt = X[:, :, ::-1, :] if flip else X
Xt = np.ascontiguousarray(np.rot90(Xt, k, axes=(1, 2)))
p = model.predict(Xt, batch_size=batch_size, verbose=0)
p = np.rot90(p, -k, axes=(1, 2))
if flip:
p = p[:, :, ::-1, :]
acc += p
n += 1
return acc / n
if USE_TTA:
print("Predicting (8-way TTA: 4 rotations x 2 flips)…")
Y_prob = predict_tta(model, X, BATCH_SIZE) # (N, 128, 128, 1)
else:
print("Predicting…")
Y_prob = model.predict(X, batch_size=BATCH_SIZE, verbose=1) # (N, 128, 128, 1)
# ─────────────────────────────────────────────────────────────────────────────
# Collect VALID-region pixels (drop zero-padding) for global + per-image metrics
# ─────────────────────────────────────────────────────────────────────────────
def metrics_from_cm(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 (worst IoU first)
# ─────────────────────────────────────────────────────────────────────────────
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 = metrics_from_cm(tp, fp, fn, tn)
# Empty GT + empty prediction = perfect agreement (correct definition;
# otherwise background-only tiles wrongly score IoU=0 and drag the mean down).
if (tp + fp + fn) == 0:
prec = rec = f1 = iou = dice = 1.0
n_px = yt.size
per_image.append({
'image_id': iid,
'orig_H': H,
'orig_W': W,
'valid_px': int(n_px),
'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),
'mean_pred_prob': round(float(yprob.mean()), 4),
'max_pred_prob': round(float(yprob.max()), 4),
})
per_image_df = pd.DataFrame(per_image).sort_values('iou', ascending=True).reset_index(drop=True)
per_image_df.to_csv(os.path.join(OUT_DIR, 'per_image_metrics.csv'), index=False)
# ─────────────────────────────────────────────────────────────────────────────
# Global threshold sweep (valid pixels only)
# ─────────────────────────────────────────────────────────────────────────────
y_true_flat = np.concatenate(all_true)
y_prob_flat = np.concatenate(all_prob)
rows = []
for t in THRESHOLDS:
yp = (y_prob_flat >= t).astype(np.int32)
tn, fp, fn, tp = confusion_matrix(y_true_flat, yp, labels=[0, 1]).ravel()
prec, rec, f1, iou, dice, acc, spec = metrics_from_cm(tp, fp, fn, tn)
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(rows)
thresh_df.to_csv(os.path.join(OUT_DIR, 'evaluation_metrics.csv'), index=False)
best_row = thresh_df.loc[thresh_df['iou'].idxmax()]
best_t = float(best_row['threshold'])
cm = confusion_matrix(y_true_flat, (y_prob_flat >= 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, 'confusion_matrix.csv'))
# ─────────────────────────────────────────────────────────────────────────────
# Summary
# ─────────────────────────────────────────────────────────────────────────────
mean_iou = per_image_df['iou'].mean()
median_iou = per_image_df['iou'].median()
worst = per_image_df.head(10)[['image_id', 'orig_H', 'orig_W', 'glacier_frac_gt', 'iou', 'precision', 'recall']]
# Machine-readable metrics for results_table.py (reliable, no text parsing).
with open(os.path.join(OUT_DIR, 'metrics.json'), 'w') as f:
json.dump({
'run_name': RUN_NAME,
'n_bands': int(INPUT_SHAPE[2]),
'patch': [int(_ph), int(_pw)],
'best_threshold': best_t,
'test_precision': round(float(best_row['precision']), 4),
'test_recall': round(float(best_row['recall']), 4),
'test_f1': round(float(best_row['f1_score']), 4),
'test_iou': round(float(best_row['iou']), 4),
'test_dice': round(float(best_row['dice']), 4),
'per_image_mean_iou': round(float(mean_iou), 4),
}, f, indent=2)
lines = [
"=" * 60,
" ROCKGLACIERNET β€” EVALUATION SUMMARY (baseline 128x128)",
"=" * 60,
"",
f"Test images : {len(ids)}",
f"Best threshold by IoU : {best_t}",
"",
"GLOBAL (pixel-pooled) metrics @ best threshold:",
f" precision : {best_row['precision']:.4f}",
f" recall : {best_row['recall']:.4f}",
f" f1 : {best_row['f1_score']:.4f}",
f" iou : {best_row['iou']:.4f}",
f" dice : {best_row['dice']:.4f}",
"",
"PER-IMAGE IoU (@ t=0.5):",
f" mean : {mean_iou:.4f}",
f" median : {median_iou:.4f}",
"",
"Threshold sweep:",
" {:<8} {:<10} {:<10} {:<10} {:<10} {:<10}".format(
"thresh", "precision", "recall", "f1", "iou", "dice"),
" " + "-" * 60,
]
for _, r in thresh_df.iterrows():
tag = " <- best IoU" if r['threshold'] == best_t else ""
lines.append(" {:<8} {:<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 += ["", "10 WORST images by IoU (@ t=0.5):", worst.to_string(index=False)]
lines += [
"",
f"Output dir: {OUT_DIR}",
"Files: per_image_metrics.csv, evaluation_metrics.csv, confusion_matrix.csv",
]
summary = "\n".join(lines)
print("\n" + summary)
with open(os.path.join(OUT_DIR, 'evaluation_summary.txt'), 'w') as f:
f.write(summary)
# ─────────────────────────────────────────────────────────────────────────────
# Prediction comparison PNGs (10 worst + 10 best)
# ─────────────────────────────────────────────────────────────────────────────
try:
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
id_to_idx = {iid: i for i, iid in enumerate(ids)}
to_plot = list(per_image_df.head(10)['image_id']) + list(per_image_df.tail(10)['image_id'])
for iid in to_plot:
i = id_to_idx[iid]
H, W, vh, vw, padded, cropped = meta[i]
fig, axes = plt.subplots(1, 3, figsize=(9, 3))
rgb = X[i, :vh, :vw, :][:, :, [2, 1, 0]]
rgb = (rgb - rgb.min()) / (rgb.max() - rgb.min() + 1e-8)
iou_val = float(per_image_df.loc[per_image_df['image_id'] == iid, 'iou'].iloc[0])
axes[0].imshow(rgb); axes[0].set_title(f'{iid} RGB'); axes[0].axis('off')
axes[1].imshow(Y[i, :vh, :vw, 0], cmap='gray'); axes[1].set_title('GT'); axes[1].axis('off')
axes[2].imshow((Y_prob[i, :vh, :vw, 0] >= best_t), cmap='gray')
axes[2].set_title(f'Pred IoU={iou_val:.2f}'); axes[2].axis('off')
plt.tight_layout()
plt.savefig(os.path.join(OUT_DIR, f'pred_{iid}.png'), dpi=80)
plt.close()
print(f"Saved {len(to_plot)} comparison PNGs (10 worst + 10 best).")
except Exception as e:
print(f"PNG save skipped: {e}")
print("\nEvaluation complete.")