| """ |
| RockGlacierNet — Stage 2: PATCH BUILDING (config-driven, any size). |
| |
| raw/{features,masks} -> patch_dataset/patches_{size}/{split}/{features,masks}/ |
| |
| What it does (replaces harmonize_dataset + split_dataset + extract_patches): |
| * pairs feature_XXXX.tif with mask_XXXX.tif, crops both to common shape |
| (safety net; the collector already aligns them), |
| * assigns each SOURCE GLACIER to train/val/test (split by image -> no |
| overlapping-patch leakage across splits), |
| * tiles every image into patch_size x patch_size GeoTIFF patches with overlap, |
| keeping all glacier patches + a sampled fraction of background patches, |
| * writes georeferenced GeoTIFF patches named feature_*.tif / mask_*.tif so the |
| existing RockGlacierDataGenerator can train on them UNCHANGED, |
| * writes metadata/ (config.json, patch_metadata.csv, dataset_summary.csv, |
| diagnostics_report.txt). |
| |
| Usage (rasterio only — runs anywhere, e.g. conda env `brats`): |
| python preprocess/build_patches.py # uses DEFAULT_PATCH (128) |
| python preprocess/build_patches.py --patch-size 64 |
| python preprocess/build_patches.py --patch-size 256 |
| # test against the OLD per-glacier tiles before re-collecting: |
| python preprocess/build_patches.py --patch-size 64 \ |
| --features-dir features --masks-dir masks --limit 40 |
| |
| All tunables live in patch_dataset/config.py. |
| """ |
|
|
| import os |
| import sys |
| import glob |
| import json |
| import math |
| import argparse |
| import numpy as np |
| import rasterio |
| import pandas as pd |
| from rasterio.windows import Window, transform as window_transform |
|
|
| sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) |
| from patch_dataset import config as C |
|
|
|
|
| |
| |
| |
|
|
| def padded_size(dim: int, patch: int, stride: int) -> int: |
| if dim <= patch: |
| return patch |
| n = math.ceil((dim - patch) / stride) |
| return n * stride + patch |
|
|
|
|
| def iter_patches(feat, mask, patch, stride, mode): |
| """Yield (feat_patch, mask_patch, row, col, padded) covering the whole image.""" |
| H, W = feat.shape[:2] |
| pH, pW = padded_size(H, patch, stride), padded_size(W, patch, stride) |
| dh, dw = pH - H, pW - W |
| padded = dh > 0 or dw > 0 |
| if padded: |
| feat = np.pad(feat, ((0, dh), (0, dw), (0, 0)), |
| mode=mode if mode == "reflect" else "constant") |
| mask = np.pad(mask, ((0, dh), (0, dw), (0, 0)), |
| mode="constant", constant_values=0) |
| for r in range(0, pH - patch + 1, stride): |
| for c in range(0, pW - patch + 1, stride): |
| yield feat[r:r+patch, c:c+patch, :], mask[r:r+patch, c:c+patch, :], r, c, padded |
|
|
|
|
| |
| |
| |
|
|
| def assign_splits(image_ids, ratios, seed): |
| rng = np.random.RandomState(seed) |
| ids = sorted(image_ids) |
| rng.shuffle(ids) |
| n = len(ids) |
| n_tr = int(n * ratios["train"]) |
| n_va = int(n * (ratios["train"] + ratios["val"])) |
| out = {} |
| for i, iid in enumerate(ids): |
| out[iid] = "train" if i < n_tr else ("val" if i < n_va else "test") |
| return out |
|
|
|
|
| |
| |
| |
|
|
| def run(patch, feat_dir, mask_dir, limit): |
| stride = C.stride_for(patch) |
| outdir = C.patches_dir(patch) |
| cfg = C.as_dict(patch) |
| cfg["features_dir_used"], cfg["masks_dir_used"] = str(feat_dir), str(mask_dir) |
|
|
| (outdir / "metadata").mkdir(parents=True, exist_ok=True) |
| for sp in C.SPLIT_RATIOS: |
| (outdir / sp / "features").mkdir(parents=True, exist_ok=True) |
| (outdir / sp / "masks").mkdir(parents=True, exist_ok=True) |
| with open(outdir / "metadata" / "config.json", "w") as f: |
| json.dump(cfg, f, indent=2) |
|
|
| feat_files = sorted(glob.glob(os.path.join(str(feat_dir), "feature_*.tif"))) |
| if limit: |
| feat_files = feat_files[:limit] |
|
|
| def img_id(fp): |
| return os.path.basename(fp).replace("feature_", "").replace(".tif", "") |
|
|
| ids = [img_id(fp) for fp in feat_files] |
| split_of = assign_splits(ids, C.SPLIT_RATIOS, C.RANDOM_SEED) |
|
|
| print(f"patch_size={patch} stride={stride} | {len(feat_files)} source tiles " |
| f"-> {outdir.name}/") |
|
|
| rng = np.random.RandomState(C.RANDOM_SEED) |
| patch_records, image_records = [], [] |
| n_gen = n_kept = 0 |
|
|
| for fp in feat_files: |
| iid = img_id(fp) |
| split = split_of[iid] |
| mp = os.path.join(str(mask_dir), f"mask_{iid}.tif") |
| if not os.path.exists(mp): |
| print(f" {iid}: no matching mask — skipped") |
| continue |
|
|
| with rasterio.open(fp) as fs: |
| feat = np.moveaxis(fs.read().astype(np.float32), 0, -1) |
| base_transform, crs = fs.transform, fs.crs |
| with rasterio.open(mp) as ms: |
| mask = ms.read(1).astype(np.uint8)[:, :, np.newaxis] |
|
|
| |
| ch, cw = min(feat.shape[0], mask.shape[0]), min(feat.shape[1], mask.shape[1]) |
| feat, mask = feat[:ch, :cw, :], mask[:ch, :cw, :] |
| feat = np.where((feat == C.NODATA_VALUE) | ~np.isfinite(feat), 0.0, feat) |
|
|
| oH, oW = feat.shape[:2] |
| img_fg = int((mask > 0).sum()) |
| n_img_kept = n_img_gen = 0 |
|
|
| for fpat, mpat, r, c, padded in iter_patches(feat, mask, patch, stride, C.PADDING_MODE): |
| g_pct = 100.0 * int((mpat > 0).sum()) / mpat.size |
| keep = (g_pct > C.MIN_GLACIER_PERCENTAGE) or (rng.rand() < C.BACKGROUND_KEEP_RATIO) |
| n_gen += 1 |
| n_img_gen += 1 |
|
|
| patch_id = f"{iid}_r{r:04d}_c{c:04d}" |
| rec = { |
| "patch_id": patch_id, "source_image_id": iid, "split": split, |
| "patch_row": r, "patch_col": c, |
| "glacier_pixel_count": int((mpat > 0).sum()), |
| "glacier_percentage": round(g_pct, 4), |
| "kept_or_filtered": "kept" if keep else "filtered", |
| "padding_applied": bool(padded), |
| } |
| for b, name in enumerate(C.BAND_NAMES): |
| rec[f"mean_{name}"] = round(float(fpat[:, :, b].mean()), 6) |
| patch_records.append(rec) |
|
|
| if not keep: |
| continue |
|
|
| ptransform = window_transform(Window(c, r, patch, patch), base_transform) |
| fout = outdir / split / "features" / f"feature_{patch_id}.tif" |
| mout = outdir / split / "masks" / f"mask_{patch_id}.tif" |
| with rasterio.open(fout, "w", driver="GTiff", height=patch, width=patch, |
| count=C.NUM_BANDS, dtype="float32", |
| crs=crs, transform=ptransform) as dst: |
| dst.write(np.moveaxis(fpat.astype(np.float32), -1, 0)) |
| with rasterio.open(mout, "w", driver="GTiff", height=patch, width=patch, |
| count=1, dtype="uint8", |
| crs=crs, transform=ptransform, nodata=0) as dst: |
| dst.write(mpat[:, :, 0].astype(np.uint8), 1) |
| n_kept += 1 |
| n_img_kept += 1 |
|
|
| image_records.append({ |
| "image_id": iid, "split": split, |
| "original_height": oH, "original_width": oW, |
| "glacier_pixels": img_fg, |
| "glacier_percentage": round(100.0 * img_fg / mask.size, 4), |
| "generated_patches": n_img_gen, "kept_patches": n_img_kept, |
| }) |
|
|
| patch_df = pd.DataFrame(patch_records) |
| image_df = pd.DataFrame(image_records) |
| patch_df.to_csv(outdir / "metadata" / "patch_metadata.csv", index=False) |
| image_df.to_csv(outdir / "metadata" / "dataset_summary.csv", index=False) |
| _diagnostics(patch_df, image_df, outdir, cfg) |
| print(f"\nDone. generated={n_gen} kept={n_kept} -> {outdir}") |
|
|
|
|
| def _diagnostics(patch_df, image_df, outdir, cfg): |
| kept = patch_df[patch_df["kept_or_filtered"] == "kept"] if len(patch_df) else patch_df |
| L = [] |
| def h(t): L.extend(["", "─" * 56, f" {t}", "─" * 56]) |
|
|
| L += ["=" * 56, " ROCKGLACIERENET — PATCH BUILD DIAGNOSTICS", "=" * 56] |
| L += [f" patch_size {cfg['patch_size']} | stride {cfg['stride']} | " |
| f"bg_keep {cfg['background_keep_ratio']} | pad {cfg['padding_mode']}"] |
|
|
| h("PATCH COUNTS") |
| for sp in C.SPLIT_RATIOS: |
| g = len(patch_df[patch_df["split"] == sp]) if len(patch_df) else 0 |
| k = len(kept[kept["split"] == sp]) if len(kept) else 0 |
| L.append(f" {sp:5s}: generated={g:6d} kept={k:6d} filtered={g-k:6d}") |
| L.append(f" TOTAL: generated={len(patch_df):6d} kept={len(kept):6d}") |
|
|
| h("IMAGE COUNTS PER SPLIT") |
| for sp in C.SPLIT_RATIOS: |
| L.append(f" {sp:5s}: {len(image_df[image_df['split']==sp]) if len(image_df) else 0} glaciers") |
|
|
| if len(kept): |
| h("PIXEL CLASS BALANCE (kept patches)") |
| tfg = int(kept["glacier_pixel_count"].sum()) |
| tpx = len(kept) * cfg["patch_size"] ** 2 |
| tbg = tpx - tfg |
| L.append(f" Glacier : {tfg:>12,} ({100*tfg/tpx:.2f}%)") |
| L.append(f" Background: {tbg:>12,} ({100*tbg/tpx:.2f}%)") |
| if tfg: |
| L.append(f" Imbalance : {tbg/tfg:.1f} : 1") |
|
|
| h("GLACIER COVERAGE OF KEPT PATCHES (%)") |
| for lo, hi in [(0, 0.1), (0.1, 5), (5, 20), (20, 50), (50, 80), (80, 100.01)]: |
| n = int(((kept["glacier_percentage"] >= lo) & (kept["glacier_percentage"] < hi)).sum()) |
| L.append(f" {lo:5.1f}-{hi:5.1f}%: {n:6d} {'█' * min(n // 5, 40)}") |
|
|
| if len(image_df): |
| h("DATASET EXPANSION (glaciers -> kept patches)") |
| for sp in C.SPLIT_RATIOS: |
| sub = image_df[image_df["split"] == sp] |
| ni, nk = len(sub), int(sub["kept_patches"].sum()) if len(sub) else 0 |
| L.append(f" {sp:5s}: {ni:4d} glaciers -> {nk:6d} patches (×{nk/ni:.1f})" |
| if ni else f" {sp:5s}: 0 glaciers") |
|
|
| report = "\n".join(L) |
| print("\n" + report) |
| with open(outdir / "metadata" / "diagnostics_report.txt", "w") as f: |
| f.write(report) |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser(description="Build GeoTIFF patches from raw tiles.") |
| ap.add_argument("--patch-size", type=int, default=C.DEFAULT_PATCH, choices=None) |
| ap.add_argument("--features-dir", default=str(C.RAW_FEATURES_DIR)) |
| ap.add_argument("--masks-dir", default=str(C.RAW_MASKS_DIR)) |
| ap.add_argument("--limit", type=int, default=None) |
| args = ap.parse_args() |
| run(args.patch_size, args.features_dir, args.masks_dir, args.limit) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
