File size: 4,465 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 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 | 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()
|