| """ |
| Publication-quality figures for RockGlacierNet (Swin-UNet, 8-way TTA). |
| |
| Per test patch -> paper_figures/test_cases/<id>.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 |
|
|
| |
| 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 |
| 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) |
|
|
| |
| 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 |
|
|
| |
| 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 |
| 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)) |
|
|
| |
| 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 |
|
|
| 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") |
|
|
| |
| 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) |
|
|
| |
| 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() |
|
|
| |
| 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") |
|
|
| |
| 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.") |
|
|
