File size: 2,917 Bytes
80e070e | 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 | """
Compute per-band normalization statistics (mean/std) over the TRAIN split of a
dataset and write them to <dataset>/metadata/norm_stats.json.
Stats are accumulated over VALID pixels only (NoData -9999 and non-finite values
are excluded). The RockGlacierDataGenerator auto-loads this file via its
_load_norm_stats() (it looks at <features_dir>/../../metadata/norm_stats.json),
so once this file exists the next training run uses the correct stats.
Usage:
python compute_norm_stats.py # defaults to output/dataset
python compute_norm_stats.py <train_features_dir> # custom
"""
import os
import sys
import glob
import json
import numpy as np
import rasterio
BAND_NAMES = ["B2", "B3", "B4", "B8", "B11", "B12",
"NDVI", "NDWI", "NDSI", "Elev", "Slope", "Aspect"]
DEFAULT_FEAT_DIR = "datacollection/output/dataset/train/features"
def main(feat_dir=DEFAULT_FEAT_DIR):
files = sorted(glob.glob(os.path.join(feat_dir, "*.tif")))
if not files:
files = sorted(glob.glob(os.path.join(feat_dir, "*.npy")))
assert files, f"No .tif/.npy features in {feat_dir}"
nb = None
ssum = ssq = cnt = None
for fp in files:
if fp.endswith(".npy"):
arr = np.load(fp).astype(np.float64) # (H, W, C)
else:
with rasterio.open(fp) as s:
arr = np.moveaxis(s.read(), 0, -1).astype(np.float64) # (H, W, C)
if nb is None:
nb = arr.shape[-1]
ssum = np.zeros(nb); ssq = np.zeros(nb); cnt = np.zeros(nb)
flat = arr.reshape(-1, nb)
valid = np.isfinite(flat) & (flat != -9999)
x = np.where(valid, flat, 0.0)
ssum += x.sum(axis=0)
ssq += (x * x).sum(axis=0)
cnt += valid.sum(axis=0)
mean = ssum / np.maximum(cnt, 1)
var = ssq / np.maximum(cnt, 1) - mean ** 2
std = np.sqrt(np.maximum(var, 1e-12))
names = BAND_NAMES[:nb] if nb <= len(BAND_NAMES) else [f"band_{i}" for i in range(nb)]
out = {
"band_names": names,
"mean": [float(m) for m in mean],
"std": [float(s) for s in std],
"num_train_images": len(files),
"num_pixels": int(cnt.max()),
}
# <feat_dir>/../../metadata/norm_stats.json (e.g. output/dataset/metadata/)
meta_dir = os.path.normpath(os.path.join(feat_dir, os.pardir, os.pardir, "metadata"))
os.makedirs(meta_dir, exist_ok=True)
out_path = os.path.join(meta_dir, "norm_stats.json")
with open(out_path, "w") as f:
json.dump(out, f, indent=2)
print(f"Computed norm stats from {len(files)} images, {int(cnt.max()):,} valid pixels/band")
print(f" {'band':6s} {'mean':>12s} {'std':>12s}")
for n, m, s in zip(names, mean, std):
print(f" {n:6s} {m:12.4f} {s:12.4f}")
print(f"\nWrote {out_path}")
if __name__ == "__main__":
main(sys.argv[1] if len(sys.argv) > 1 else DEFAULT_FEAT_DIR)
|