| """ |
| RockGlacierNet — Region collection (whole-Sikkim grid). [lives in patch_dataset/] |
| |
| Grids the AOI (bounding box of all inventory polygons + AOI_MARGIN_M ~= the Sikkim |
| glacier belt) into RAW_TILE_PX tiles and downloads, per tile: |
| * a 12-band Sentinel-2 + terrain GeoTIFF, and |
| * a 1-band mask rasterized from ALL polygons onto that tile's exact grid. |
| |
| REGION_MODE = "sampled" : every tile touching a glacier + REGION_BG_KEEP background tiles. |
| REGION_MODE = "dense" : every tile in the AOI (full mosaic, multi-GB). |
| |
| Robustness: the median composite is kept light with a scene-level cloud filter |
| (S2_CLOUD_PCT); if a tile still trips GEE "User memory limit exceeded", it is |
| auto-split into a finer sub-grid, each part downloaded and mosaicked in memory. |
| |
| Outputs (nothing outside patch_dataset/): |
| patch_dataset/raw_region/features/tile_RRRR_CCCC.tif |
| patch_dataset/raw_region/masks/tile_RRRR_CCCC.tif |
| |
| Run (env with ee+geemap+geopandas+rasterio, e.g. `brats`): |
| python patch_dataset/collect_region.py # sampled (default) |
| python patch_dataset/collect_region.py --mode dense |
| python patch_dataset/collect_region.py --limit 3 # smoke test |
| All tunables live in patch_dataset/config.py. |
| """ |
|
|
| import os |
| import sys |
| import time |
| import argparse |
| import requests |
| import numpy as np |
|
|
| sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) |
| import config as C |
|
|
| import ee |
| import geopandas as gpd |
| import rasterio |
| from rasterio.io import MemoryFile |
| from rasterio.merge import merge as rio_merge |
| from rasterio.features import rasterize |
| from shapely.geometry import box |
|
|
|
|
| |
| |
| |
|
|
| def feature_image(roi): |
| coll = (ee.ImageCollection(C.S2_COLLECTION) |
| .filterBounds(roi) |
| .filterDate(C.S2_DATE_START, C.S2_DATE_END) |
| .filter(ee.Filter.lt("CLOUDY_PIXEL_PERCENTAGE", C.S2_CLOUD_PCT))) |
|
|
| def mask_clouds(im): |
| qa = im.select("QA60") |
| m = qa.bitwiseAnd(1 << 10).eq(0).And(qa.bitwiseAnd(1 << 11).eq(0)) |
| return im.updateMask(m).divide(C.S2_SR_DIVISOR) |
|
|
| s2 = coll.map(mask_clouds).median() |
| ndvi = s2.normalizedDifference(["B8", "B4"]).rename("NDVI") |
| ndwi = s2.normalizedDifference(["B3", "B8"]).rename("NDWI") |
| ndsi = s2.normalizedDifference(["B3", "B11"]).rename("NDSI") |
| s2 = s2.select(["B2", "B3", "B4", "B8", "B11", "B12"]).addBands([ndvi, ndwi, ndsi]) |
|
|
| dem = ee.ImageCollection(C.DEM_COLLECTION).select("DEM").mosaic() |
| |
| |
| |
| dem_m = dem.reproject(crs=C.TARGET_CRS, scale=30) |
| terr = ee.Terrain.products(dem_m) |
| img = (s2.addBands(dem_m.rename("Elev")) |
| .addBands(terr.select("slope").rename("Slope")) |
| .addBands(terr.select("aspect").rename("Aspect"))) |
| return img.clip(roi) |
|
|
|
|
| |
| |
| |
|
|
| def _utm_to_ee_roi(bbox_utm): |
| ll = gpd.GeoSeries([box(*bbox_utm)], crs=C.TARGET_CRS).to_crs("EPSG:4326").total_bounds |
| return ee.Geometry.Rectangle([ll[0], ll[1], ll[2], ll[3]]) |
|
|
|
|
| def _download_bytes(bbox_utm, tries=3): |
| """Download one bbox as GeoTIFF bytes. Raises MemoryError on GEE memory limit.""" |
| roi = _utm_to_ee_roi(bbox_utm) |
| url = feature_image(roi).getDownloadURL( |
| {"region": roi, "scale": C.SCALE_M, "crs": C.TARGET_CRS, "format": "GEO_TIFF"}) |
| last = None |
| for k in range(tries): |
| r = requests.get(url, timeout=300) |
| if r.status_code == 200 and r.content[:2] in (b"II", b"MM"): |
| return r.content |
| last = r.content[:200] |
| if r.status_code == 400 and b"memory limit" in r.content.lower(): |
| raise MemoryError(last) |
| time.sleep(2 * (k + 1)) |
| raise RuntimeError(f"download failed: {last}") |
|
|
|
|
| def fetch_tile(bbox_utm, max_split=4): |
| """Download bbox; on GEE memory error, recursively split into a finer grid and |
| mosaic the parts in memory. Returns (array[C,H,W], transform, crs). |
| |
| NB: MemoryFile must stay open while its dataset is read, so we copy the array |
| out before closing (otherwise GDAL reads a freed buffer).""" |
| split = 1 |
| while True: |
| minx, miny, maxx, maxy = bbox_utm |
| xs = np.linspace(minx, maxx, split + 1) |
| ys = np.linspace(miny, maxy, split + 1) |
| try: |
| chunks = [_download_bytes((xs[i], ys[j], xs[i+1], ys[j+1])) |
| for i in range(split) for j in range(split)] |
| except MemoryError: |
| if split >= max_split: |
| raise RuntimeError("memory limit even at max split") |
| split *= 2 |
| continue |
|
|
| mfs, dsets = [], [] |
| try: |
| for b in chunks: |
| mf = MemoryFile(b); mfs.append(mf); dsets.append(mf.open()) |
| if len(dsets) == 1: |
| ds = dsets[0] |
| return ds.read().copy(), ds.transform, ds.crs |
| arr, transform = rio_merge(dsets) |
| return arr.copy(), transform, dsets[0].crs |
| finally: |
| for d in dsets: d.close() |
| for mf in mfs: mf.close() |
|
|
|
|
| |
| |
| |
|
|
| def main(): |
| ap = argparse.ArgumentParser(description="Whole-Sikkim region grid collection.") |
| ap.add_argument("--shapefile", default=str(C.SHAPEFILE_PATH)) |
| ap.add_argument("--mode", choices=["sampled", "dense"], default=C.REGION_MODE) |
| ap.add_argument("--limit", type=int, default=None, help="download at most N tiles") |
| args = ap.parse_args() |
|
|
| try: |
| ee.Initialize(project=C.GEE_PROJECT) |
| except Exception: |
| ee.Authenticate(); ee.Initialize(project=C.GEE_PROJECT) |
|
|
| C.RAW_REGION_FEATURES_DIR.mkdir(parents=True, exist_ok=True) |
| C.RAW_REGION_MASKS_DIR.mkdir(parents=True, exist_ok=True) |
|
|
| gdf = gpd.read_file(args.shapefile).to_crs(C.TARGET_CRS) |
| gdf = gdf[gdf.geometry.notna() & ~gdf.geometry.is_empty].copy() |
| sidx = gdf.sindex |
|
|
| minx, miny, maxx, maxy = gdf.total_bounds |
| m, ts = C.AOI_MARGIN_M, C.RAW_TILE_PX * C.SCALE_M |
| minx, miny, maxx, maxy = minx - m, miny - m, maxx + m, maxy + m |
| xs = np.arange(minx, maxx, ts) |
| ys = np.arange(miny, maxy, ts) |
| rng = np.random.RandomState(C.RANDOM_SEED) |
|
|
| |
| plan = [] |
| for r, y in enumerate(ys): |
| for c, x in enumerate(xs): |
| bb = (x, y, x + ts, y + ts) |
| cand = list(sidx.intersection(bb)) |
| has = bool(cand) and gdf.geometry.iloc[cand].intersects(box(*bb)).any() |
| if args.mode == "dense" or has or rng.rand() < C.REGION_BG_KEEP: |
| plan.append((r, c, bb, has)) |
| if args.limit: |
| plan = plan[:args.limit] |
|
|
| n_g = sum(1 for *_, h in plan if h) |
| print(f"AOI {(maxx-minx)/1000:.0f}x{(maxy-miny)/1000:.0f} km | grid {len(xs)}x{len(ys)} " |
| f"= {len(xs)*len(ys)} cells | mode={args.mode} | downloading {len(plan)} tiles " |
| f"({n_g} glacier, {len(plan)-n_g} background)") |
|
|
| done = skip = fail = 0 |
| for n, (r, c, bb, has) in enumerate(plan): |
| fpath = C.RAW_REGION_FEATURES_DIR / f"tile_{r:04d}_{c:04d}.tif" |
| mpath = C.RAW_REGION_MASKS_DIR / f"tile_{r:04d}_{c:04d}.tif" |
| if fpath.exists() and mpath.exists(): |
| skip += 1; continue |
| try: |
| arr, transform, crs = fetch_tile(bb) |
| H, W = arr.shape[1], arr.shape[2] |
| with rasterio.open(fpath, "w", driver="GTiff", height=H, width=W, |
| count=arr.shape[0], dtype="float32", |
| crs=crs, transform=transform) as dst: |
| dst.write(arr.astype(np.float32)) |
| |
| cand = list(sidx.intersection(bb)) |
| geoms = [g for g in gdf.geometry.iloc[cand] if g.intersects(box(*bb))] |
| mask = (rasterize([(g, 1) for g in geoms], out_shape=(H, W), |
| transform=transform, fill=0, dtype="uint8") |
| if geoms else np.zeros((H, W), "uint8")) |
| with rasterio.open(mpath, "w", driver="GTiff", height=H, width=W, count=1, |
| dtype="uint8", crs=crs, transform=transform, nodata=0) as dst: |
| dst.write(mask, 1) |
| done += 1 |
| if done % 20 == 0 or n == len(plan) - 1: |
| print(f" {n+1}/{len(plan)} done={done} skip={skip} fail={fail} " |
| f"(last tile {H}x{W}, fg={int(mask.sum())})") |
| except Exception as e: |
| fail += 1 |
| print(f" tile_{r:04d}_{c:04d} FAILED: {e}") |
|
|
| print(f"\nRegion collection complete. done={done} skipped={skip} failed={fail}") |
| print(f"Features -> {C.RAW_REGION_FEATURES_DIR}") |
| print(f"Masks -> {C.RAW_REGION_MASKS_DIR}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
