Blame for make_paper_figures.py

"""
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

# ---------------- 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.")