| """ |
| RockGlacierNet — Stage 1: feature/mask COLLECTION (Option A: enlarged window). |
| |
| For every rock-glacier polygon in the inventory shapefile: |
| 1. compute its bbox in UTM and pad by CONTEXT_BUFFER_M on every side, |
| 2. download a 12-band Sentinel-2 + terrain feature stack over that window, |
| 3. rasterize the polygon to the EXACT grid of the downloaded feature tile, |
| so feature and mask are co-registered by construction (no harmonize step). |
| |
| Outputs (and nothing outside patch_dataset/): |
| patch_dataset/raw/features/feature_XXXX.tif (12-band, float32, 10 m, UTM45N) |
| patch_dataset/raw/masks/mask_XXXX.tif (1-band uint8, same grid) |
| |
| ENVIRONMENT: needs `earthengine-api`, `geemap`, `geopandas`, `rasterio`. |
| Run on the machine that has them (your Windows box): |
| python preprocess/collect_features.py |
| python preprocess/collect_features.py --limit 5 # smoke test |
| python preprocess/collect_features.py --shapefile /path/to/your.shp |
| |
| All tunables live in patch_dataset/config.py. |
| """ |
|
|
| import os |
| import sys |
| import argparse |
|
|
| sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) |
| from patch_dataset import config as C |
|
|
| import numpy as np |
| import geopandas as gpd |
| import rasterio |
| from rasterio.features import rasterize |
| import ee |
| import geemap |
| from shapely.geometry import box |
|
|
|
|
| |
| |
| |
|
|
| def get_s2_image(roi): |
| coll = (ee.ImageCollection(C.S2_COLLECTION) |
| .filterBounds(roi) |
| .filterDate(C.S2_DATE_START, C.S2_DATE_END)) |
|
|
| def mask_clouds(image): |
| qa = image.select("QA60") |
| cloud = 1 << 10 |
| cirrus = 1 << 11 |
| m = qa.bitwiseAnd(cloud).eq(0).And(qa.bitwiseAnd(cirrus).eq(0)) |
| return image.updateMask(m).divide(C.S2_SR_DIVISOR) |
|
|
| s2 = coll.map(mask_clouds).median().clip(roi) |
| ndvi = s2.normalizedDifference(["B8", "B4"]).rename("NDVI") |
| ndwi = s2.normalizedDifference(["B3", "B8"]).rename("NDWI") |
| ndsi = s2.normalizedDifference(["B3", "B11"]).rename("NDSI") |
| |
| return s2.select(["B2", "B3", "B4", "B8", "B11", "B12"]).addBands([ndvi, ndwi, ndsi]) |
|
|
|
|
| def get_terrain_image(roi): |
| |
| def add_terrain(img): |
| slope = ee.Terrain.slope(img).rename("slope") |
| aspect = ee.Terrain.aspect(img).rename("aspect") |
| return img.addBands([slope, aspect]) |
|
|
| terrain = (ee.ImageCollection(C.DEM_COLLECTION).select("DEM") |
| .map(add_terrain).mosaic().clip(roi)) |
| elevation = terrain.select("DEM").rename("elevation") |
| slope = terrain.select("slope") |
| aspect = terrain.select("aspect") |
| return elevation.addBands([slope, aspect]) |
|
|
|
|
| def build_feature_image(roi): |
| """12-band stack in the contractual order C.BAND_NAMES.""" |
| return get_s2_image(roi).addBands(get_terrain_image(roi)) |
|
|
|
|
| |
| |
| |
|
|
| def padded_roi(geom_utm): |
| """Return (ee_roi_4326, bounds_utm) for a single UTM geometry + buffer.""" |
| minx, miny, maxx, maxy = geom_utm.bounds |
| b = C.CONTEXT_BUFFER_M |
| minx, miny, maxx, maxy = minx - b, miny - b, maxx + b, maxy + b |
| box_utm = gpd.GeoDataFrame(geometry=[box(minx, miny, maxx, maxy)], crs=C.TARGET_CRS) |
| box_4326 = box_utm.to_crs("EPSG:4326").total_bounds |
| roi = ee.Geometry.BBox(box_4326[0], box_4326[1], box_4326[2], box_4326[3]) |
| return roi, (minx, miny, maxx, maxy) |
|
|
|
|
| def rasterize_aligned_mask(geom_utm, feat_tif_path, out_mask_path): |
| """Rasterize geom onto the EXACT grid of the downloaded feature tile.""" |
| with rasterio.open(feat_tif_path) as src: |
| transform, height, width, crs = src.transform, src.height, src.width, src.crs |
| mask = rasterize( |
| [(geom_utm, 1)], |
| out_shape=(height, width), |
| transform=transform, |
| fill=0, |
| dtype="uint8", |
| ) |
| with rasterio.open( |
| out_mask_path, "w", driver="GTiff", |
| height=height, width=width, count=1, dtype="uint8", |
| crs=crs, transform=transform, nodata=0, |
| ) as dst: |
| dst.write(mask, 1) |
| return int(mask.sum()) |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser(description="Collect enlarged-window feature/mask tiles.") |
| ap.add_argument("--shapefile", default=str(C.SHAPEFILE_PATH)) |
| ap.add_argument("--limit", type=int, default=None, help="process only first N polygons") |
| args = ap.parse_args() |
|
|
| |
| try: |
| ee.Initialize(project=C.GEE_PROJECT) |
| except Exception: |
| ee.Authenticate() |
| ee.Initialize(project=C.GEE_PROJECT) |
|
|
| C.RAW_FEATURES_DIR.mkdir(parents=True, exist_ok=True) |
| C.RAW_MASKS_DIR.mkdir(parents=True, exist_ok=True) |
|
|
| gdf = gpd.read_file(args.shapefile).to_crs(C.TARGET_CRS) |
| n = len(gdf) if args.limit is None else min(args.limit, len(gdf)) |
| print(f"Loaded {len(gdf)} polygons from {args.shapefile} | processing {n} " |
| f"| buffer={C.CONTEXT_BUFFER_M:.0f} m | scale={C.SCALE_M:.0f} m") |
|
|
| n_done = n_skip = n_fail = 0 |
| for idx in range(n): |
| geom = gdf.geometry.iloc[idx] |
| if geom is None or geom.is_empty: |
| print(f" [{idx:04d}] empty geometry — skipped") |
| n_skip += 1 |
| continue |
|
|
| feat_path = C.RAW_FEATURES_DIR / f"feature_{idx:04d}.tif" |
| mask_path = C.RAW_MASKS_DIR / f"mask_{idx:04d}.tif" |
| if feat_path.exists() and mask_path.exists(): |
| n_skip += 1 |
| continue |
|
|
| try: |
| roi, _ = padded_roi(geom) |
| img = build_feature_image(roi) |
| geemap.ee_export_image( |
| img, filename=str(feat_path), scale=C.SCALE_M, |
| region=roi, crs=C.TARGET_CRS, file_per_band=False, |
| ) |
| fg = rasterize_aligned_mask(geom, str(feat_path), str(mask_path)) |
| n_done += 1 |
| if n_done % 25 == 0 or idx == n - 1: |
| print(f" [{idx:04d}] done={n_done} skip={n_skip} fail={n_fail} " |
| f"(last fg px={fg})") |
| except Exception as e: |
| n_fail += 1 |
| print(f" [{idx:04d}] FAILED: {e}") |
|
|
| print(f"\nCollection complete. done={n_done} skipped={n_skip} failed={n_fail}") |
| print(f"Features -> {C.RAW_FEATURES_DIR}") |
| print(f"Masks -> {C.RAW_MASKS_DIR}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
