rockglacierNET / preprocess /slice_patches.py

Upload folder using huggingface_hub

820fab0 verified 4 months ago

3.07 kB

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()

	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()