File size: 3,073 Bytes
820fab0 | 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 | import os
import rasterio
from rasterio.windows import Window
import numpy as np
import random
import glob
FEATURE_STACK = "mixed_res_dataset/stacked_features.tif"
MASK_STACK = "mixed_res_dataset/stacked_mask.tif"
OUTPUT_DIR = "dataset/extracted_patches"
PATCH_SIZE = 128
STRIDE = 64 # 50% overlap
def main():
if not os.path.exists(FEATURE_STACK) or not os.path.exists(MASK_STACK):
print("Stacked feature or mask not found. Please run align_and_stack.py first.")
return
os.makedirs(os.path.join(OUTPUT_DIR, "features"), exist_ok=True)
os.makedirs(os.path.join(OUTPUT_DIR, "masks"), exist_ok=True)
with rasterio.open(FEATURE_STACK) as src_feat, rasterio.open(MASK_STACK) as src_mask:
height = src_feat.height
width = src_feat.width
num_bands = src_feat.count
print(f"Opened stacked files. Size: {width}x{height}, Bands: {num_bands}")
patch_idx = 0
valid_patches = 0
# Read the whole mask into memory to fast-check if a patch has any glacier
# It's uint8, so for 11477x8969 it's ~100MB, easily fits in RAM
print("Loading full mask into memory for filtering...")
full_mask = src_mask.read(1)
print("Extracting patches...")
for y in range(0, height - PATCH_SIZE + 1, STRIDE):
for x in range(0, width - PATCH_SIZE + 1, STRIDE):
# Check mask first
patch_mask = full_mask[y:y+PATCH_SIZE, x:x+PATCH_SIZE]
# Filter: Keep if it contains any glacier OR 10% chance of background
is_glacier = np.any(patch_mask > 0)
if not is_glacier and random.random() > 0.1:
continue
# Read feature patch
window = Window(x, y, PATCH_SIZE, PATCH_SIZE)
patch_feat = src_feat.read(window=window)
# patch_feat shape is (Bands, H, W) -> we need (H, W, Bands) for generator
patch_feat = np.moveaxis(patch_feat, 0, -1)
patch_mask_expanded = np.expand_dims(patch_mask, axis=-1)
# Save as .npy
feat_out = os.path.join(OUTPUT_DIR, "features", f"patch_{patch_idx:06d}.npy")
mask_out = os.path.join(OUTPUT_DIR, "masks", f"patch_{patch_idx:06d}.npy")
np.save(feat_out, patch_feat)
np.save(mask_out, patch_mask_expanded)
patch_idx += 1
if is_glacier:
valid_patches += 1
if patch_idx % 100 == 0:
print(f" Extracted {patch_idx} patches (Glacier: {valid_patches})")
print(f"\nDone! Extracted total {patch_idx} patches.")
print(f"Patches containing glacier: {valid_patches}")
print(f"Background patches (kept): {patch_idx - valid_patches}")
print(f"Saved to {OUTPUT_DIR}")
if __name__ == '__main__':
main()
|