| """Synthetic sanity check: does the CHMv2 height-drop cue improve the 0.5776 ensemble's tree_removal score? |
| |
| usage: python chm_eval.py N (run in /workspace with pkg/assets populated) |
| """ |
| import glob |
| import sys |
|
|
| import cv2 |
| import numpy as np |
| import torch |
|
|
| N = int(sys.argv[1]) if len(sys.argv) > 1 else 600 |
| sys.argv = sys.argv[:1] |
| sys.path.insert(0, "/workspace/pkg/assets") |
| import segmentation_models_pytorch as smp |
| from model_v2 import SiamCD |
| from chm_fuse import CHM, cue_abs, cue_diff |
| from train3b import Synth3 |
|
|
| dev = "cuda" |
| MEAN = np.array([0.485, 0.456, 0.406], np.float32); STD = np.array([0.229, 0.224, 0.225], np.float32) |
| u = smp.Unet(encoder_name="resnet18", encoder_weights=None, in_channels=6, classes=3) |
| u.load_state_dict(torch.load("/workspace/pkg/assets/model/unet_synth.pt", map_location="cpu", weights_only=True)["state_dict"]) |
| u = u.to(dev).eval() |
| sm = [] |
| for f in sorted(glob.glob("/workspace/pkg/assets/model/siam_*.pt")): |
| s = SiamCD(); s.load_state_dict(torch.load(f, map_location="cpu", weights_only=True)["state_dict"]); sm.append(s.to(dev).eval()) |
| print("satlas models", len(sm), flush=True) |
| chms = {s: CHM("/workspace/pkg/assets/chmv2", dev, scale=s) for s in (1.0, 1.5)} |
|
|
|
|
| @torch.no_grad() |
| def ens(pres, posts): |
| a = torch.from_numpy(np.stack(pres)).permute(0, 3, 1, 2).to(dev).float() / 255 |
| b = torch.from_numpy(np.stack(posts)).permute(0, 3, 1, 2).to(dev).float() / 255 |
| m = torch.tensor(MEAN, device=dev).view(1, 3, 1, 1); sd = torch.tensor(STD, device=dev).view(1, 3, 1, 1) |
| prob = 0 |
| for dims in ([], [3], [2], [2, 3]): |
| ai, bi = (torch.flip(a, dims), torch.flip(b, dims)) if dims else (a, b) |
| pu = torch.softmax(u(torch.cat([(ai - m) / sd, (bi - m) / sd], 1)), 1)[:, 1:3] |
| cs = 0 |
| for sk in sm: |
| with torch.autocast("cuda", dtype=torch.float16): |
| ch, _, _, _ = sk(ai, bi) |
| cs = cs + ch.float().sigmoid() / len(sm) |
| ci = 0.5 * pu.float() + 0.5 * cs |
| prob = prob + (torch.flip(ci, dims) if dims else ci) / 4 |
| return prob |
|
|
|
|
| ds = Synth3(10**7) |
| P, HP, HQ, HP2, HQ2, LB, LT, SEM, ND = ([] for _ in range(9)) |
| for s in range(0, N, 16): |
| items = [ds[i] for i in range(s, min(s + 16, N))] |
| pres = [it[0].numpy().transpose(1, 2, 0).copy() for it in items] |
| posts = [it[1].numpy().transpose(1, 2, 0).copy() for it in items] |
| P.append(ens(pres, posts).cpu().numpy()) |
| HP.append(chms[1.0].heights(pres)); HQ.append(chms[1.0].heights(posts)) |
| HP2.append(chms[1.5].heights(pres)); HQ2.append(chms[1.5].heights(posts)) |
| LB += [it[2][0].numpy() > 0.5 for it in items]; LT += [it[2][1].numpy() > 0.5 for it in items] |
| SEM += [it[4].numpy() for it in items] |
| ND += [(a.max(2) == 0) | (b.max(2) == 0) for a, b in zip(pres, posts)] |
| P = np.concatenate(P); HP = torch.cat(HP); HQ = torch.cat(HQ); HP2 = torch.cat(HP2); HQ2 = torch.cat(HQ2); ND = np.stack(ND) |
| P[:, 0][ND] = 0; P[:, 1][ND] = 0 |
| sem = np.stack(SEM); lt = np.stack(LT); hp = HP.cpu().numpy(); hq = HQ.cpu().numpy() |
| print(f"CHM metres pre tree px {hp[sem == 2].mean():.2f} pre non-tree px {hp[(sem != 2) & (sem != 255)].mean():.2f} " |
| f"removed px pre {hp[lt].mean():.2f} post {hq[lt].mean():.2f} (x1.5: tree {HP2.cpu().numpy()[sem == 2].mean():.2f} removed pre {HP2.cpu().numpy()[lt].mean():.2f} post {HQ2.cpu().numpy()[lt].mean():.2f}) overall p50/p90/p99 {np.percentile(hp, [50, 90, 99]).round(2).tolist()}", flush=True) |
|
|
|
|
| def clean(m, amin=30): |
| n, cc, st, _ = cv2.connectedComponentsWithStats(m.astype(np.uint8), 8) |
| keep = np.zeros(n, bool); keep[1:] = st[1:, cv2.CC_STAT_AREA] >= amin |
| return keep[cc] |
|
|
|
|
| def score(prob, gts): |
| k3 = np.ones((3, 3), np.uint8); tp = fp = fn = tn = 0; shp = [] |
| for p, g in zip(prob, gts): |
| pm = clean(p > 0.5); pp, gp = pm.sum() >= 20, g.sum() >= 20 |
| tp += pp and gp; fp += pp and not gp; fn += gp and not pp; tn += (not pp) and (not gp) |
| if gp: |
| if not pp: |
| shp.append(0.0); continue |
| gd = cv2.dilate(g.astype(np.uint8), k3).astype(bool); pd = cv2.dilate(pm.astype(np.uint8), k3).astype(bool) |
| pr = (pm & gd).sum() / pm.sum(); rc = (g & pd).sum() / g.sum() |
| shp.append(2 * pr * rc / (pr + rc + 1e-9)) |
| f1p = 2 * tp / max(2 * tp + fp + fn, 1); f1n = 2 * tn / max(2 * tn + fn + fp, 1) |
| return 0.5 * (f1p + f1n) / 2 + 0.5 * (float(np.mean(shp)) if shp else 0.0), (tp, fp, fn, tn) |
|
|
|
|
| np.savez("/workspace/eval_arr.npz", pt=P[:, 1].astype(np.float16), hp=HP.cpu().numpy().astype(np.float16), hq=HQ.cpu().numpy().astype(np.float16), hp2=HP2.cpu().numpy().astype(np.float16), hq2=HQ2.cpu().numpy().astype(np.float16), lt=lt, nd=ND) |
| print("SAVED", flush=True) |
| sb, cb = score(P[:, 0], LB) |
| print(f"building {sb:.4f} {cb}", flush=True) |
| cues = {"abs1.0": cue_abs(HP, HQ), "diff1.0": cue_diff(HP, HQ), "abs1.5": cue_abs(HP2, HQ2), "diff1.5": cue_diff(HP2, HQ2)} |
| res = {} |
| for name, c in cues.items(): |
| c = c.cpu().numpy(); c[ND] = 0 |
| for a in (0.0, 0.2, 0.3, 0.4, 0.5, 1.0): |
| st, ct = score((1 - a) * P[:, 1] + a * c, LT) |
| res[(name, a)] = st |
| print(f"tree cue={name} a={a:.1f} {st:.4f} tp/fp/fn/tn {ct}", flush=True) |
| print("EVAL_DONE", flush=True) |
|
|
