| """ |
| 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) |
| else: |
| with rasterio.open(fp) as s: |
| arr = np.moveaxis(s.read(), 0, -1).astype(np.float64) |
| 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()), |
| } |
|
|
| |
| 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) |
|
|
