rockglacierNET / preprocess /align_and_stack.py

Upload folder using huggingface_hub

820fab0 verified 4 months ago

4.47 kB

import os
import glob
import rasterio
from rasterio.warp import reproject, Resampling
import numpy as np
import geopandas as gpd
from rasterio.features import rasterize

INPUT_DIR = "mixed_res_dataset"
SHAPEFILE_PATH = "Final_merged01/Final_merged01.shp"
OUTPUT_FEATURE = "mixed_res_dataset/stacked_features.tif"
OUTPUT_MASK = "mixed_res_dataset/stacked_mask.tif"

def main():
    # Find all downloaded bands
    band_files = sorted(glob.glob(os.path.join(INPUT_DIR, "Band*.tif")))
    if not band_files:
        print(f"No band files found in {INPUT_DIR}")
        return

    print(f"Found {len(band_files)} band files.")

    # We will use Band01 as the reference geometry (10m resolution)
    ref_file = [f for f in band_files if "Band01" in f][0]
    
    with rasterio.open(ref_file) as src_ref:
        ref_meta = src_ref.meta.copy()
        ref_transform = src_ref.transform
        ref_crs = src_ref.crs
        ref_height = src_ref.height
        ref_width = src_ref.width

    print(f"Reference geometry based on {os.path.basename(ref_file)}:")
    print(f"  Shape: ({ref_height}, {ref_width})")
    print(f"  Transform: {ref_transform}")
    print(f"  CRS: {ref_crs}")

    # Prepare output metadata for features
    ref_meta.update({
        "count": len(band_files),
        "dtype": "float32",
        "nodata": 0.0,
        "compress": "lzw",
        "interleave": "band", # band sequential for easy reading
        "BIGTIFF": "YES"
    })

    print(f"\n1. Stacking and aligning features to {OUTPUT_FEATURE}")
    if not os.path.exists(OUTPUT_FEATURE):
        with rasterio.open(OUTPUT_FEATURE, "w", **ref_meta) as dst:
            for i, b_path in enumerate(band_files, start=1):
                b_name = os.path.basename(b_path)
                print(f"  Processing [{i}/{len(band_files)}] {b_name}...")
                
                with rasterio.open(b_path) as src_b:
                    # Read data
                    data = src_b.read(1)
                    
                    # If shapes match exactly, just cast and write
                    if src_b.shape == (ref_height, ref_width) and src_b.transform == ref_transform:
                        data = data.astype(np.float32)
                        dst.write(data, i)
                    else:
                        # Resample needed
                        print(f"    Resampling {b_name} from {src_b.shape} to ({ref_height}, {ref_width})")
                        resampled_data = np.empty((ref_height, ref_width), dtype=np.float32)
                        reproject(
                            source=data,
                            destination=resampled_data,
                            src_transform=src_b.transform,
                            src_crs=src_b.crs,
                            dst_transform=ref_transform,
                            dst_crs=ref_crs,
                            resampling=Resampling.bilinear
                        )
                        dst.write(resampled_data, i)
    else:
        print(f"  {OUTPUT_FEATURE} already exists. Skipping feature stacking.")

    print("\n2. Generating mask from shapefile")
    if os.path.exists(SHAPEFILE_PATH):
        try:
            gdf = gpd.read_file(SHAPEFILE_PATH)
            # Ensure the CRS matches the reference raster
            if gdf.crs != ref_crs:
                print(f"  Reprojecting shapefile from {gdf.crs} to {ref_crs}")
                gdf = gdf.to_crs(ref_crs)
            
            # Rasterize
            mask_meta = ref_meta.copy()
            mask_meta.update({
                "count": 1,
                "dtype": "uint8",
                "nodata": 0
            })
            
            print(f"  Rasterizing to {OUTPUT_MASK}")
            mask_array = rasterize(
                [(geom, 1) for geom in gdf.geometry if geom is not None and geom.is_valid],
                out_shape=(ref_height, ref_width),
                transform=ref_transform,
                fill=0,
                dtype="uint8"
            )
            
            with rasterio.open(OUTPUT_MASK, "w", **mask_meta) as dst_mask:
                dst_mask.write(mask_array, 1)
            print("  Mask generation completed.")
            
        except Exception as e:
            print(f"  Failed to rasterize mask: {e}")
    else:
        print(f"  Shapefile not found at {SHAPEFILE_PATH}. Skipping mask generation.")

    print("\nAll done!")

if __name__ == '__main__':
    main()

	import os
	import glob
	import rasterio
	from rasterio.warp import reproject, Resampling
	import numpy as np
	import geopandas as gpd
	from rasterio.features import rasterize

	INPUT_DIR = "mixed_res_dataset"
	SHAPEFILE_PATH = "Final_merged01/Final_merged01.shp"
	OUTPUT_FEATURE = "mixed_res_dataset/stacked_features.tif"
	OUTPUT_MASK = "mixed_res_dataset/stacked_mask.tif"

	def main():
	# Find all downloaded bands
	band_files = sorted(glob.glob(os.path.join(INPUT_DIR, "Band*.tif")))
	if not band_files:
	print(f"No band files found in {INPUT_DIR}")
	return

	print(f"Found {len(band_files)} band files.")

	# We will use Band01 as the reference geometry (10m resolution)
	ref_file = [f for f in band_files if "Band01" in f][0]

	with rasterio.open(ref_file) as src_ref:
	ref_meta = src_ref.meta.copy()
	ref_transform = src_ref.transform
	ref_crs = src_ref.crs
	ref_height = src_ref.height
	ref_width = src_ref.width

	print(f"Reference geometry based on {os.path.basename(ref_file)}:")
	print(f" Shape: ({ref_height}, {ref_width})")
	print(f" Transform: {ref_transform}")
	print(f" CRS: {ref_crs}")

	# Prepare output metadata for features
	ref_meta.update({
	"count": len(band_files),
	"dtype": "float32",
	"nodata": 0.0,
	"compress": "lzw",
	"interleave": "band", # band sequential for easy reading
	"BIGTIFF": "YES"
	})

	print(f"\n1. Stacking and aligning features to {OUTPUT_FEATURE}")
	if not os.path.exists(OUTPUT_FEATURE):
	with rasterio.open(OUTPUT_FEATURE, "w", **ref_meta) as dst:
	for i, b_path in enumerate(band_files, start=1):
	b_name = os.path.basename(b_path)
	print(f" Processing [{i}/{len(band_files)}] {b_name}...")

	with rasterio.open(b_path) as src_b:
	# Read data
	data = src_b.read(1)

	# If shapes match exactly, just cast and write
	if src_b.shape == (ref_height, ref_width) and src_b.transform == ref_transform:
	data = data.astype(np.float32)
	dst.write(data, i)
	else:
	# Resample needed
	print(f" Resampling {b_name} from {src_b.shape} to ({ref_height}, {ref_width})")
	resampled_data = np.empty((ref_height, ref_width), dtype=np.float32)
	reproject(
	source=data,
	destination=resampled_data,
	src_transform=src_b.transform,
	src_crs=src_b.crs,
	dst_transform=ref_transform,
	dst_crs=ref_crs,
	resampling=Resampling.bilinear
	)
	dst.write(resampled_data, i)
	else:
	print(f" {OUTPUT_FEATURE} already exists. Skipping feature stacking.")

	print("\n2. Generating mask from shapefile")
	if os.path.exists(SHAPEFILE_PATH):
	try:
	gdf = gpd.read_file(SHAPEFILE_PATH)
	# Ensure the CRS matches the reference raster
	if gdf.crs != ref_crs:
	print(f" Reprojecting shapefile from {gdf.crs} to {ref_crs}")
	gdf = gdf.to_crs(ref_crs)

	# Rasterize
	mask_meta = ref_meta.copy()
	mask_meta.update({
	"count": 1,
	"dtype": "uint8",
	"nodata": 0
	})

	print(f" Rasterizing to {OUTPUT_MASK}")
	mask_array = rasterize(
	[(geom, 1) for geom in gdf.geometry if geom is not None and geom.is_valid],
	out_shape=(ref_height, ref_width),
	transform=ref_transform,
	fill=0,
	dtype="uint8"
	)

	with rasterio.open(OUTPUT_MASK, "w", **mask_meta) as dst_mask:
	dst_mask.write(mask_array, 1)
	print(" Mask generation completed.")

	except Exception as e:
	print(f" Failed to rasterize mask: {e}")
	else:
	print(f" Shapefile not found at {SHAPEFILE_PATH}. Skipping mask generation.")

	print("\nAll done!")

	if __name__ == '__main__':
	main()