rockglacierNET / preprocess /generate_masks.py

Upload folder using huggingface_hub

80e070e verified 5 months ago

4.39 kB

import os
import glob
import geopandas as gpd
import pandas as pd
import rasterio
from rasterio.features import rasterize
from rasterio.transform import from_origin
import matplotlib.pyplot as plt
import matplotlib
matplotlib.use('Agg')

INPUT_DIR = r'D:\RockGlacier_Project\EDA\individual_polygons'
OUTPUT_MASK_DIR = r'D:\RockGlacier_Project\masks'
REPORT_PATH = os.path.join(OUTPUT_MASK_DIR, 'mask_report.csv')

TARGET_CRS = 'EPSG:32645'  # UTM Zone 45N for Sikkim
RESOLUTION = 10.0  # 10 meters per pixel
PADDING = 200.0  # 200 meters padding around the polygon

def generate_masks():
    os.makedirs(OUTPUT_MASK_DIR, exist_ok=True)
    
    geojson_files = sorted(glob.glob(os.path.join(INPUT_DIR, '*.geojson')))
    if not geojson_files:
        print(f"No GeoJSON files found in {INPUT_DIR}")
        return

    print(f"Found {len(geojson_files)} polygons. Starting rasterization at {RESOLUTION}m/px...")
    
    report_data = []
    
    for filepath in geojson_files:
        filename = os.path.basename(filepath)
        # Extracts index from 'polygon_0.geojson'
        idx_str = filename.split('_')[1].split('.')[0]
        idx = int(idx_str)
        
        # 1. Read polygon
        gdf = gpd.read_file(filepath)
        
        # 2. Reproject to UTM Zone 45N for meter-based math
        gdf_proj = gdf.to_crs(TARGET_CRS)
        
        # 3. Compute Bounds with Padding
        bounds = gdf_proj.total_bounds
        import numpy as np
        if np.isnan(bounds).any() or gdf_proj.empty or gdf_proj.geometry.is_empty.all() or gdf_proj.geometry.isna().all():
            print(f"Skipping invalid/empty geometry for polygon {idx}")
            continue
            
        minx, miny, maxx, maxy = bounds
        
        minx -= PADDING
        miny -= PADDING
        maxx += PADDING
        maxy += PADDING
        
        # 4. Calculate Raster Dimensions
        width = int((maxx - minx) / RESOLUTION)
        height = int((maxy - miny) / RESOLUTION)
        
        # Adjust maxx and miny slightly so resolution is exactly 10.0
        maxx = minx + (width * RESOLUTION)
        miny = maxy - (height * RESOLUTION)
        
        # 5. Create Affine Transform
        # from_origin(west, north, xsize, ysize)
        transform = from_origin(minx, maxy, RESOLUTION, RESOLUTION)
        
        # 6. Rasterize
        # shapes requires an iterable of (geometry, value) pairs
        shapes = ((geom, 1) for geom in gdf_proj.geometry)
        mask = rasterize(
            shapes,
            out_shape=(height, width),
            transform=transform,
            fill=0,
            dtype='uint8'
        )
        
        # 7. Compute Stats
        fg_pixels = int(mask.sum())
        bg_pixels = int((width * height) - fg_pixels)
        rasterized_area = fg_pixels * (RESOLUTION ** 2)
        polygon_area = gdf_proj.geometry.area.sum()
        
        # 8. Save GeoTIFF
        out_tif = os.path.join(OUTPUT_MASK_DIR, f"mask_{idx:04d}.tif")
        with rasterio.open(
            out_tif,
            'w',
            driver='GTiff',
            height=height,
            width=width,
            count=1,
            dtype=mask.dtype,
            crs=TARGET_CRS,
            transform=transform,
        ) as dst:
            dst.write(mask, 1)
            
        # 9. Plot first 5 masks
        if idx < 5:
            plt.figure(figsize=(6, 6))
            plt.imshow(mask, cmap='gray')
            plt.title(f"Mask {idx:04d} ({width}x{height} pixels)")
            plt.axis('off')
            plt.savefig(os.path.join(OUTPUT_MASK_DIR, f"plot_mask_{idx:04d}.png"), dpi=150, bbox_inches='tight')
            plt.close()
            
        report_data.append({
            'Polygon_ID': idx,
            'Width': width,
            'Height': height,
            'Foreground_Pixels': fg_pixels,
            'Background_Pixels': bg_pixels,
            'Polygon_Area_m2': round(polygon_area, 2),
            'Rasterized_Area_m2': round(rasterized_area, 2)
        })
        
        if idx % 50 == 0:
            print(f"Processed {idx+1}/{len(geojson_files)} masks...")

    # Save Report
    df_report = pd.DataFrame(report_data)
    df_report.to_csv(REPORT_PATH, index=False)
    print(f"\nSuccess! Generated {len(geojson_files)} masks.")
    print(f"Validation report saved to {REPORT_PATH}")

if __name__ == '__main__':
    generate_masks()

	import os
	import glob
	import geopandas as gpd
	import pandas as pd
	import rasterio
	from rasterio.features import rasterize
	from rasterio.transform import from_origin
	import matplotlib.pyplot as plt
	import matplotlib
	matplotlib.use('Agg')

	INPUT_DIR = r'D:\RockGlacier_Project\EDA\individual_polygons'
	OUTPUT_MASK_DIR = r'D:\RockGlacier_Project\masks'
	REPORT_PATH = os.path.join(OUTPUT_MASK_DIR, 'mask_report.csv')

	TARGET_CRS = 'EPSG:32645' # UTM Zone 45N for Sikkim
	RESOLUTION = 10.0 # 10 meters per pixel
	PADDING = 200.0 # 200 meters padding around the polygon

	def generate_masks():
	os.makedirs(OUTPUT_MASK_DIR, exist_ok=True)

	geojson_files = sorted(glob.glob(os.path.join(INPUT_DIR, '*.geojson')))
	if not geojson_files:
	print(f"No GeoJSON files found in {INPUT_DIR}")
	return

	print(f"Found {len(geojson_files)} polygons. Starting rasterization at {RESOLUTION}m/px...")

	report_data = []

	for filepath in geojson_files:
	filename = os.path.basename(filepath)
	# Extracts index from 'polygon_0.geojson'
	idx_str = filename.split('_')[1].split('.')[0]
	idx = int(idx_str)

	# 1. Read polygon
	gdf = gpd.read_file(filepath)

	# 2. Reproject to UTM Zone 45N for meter-based math
	gdf_proj = gdf.to_crs(TARGET_CRS)

	# 3. Compute Bounds with Padding
	bounds = gdf_proj.total_bounds
	import numpy as np
	if np.isnan(bounds).any() or gdf_proj.empty or gdf_proj.geometry.is_empty.all() or gdf_proj.geometry.isna().all():
	print(f"Skipping invalid/empty geometry for polygon {idx}")
	continue

	minx, miny, maxx, maxy = bounds

	minx -= PADDING
	miny -= PADDING
	maxx += PADDING
	maxy += PADDING

	# 4. Calculate Raster Dimensions
	width = int((maxx - minx) / RESOLUTION)
	height = int((maxy - miny) / RESOLUTION)

	# Adjust maxx and miny slightly so resolution is exactly 10.0
	maxx = minx + (width * RESOLUTION)
	miny = maxy - (height * RESOLUTION)

	# 5. Create Affine Transform
	# from_origin(west, north, xsize, ysize)
	transform = from_origin(minx, maxy, RESOLUTION, RESOLUTION)

	# 6. Rasterize
	# shapes requires an iterable of (geometry, value) pairs
	shapes = ((geom, 1) for geom in gdf_proj.geometry)
	mask = rasterize(
	shapes,
	out_shape=(height, width),
	transform=transform,
	fill=0,
	dtype='uint8'
	)

	# 7. Compute Stats
	fg_pixels = int(mask.sum())
	bg_pixels = int((width * height) - fg_pixels)
	rasterized_area = fg_pixels * (RESOLUTION ** 2)
	polygon_area = gdf_proj.geometry.area.sum()

	# 8. Save GeoTIFF
	out_tif = os.path.join(OUTPUT_MASK_DIR, f"mask_{idx:04d}.tif")
	with rasterio.open(
	out_tif,
	'w',
	driver='GTiff',
	height=height,
	width=width,
	count=1,
	dtype=mask.dtype,
	crs=TARGET_CRS,
	transform=transform,
	) as dst:
	dst.write(mask, 1)

	# 9. Plot first 5 masks
	if idx < 5:
	plt.figure(figsize=(6, 6))
	plt.imshow(mask, cmap='gray')
	plt.title(f"Mask {idx:04d} ({width}x{height} pixels)")
	plt.axis('off')
	plt.savefig(os.path.join(OUTPUT_MASK_DIR, f"plot_mask_{idx:04d}.png"), dpi=150, bbox_inches='tight')
	plt.close()

	report_data.append({
	'Polygon_ID': idx,
	'Width': width,
	'Height': height,
	'Foreground_Pixels': fg_pixels,
	'Background_Pixels': bg_pixels,
	'Polygon_Area_m2': round(polygon_area, 2),
	'Rasterized_Area_m2': round(rasterized_area, 2)
	})

	if idx % 50 == 0:
	print(f"Processed {idx+1}/{len(geojson_files)} masks...")

	# Save Report
	df_report = pd.DataFrame(report_data)
	df_report.to_csv(REPORT_PATH, index=False)
	print(f"\nSuccess! Generated {len(geojson_files)} masks.")
	print(f"Validation report saved to {REPORT_PATH}")

	if __name__ == '__main__':
	generate_masks()