File size: 4,517 Bytes
4b20d1b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 | """Re-score a Blue Marble Flood model on Sen1Floods11: the official hand-labelled test split and the held-out Bolivia split.
pip install numpy onnxruntime rasterio huggingface_hub
python evaluate.py --model nano/model_fp32.onnx # or nano/model_int8.onnx
Data: Sen1Floods11 v1.1 (hand-labelled Sentinel-2 L1C chips and labels, official splits), about 1.8 GB, downloaded once."""
import numpy as np
# Sentinel-2 L1C, 13 bands (B01, B02, B03, B04, B05, B06, B07, B08, B8A, B09, B10, B11, B12)
MEAN = np.array([2357.089, 2137.385, 2018.788, 2082.986, 2295.651, 2854.537, 3122.849, 3040.560, 3306.481, 1473.847, 506.070, 2472.825, 1838.929], np.float32)
STD = np.array([1624.683, 1675.806, 1557.708, 1833.702, 1823.738, 1733.977, 1732.131, 1679.732, 1727.26, 1024.687, 442.165, 1331.411, 1160.419], np.float32)
def normalize(dn, bands=None):
"""dn: [C, H, W] digital numbers for the 13 bands, or for the band indices in `bands`. Missing bands are 0 after normalisation."""
out = np.zeros((13, *dn.shape[1:]), np.float32)
idx = list(range(13)) if bands is None else list(bands)
out[idx] = (dn.astype(np.float32) - MEAN[idx, None, None]) / STD[idx, None, None]
return out
def session(path, threads=0):
import onnxruntime as ort
so = ort.SessionOptions()
if threads: so.intra_op_num_threads = threads
prov = [p for p in ("CUDAExecutionProvider", "CPUExecutionProvider") if p in ort.get_available_providers()]
return ort.InferenceSession(path, so, providers=prov)
def predict(sess, x):
"""Whole image in one pass: pad H and W to multiples of 16 by edge replication, crop the logits back."""
h, w = x.shape[-2:]
xp = np.pad(x, ((0, 0), (0, (-h) % 16), (0, (-w) % 16)), mode="edge")[None].astype(np.float32)
return sess.run(None, {sess.get_inputs()[0].name: xp})[0][0, :, :h, :w].argmax(0)
class MIoU:
def __init__(self, classes):
self.classes, n = classes, len(classes)
self.cm = np.zeros((n, n), np.int64)
def add(self, y, p):
n = len(self.classes); m = y >= 0
self.cm += np.bincount(n * y[m].astype(np.int64) + p[m], minlength=n * n).reshape(n, n)
def report(self, title):
inter = np.diag(self.cm).astype(float); union = self.cm.sum(0) + self.cm.sum(1) - inter
iou = inter / np.maximum(union, 1)
print(f"{title}: mIoU {100 * iou.mean():.2f} (" + ", ".join(f"{c} {100 * v:.2f}" for c, v in zip(self.classes, iou)) + ")", flush=True)
return float(iou.mean())
import argparse, os, tarfile
import numpy as np, rasterio
SPLITS = {"test": "flood_test_data", "bolivia": "flood_bolivia_data"}
def fetch(data):
root = f"{data}/sen1floods11_v1.1"
if os.path.exists(f"{root}/.ok"): return root
from huggingface_hub import hf_hub_download
arc = f"{data}/sen1floods11_v1_1.tar.gz"
if not os.path.exists(arc): arc = hf_hub_download("blumenstiel/Sen1Floods11", "sen1floods11_v1_1.tar.gz", repo_type="dataset", local_dir=data)
with tarfile.open(arc, "r|gz") as t: # pass 1: split lists
for m in t:
if "/splits/" in m.name and m.isfile(): t.extract(m, data, filter="data")
ids = {i.strip() for s in SPLITS.values() for i in open(f"{root}/splits/{s}.txt") if i.strip()}
want = {f"{i}_S2Hand.tif" for i in ids} | {f"{i}_LabelHand.tif" for i in ids}
with tarfile.open(arc, "r|gz") as t: # pass 2: only the chips of the evaluated splits
for m in t:
if os.path.basename(m.name) in want and ("/S2L1CHand/" in m.name or "/LabelHand/" in m.name): t.extract(m, data, filter="data")
open(f"{root}/.ok", "w").close()
return root
def read(p):
with rasterio.open(p) as s: return s.read()
def main():
ap = argparse.ArgumentParser(); ap.add_argument("--model", required=True); ap.add_argument("--data", default="data")
a = ap.parse_args(); os.makedirs(a.data, exist_ok=True)
root = fetch(a.data); sess = session(a.model)
for split, f in SPLITS.items():
ids = [l.strip() for l in open(f"{root}/splits/{f}.txt") if l.strip()]
m = MIoU(["not water", "water"])
for i in ids:
dn = np.clip(np.nan_to_num(read(f"{root}/data/S2L1CHand/{i}_S2Hand.tif")).astype(np.int32), 0, 65535).astype(np.uint16)
y = read(f"{root}/data/LabelHand/{i}_LabelHand.tif")[0].astype(np.int64) # -1 = no label
m.add(y, predict(sess, normalize(dn)))
m.report(f"Sen1Floods11 {split} ({len(ids)} chips)")
if __name__ == "__main__": main()
|