""" Publication-quality figures for RockGlacierNet (Swin-UNet, 8-way TTA). Per test patch -> paper_figures/test_cases/.png : a georeferenced 5-panel figure (a) RGB (b) Ground truth (c) Prediction (d) Probability (e) Error map with lat/lon axes, scale bar, north arrow, panel labels, colourbar and legend. Also: paper_figures/ranking.csv - test patches ranked by Dice paper_figures/montage_best.png - 3x3 RGB gallery with GT (yellow) + pred (cyan) outlines paper_figures/montage_worst.png paper_figures/training_history.png Run on CPU to avoid GPU contention: CUDA_VISIBLE_DEVICES="" python3 make_paper_figures.py """ import os, glob, json, csv import numpy as np import matplotlib; matplotlib.use("Agg") import matplotlib.pyplot as plt from matplotlib.colors import ListedColormap from matplotlib.ticker import MaxNLocator, FuncFormatter from matplotlib.patches import FancyArrow, Rectangle from matplotlib.lines import Line2D import tensorflow as tf from models.swin_unet import build_model # ---------------- config ---------------- DATA = "dataset/patches_v2" POLY_DIR = "EDA/individual_polygons" CKPT_DIR = os.environ.get("CKPT_DIR", "bandselection/run_v2_cosine250") CKPT = f"{CKPT_DIR}/best_model.weights.h5" OUT = "paper_figures" BAND_IDX = [25,18,22,23,10,6,24,27,17,11,13,14,4,3,9,26,8,5,15,7] RGB_IDX = [2,1,0] THRESH = float(os.environ.get("THRESH", 0.35)) PATCH_M = 1280.0 # 128 px * 10 m UTM_CRS = "EPSG:32645" plt.rcParams.update({ "font.size": 9, "font.family": "serif", "axes.linewidth": 0.7, "xtick.major.size": 3, "ytick.major.size": 3, "figure.dpi": 120, }) os.makedirs(f"{OUT}/test_cases", exist_ok=True) # ---------------- geo lookup: polygon centroid -> patch extent ---------------- def build_centroids(): try: import geopandas as gpd from pyproj import Transformer except Exception as e: print("[warn] geopandas/pyproj unavailable -> km axes fallback:", e) return None, None cents = {} for f in glob.glob(os.path.join(POLY_DIR, "*.geojson")): gid = os.path.splitext(os.path.basename(f))[0] try: g = gpd.read_file(f).to_crs(UTM_CRS) g["geometry"] = g.geometry.buffer(0) c = g.unary_union.centroid cents[gid] = (c.x, c.y) except Exception: continue tf_ = Transformer.from_crs(UTM_CRS, "EPSG:4326", always_xy=True) return cents, tf_ CENTROIDS, TF_ = build_centroids() def geo_extent(pid): """Return (extent_lonlat, ok). extent = [lon_w, lon_e, lat_s, lat_n].""" if CENTROIDS and pid in CENTROIDS: cx, cy = CENTROIDS[pid] lon_w, lat_s = TF_.transform(cx - PATCH_M/2, cy - PATCH_M/2) lon_e, lat_n = TF_.transform(cx + PATCH_M/2, cy + PATCH_M/2) return [lon_w, lon_e, lat_s, lat_n], True return [0, PATCH_M/1000.0, 0, PATCH_M/1000.0], False # km fallback # ---------------- model ---------------- print("Building model + loading weights...") model = build_model(input_shape=(128,128,len(BAND_IDX)), num_classes=1, num_heads=8) model.load_weights(CKPT) st = json.load(open(f"{CKPT_DIR}/norm_stats.json")) MEAN = np.array(st["mean"], np.float32).reshape(1,1,-1) STD = np.array(st["std"], np.float32).reshape(1,1,-1) def predict_tta(x): acc = np.zeros((1,128,128,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, verbose=0); p = np.rot90(p, -k, axes=(1,2)) if flip: p = p[:, :, ::-1, :] acc += p; n += 1 return acc/n def stretch(rgb): o = np.zeros_like(rgb, np.float32) for c in range(3): b = rgb[:,:,c]; lo,hi = np.percentile(b,2), np.percentile(b,98) o[:,:,c] = np.clip((b-lo)/(hi-lo+1e-8),0,1) return o def deg_fmt(v, _): return f"{abs(v):.3f}°" def add_geo_axes(ax, extent, ok, show_labels): ax.set_xlim(extent[0], extent[1]); ax.set_ylim(extent[2], extent[3]) ax.xaxis.set_major_locator(MaxNLocator(3)); ax.yaxis.set_major_locator(MaxNLocator(3)) if ok: ax.xaxis.set_major_formatter(FuncFormatter(deg_fmt)) ax.yaxis.set_major_formatter(FuncFormatter(deg_fmt)) if show_labels: ax.set_xlabel("Longitude (°E)"); ax.set_ylabel("Latitude (°N)") ax.tick_params(labelsize=7) else: ax.set_xticklabels([]); ax.set_yticklabels([]) ax.tick_params(length=2) def add_scalebar(ax, extent, ok): span = extent[1]-extent[0] bar = span*(500.0/PATCH_M) if ok else 0.5 # 500 m (or 0.5 km fallback) x0 = extent[0] + span*0.06; y0 = extent[2] + (extent[3]-extent[2])*0.06 ax.add_line(Line2D([x0,x0+bar],[y0,y0],color="white",lw=3,solid_capstyle="butt")) ax.text(x0+bar/2, y0+(extent[3]-extent[2])*0.03, "500 m", color="white", ha="center", va="bottom", fontsize=6.5, weight="bold") def add_north(ax, extent): x = extent[1]-(extent[1]-extent[0])*0.08; y0 = extent[3]-(extent[3]-extent[2])*0.22 dy = (extent[3]-extent[2])*0.12 ax.add_patch(FancyArrow(x, y0, 0, dy, width=0, head_width=(extent[1]-extent[0])*0.03, head_length=dy*0.5, color="white", length_includes_head=True)) ax.text(x, y0+dy+ (extent[3]-extent[2])*0.02, "N", color="white", ha="center", va="bottom", fontsize=8, weight="bold") def panel_label(ax, s): ax.text(0.03,0.95,s, transform=ax.transAxes, fontsize=10, weight="bold", color="white", va="top", ha="left", bbox=dict(boxstyle="round,pad=0.15", fc="black", ec="none", alpha=0.55)) # ---------------- render ---------------- err_cmap = ListedColormap([[0,0,0,0],[0.15,0.8,0.15,1],[0.9,0.15,0.15,1],[0.15,0.45,0.95,1]]) feat_files = sorted(glob.glob(f"{DATA}/test/features/*.npy")) rows = [] print(f"Rendering {len(feat_files)} test cases (geo-referenced)...") for fp in feat_files: pid = os.path.basename(fp).replace(".npy","") full = np.load(fp).astype(np.float32) mask = np.load(fp.replace("/features/","/masks/")).squeeze().astype(np.int32) ext, ok = geo_extent(pid) rgb = stretch(full[:,:,RGB_IDX]) x = full[..., BAND_IDX]; nod=(x==-9999)|~np.isfinite(x) x = np.where(nod,0.0,x); x=(x-MEAN)/(STD+1e-7); x=np.where(nod,0.0,x)[None,...] prob = predict_tta(x)[0,:,:,0]; pred=(prob>=THRESH).astype(np.int32) tp=int(((pred==1)&(mask==1)).sum()); fp=int(((pred==1)&(mask==0)).sum()) fn=int(((pred==0)&(mask==1)).sum()) dice=2*tp/(2*tp+fp+fn+1e-9); iou=tp/(tp+fp+fn+1e-9) prec=tp/(tp+fp+1e-9); rec=tp/(tp+fn+1e-9) rows.append((pid,round(dice,4),round(iou,4),round(prec,4),round(rec,4),round(float(mask.mean()),4))) err=np.zeros_like(mask); err[(pred==1)&(mask==1)]=1; err[(pred==1)&(mask==0)]=2; err[(pred==0)&(mask==1)]=3 imk=dict(extent=ext, origin="upper", aspect="equal") fig,ax=plt.subplots(1,5,figsize=(15,3.5)) ax[0].imshow(rgb, **imk); add_scalebar(ax[0],ext,ok); add_north(ax[0],ext) ax[0].set_title("(a) RGB (B4/B3/B2)") ax[1].imshow(mask,cmap="Greys_r",vmin=0,vmax=1,**imk); ax[1].set_title("(b) Ground truth") ax[2].imshow(pred,cmap="Greys_r",vmin=0,vmax=1,**imk); ax[2].set_title(f"(c) Prediction (t={THRESH})") im=ax[3].imshow(prob,cmap="magma",vmin=0,vmax=1,**imk); ax[3].set_title("(d) Probability") ax[4].imshow(rgb,**imk); ax[4].imshow(err,cmap=err_cmap,vmin=0,vmax=3,**imk); ax[4].set_title("(e) Error map") add_geo_axes(ax[0],ext,ok,show_labels=True) for a in ax[1:]: add_geo_axes(a,ext,ok,show_labels=False) for a,l in zip(ax,"abcde"): panel_label(a,f"({l})") if False else None # titles already labelled cb=fig.colorbar(im, ax=ax[3], fraction=0.046, pad=0.03); cb.set_label("P(glacier)",fontsize=7); cb.ax.tick_params(labelsize=6) leg=[Line2D([0],[0],marker="s",ls="",mfc=(0.15,0.8,0.15),mec="none",label="TP"), Line2D([0],[0],marker="s",ls="",mfc=(0.9,0.15,0.15),mec="none",label="FP"), Line2D([0],[0],marker="s",ls="",mfc=(0.15,0.45,0.95),mec="none",label="FN")] ax[4].legend(handles=leg, loc="lower right", fontsize=6, framealpha=0.85, handletextpad=0.3, borderpad=0.3) fig.suptitle(f"{pid} Dice = {dice:.3f} IoU = {iou:.3f} Precision = {prec:.3f} Recall = {rec:.3f}", fontsize=11, y=1.03) plt.tight_layout() plt.savefig(f"{OUT}/test_cases/{pid}.png", dpi=200, bbox_inches="tight"); plt.close() rows.sort(key=lambda r:-r[1]) with open(f"{OUT}/ranking.csv","w") as f: f.write("patch_id,dice,iou,precision,recall,glacier_frac\n") for r in rows: f.write(",".join(map(str,r))+"\n") print(f"Saved {len(rows)} georeferenced figures + ranking.csv") # ---------------- grid-axis helper (lat/lon on outer panels only) ---------------- def grid_axes(a, ext, ok, show_x, show_y): a.set_xlim(ext[0],ext[1]); a.set_ylim(ext[2],ext[3]) a.xaxis.set_major_locator(MaxNLocator(3)); a.yaxis.set_major_locator(MaxNLocator(3)) if ok: a.xaxis.set_major_formatter(FuncFormatter(deg_fmt)) a.yaxis.set_major_formatter(FuncFormatter(deg_fmt)) a.tick_params(labelsize=6, length=2) if not show_x: a.set_xticklabels([]) if not show_y: a.set_yticklabels([]) if show_x: a.set_xlabel("Lon (°E)", fontsize=7) if show_y: a.set_ylabel("Lat (°N)", fontsize=7) def _pred(pid): full=np.load(f"{DATA}/test/features/{pid}.npy").astype(np.float32) mask=np.load(f"{DATA}/test/masks/{pid}.npy").squeeze().astype(np.int32) x=full[...,BAND_IDX]; nod=(x==-9999)|~np.isfinite(x) xx=np.where(nod,0.0,x); xx=(xx-MEAN)/(STD+1e-7); xx=np.where(nod,0.0,xx)[None,...] prob=predict_tta(xx)[0,:,:,0] return full, mask, prob, (prob>=THRESH).astype(np.int32) # ---------------- montage: RGB + GT (yellow) + pred (cyan) outlines ---------------- def montage(cases, title, out): fig,axes=plt.subplots(3,3,figsize=(11,11.8)) for idx,(a,(pid,dice,*_)) in enumerate(zip(axes.ravel(),cases)): r,c=divmod(idx,3) full,mask,_,pr=_pred(pid); ext,ok=geo_extent(pid) a.imshow(stretch(full[:,:,RGB_IDX]),extent=ext,origin="upper",aspect="equal") a.contour(mask,levels=[0.5],colors="yellow",linewidths=1.3,extent=ext,origin="upper") a.contour(pr, levels=[0.5],colors="cyan", linewidths=1.1,extent=ext,origin="upper") grid_axes(a,ext,ok, show_x=(r==2), show_y=(c==0)) a.set_title(f"({chr(97+idx)}) Dice = {dice:.2f}", fontsize=9) for a in axes.ravel()[len(cases):]: a.axis("off") fig.legend(handles=[Line2D([0],[0],color="gold",lw=2.5,label="Ground truth"), Line2D([0],[0],color="cyan",lw=2.5,label="Prediction")], loc="lower center", ncol=2, fontsize=10, framealpha=0.95, bbox_to_anchor=(0.5, 0.005)) fig.suptitle(title, fontsize=13, y=0.98) plt.tight_layout(rect=[0,0.035,1,0.955]) plt.savefig(out,dpi=200,bbox_inches="tight"); plt.close() # ---------------- top-5 detailed: 5 rows x 5 panels ---------------- def top5_detailed(cases, out): n=len(cases) fig,axes=plt.subplots(n,5,figsize=(15, 3.05*n)) cols=["RGB (B4/B3/B2)","Ground truth",f"Prediction (t={THRESH})","Probability","Error map"] im=None for i,(pid,dice,iou,prec,rec,gf) in enumerate(cases): full,mask,prob,pred=_pred(pid); ext,ok=geo_extent(pid) rgb=stretch(full[:,:,RGB_IDX]) err=np.zeros_like(mask); err[(pred==1)&(mask==1)]=1; err[(pred==1)&(mask==0)]=2; err[(pred==0)&(mask==1)]=3 row=axes[i] if n>1 else axes row[0].imshow(rgb,extent=ext,origin="upper",aspect="equal"); add_scalebar(row[0],ext,ok); add_north(row[0],ext) row[1].imshow(mask,cmap="Greys_r",vmin=0,vmax=1,extent=ext,origin="upper",aspect="equal") row[2].imshow(pred,cmap="Greys_r",vmin=0,vmax=1,extent=ext,origin="upper",aspect="equal") im=row[3].imshow(prob,cmap="magma",vmin=0,vmax=1,extent=ext,origin="upper",aspect="equal") row[4].imshow(rgb,extent=ext,origin="upper",aspect="equal") row[4].imshow(err,cmap=err_cmap,vmin=0,vmax=3,extent=ext,origin="upper",aspect="equal") for j,a in enumerate(row): grid_axes(a,ext,ok, show_x=(i==n-1), show_y=(j==0)) if i==0: a.set_title(cols[j], fontsize=9.5) row[0].set_ylabel(f"({chr(97+i)}) Dice = {dice:.2f}", fontsize=9) leg=[Line2D([0],[0],marker="s",ls="",mfc=(0.15,0.8,0.15),mec="none",label="TP"), Line2D([0],[0],marker="s",ls="",mfc=(0.9,0.15,0.15),mec="none",label="FP"), Line2D([0],[0],marker="s",ls="",mfc=(0.15,0.45,0.95),mec="none",label="FN")] (axes[0] if n>1 else axes)[4].legend(handles=leg, loc="lower right", fontsize=6.5, framealpha=0.85, handletextpad=0.3, borderpad=0.3) plt.tight_layout() plt.savefig(out,dpi=200,bbox_inches="tight"); plt.close() glac=[r for r in rows if r[5]>0.01] if len(glac)>=9: montage(glac[:9], "Swin-UNet prediction vs. ground truth — best test cases", f"{OUT}/montage_best.png") montage(glac[-9:],"Swin-UNet prediction vs. ground truth — hardest test cases", f"{OUT}/montage_worst.png") print("Saved montage_best.png / montage_worst.png") if len(glac)>=5: top5_detailed(glac[:5], f"{OUT}/top5_detailed.png") print("Saved top5_detailed.png") # ---------------- training history ---------------- hp=f"{CKPT_DIR}/training_history.csv" if os.path.exists(hp): H=[r for r in csv.DictReader(open(hp)) if int(float(r['epoch']))<10000] ep=[int(float(r['epoch'])) for r in H] fig,ax=plt.subplots(1,2,figsize=(11,3.8)) ax[0].plot(ep,[float(r['dice_coef']) for r in H],lw=1.4,label="Train") ax[0].plot(ep,[float(r['val_dice_coef']) for r in H],lw=1.4,label="Validation") ax[0].set_xlabel("Epoch"); ax[0].set_ylabel("Dice coefficient"); ax[0].legend(frameon=False) ax[0].set_title("(a) Dice"); ax[0].grid(alpha=0.3) ax[1].plot(ep,[float(r['loss']) for r in H],lw=1.4,label="Train") ax[1].plot(ep,[float(r['val_loss']) for r in H],lw=1.4,label="Validation") ax[1].set_xlabel("Epoch"); ax[1].set_ylabel("Loss"); ax[1].legend(frameon=False) ax[1].set_title("(b) Loss"); ax[1].grid(alpha=0.3) plt.tight_layout(); plt.savefig(f"{OUT}/training_history.png",dpi=200,bbox_inches="tight"); plt.close() print("Saved training_history.png") print("\nDONE. Figures in", OUT, "\nPick patches from ranking.csv and tell me which to use.")