Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import os | |
| import tempfile | |
| import zipfile | |
| import uuid | |
| import rasterio | |
| import numpy as np | |
| from rasterio.mask import mask | |
| import geopandas as gpd | |
| import matplotlib.pyplot as plt | |
| from io import BytesIO | |
| from PIL import Image | |
| import shutil | |
| def clip_geotiff(tif_file, shapefile_zip): | |
| """ | |
| Clips a GeoTIFF file using a shapefile | |
| """ | |
| temp_dir = None | |
| try: | |
| # Check if files were provided | |
| if tif_file is None or shapefile_zip is None: | |
| return None, None | |
| # Create unique temporary directory | |
| temp_dir = tempfile.mkdtemp(prefix="geotiff_clip_") | |
| img_dir = os.path.join(temp_dir, "image") | |
| shp_dir = os.path.join(temp_dir, "shapefile") | |
| out_dir = os.path.join(temp_dir, "output") | |
| # Create subdirectories | |
| for directory in [img_dir, shp_dir, out_dir]: | |
| os.makedirs(directory, exist_ok=True) | |
| # Handle TIFF file - tif_file is now a file path (string) | |
| tif_path = os.path.join(img_dir, f"input_{uuid.uuid4().hex}.tif") | |
| shutil.copy2(tif_file, tif_path) | |
| # Handle shapefile ZIP - shapefile_zip is now a file path (string) | |
| zip_path = os.path.join(shp_dir, f"shapefile_{uuid.uuid4().hex}.zip") | |
| shutil.copy2(shapefile_zip, zip_path) | |
| # Extract ZIP | |
| with zipfile.ZipFile(zip_path, 'r') as zip_ref: | |
| zip_ref.extractall(shp_dir) | |
| # Find .shp file | |
| shp_files = [f for f in os.listdir(shp_dir) if f.endswith(".shp")] | |
| if not shp_files: | |
| raise FileNotFoundError(".shp file not found in the provided ZIP.") | |
| shp_file = os.path.join(shp_dir, shp_files[0]) | |
| # Read shapefile using geopandas | |
| gdf = gpd.read_file(shp_file) | |
| # Check if shapefile has valid geometries | |
| if gdf.empty or gdf.geometry.isna().all(): | |
| raise ValueError("Shapefile does not contain valid geometries.") | |
| # Open and process GeoTIFF file | |
| with rasterio.open(tif_path) as src: | |
| # Check for overlap between shapefile and raster | |
| gdf_proj = gdf.to_crs(src.crs) | |
| # Perform clipping | |
| out_image, out_transform = mask(src, gdf_proj.geometry, crop=True, nodata=src.nodata) | |
| # Update metadata | |
| out_meta = src.meta.copy() | |
| out_meta.update({ | |
| "height": out_image.shape[1], | |
| "width": out_image.shape[2], | |
| "transform": out_transform, | |
| "nodata": src.nodata | |
| }) | |
| # Create output file in a persistent temporary location | |
| output_filename = f"clipped_{uuid.uuid4().hex}.tif" | |
| # Use tempfile.NamedTemporaryFile to create a file that Gradio can manage | |
| output_temp_file = tempfile.NamedTemporaryFile( | |
| suffix=".tif", | |
| prefix="clipped_", | |
| delete=False # Don't auto-delete, let Gradio handle it | |
| ) | |
| output_tif_path = output_temp_file.name | |
| output_temp_file.close() # Close the file handle so rasterio can write to it | |
| with rasterio.open(output_tif_path, "w", **out_meta) as dest: | |
| dest.write(out_image) | |
| # Create PNG visualization in memory | |
| # Prepare data for visualization | |
| if out_image.shape[0] >= 3: | |
| # If has 3 or more bands, use first 3 (RGB) | |
| preview_array = out_image[:3] | |
| else: | |
| # If has less than 3 bands, repeat first band | |
| preview_array = np.repeat(out_image[0:1], 3, axis=0) | |
| # Rearrange dimensions (bands, height, width) -> (height, width, bands) | |
| preview_array = np.moveaxis(preview_array, 0, -1) | |
| # Normalize values to 0-255 if necessary | |
| if preview_array.dtype != np.uint8: | |
| # Normalize to 0-255 | |
| preview_min = np.nanmin(preview_array) | |
| preview_max = np.nanmax(preview_array) | |
| if preview_max > preview_min: | |
| preview_array = ((preview_array - preview_min) / (preview_max - preview_min) * 255).astype(np.uint8) | |
| else: | |
| preview_array = np.zeros_like(preview_array, dtype=np.uint8) | |
| # Handle nodata values | |
| if out_meta.get('nodata') is not None: | |
| nodata_mask = np.any(out_image == out_meta['nodata'], axis=0) | |
| preview_array[nodata_mask] = [0, 0, 0] # Black for nodata | |
| # Create matplotlib figure | |
| plt.style.use('default') # Ensure default style | |
| fig, ax = plt.subplots(figsize=(10, 8), dpi=100) | |
| ax.imshow(preview_array) | |
| ax.set_title(f"Clipped GeoTIFF - {output_filename}", fontsize=12, pad=20) | |
| ax.axis('off') | |
| # Save to memory buffer | |
| buf = BytesIO() | |
| plt.savefig(buf, format='png', bbox_inches='tight', pad_inches=0.1, dpi=150) | |
| plt.close(fig) # Important: close figure to free memory | |
| buf.seek(0) | |
| # Convert to PIL image | |
| pil_image = Image.open(buf).convert("RGB") | |
| buf.close() | |
| # Return the PIL image and the path to the persistent temporary file | |
| return pil_image, output_tif_path | |
| except Exception as e: | |
| error_msg = f"Error during processing: {str(e)}" | |
| print(f"[ERROR] {error_msg}") | |
| # Create error image | |
| error_image = Image.new('RGB', (400, 200), color='white') | |
| from PIL import ImageDraw, ImageFont | |
| draw = ImageDraw.Draw(error_image) | |
| try: | |
| font = ImageFont.truetype("arial.ttf", 16) | |
| except: | |
| font = ImageFont.load_default() | |
| draw.text((10, 10), "Processing Error:", fill='red', font=font) | |
| draw.text((10, 40), str(e)[:50] + "..." if len(str(e)) > 50 else str(e), fill='black', font=font) | |
| return error_image, None | |
| finally: | |
| # Clean up temporary directory (but not the output file) | |
| if temp_dir and os.path.exists(temp_dir): | |
| try: | |
| shutil.rmtree(temp_dir) | |
| except: | |
| pass # Ignore cleanup errors | |
| # Gradio Interface | |
| with gr.Blocks(title="GeoTIFF Clipper", theme=gr.themes.Soft()) as app: | |
| gr.Markdown(""" | |
| # 🛰️ GeoTIFF Clipping with Shapefile | |
| This tool allows you to clip GeoTIFF images using shapefiles as clipping masks. | |
| **Instructions:** | |
| 1. Upload a GeoTIFF file (.tif) | |
| 2. Upload a ZIP file containing the shapefile (.shp, .dbf, .shx, .prj) | |
| 3. Click "Execute Clipping" | |
| 4. View the result and download the clipped file | |
| """) | |
| with gr.Row(): | |
| with gr.Column(): | |
| geotiff_input = gr.File( | |
| label="📁 GeoTIFF File (.tif)", | |
| file_types=[".tif", ".tiff"], | |
| type="filepath" | |
| ) | |
| shapefile_input = gr.File( | |
| label="📁 Shapefile ZIP (.zip)", | |
| file_types=[".zip"], | |
| type="filepath" | |
| ) | |
| run_button = gr.Button("🚀 Execute Clipping", variant="primary", size="lg") | |
| with gr.Row(): | |
| with gr.Column(): | |
| preview_output = gr.Image( | |
| label="🖼️ Result Preview", | |
| type="pil", | |
| height=400 | |
| ) | |
| with gr.Column(): | |
| tif_output = gr.File( | |
| label="💾 Download Clipped GeoTIFF", | |
| type="filepath" | |
| ) | |
| # Connect function to button | |
| run_button.click( | |
| fn=clip_geotiff, | |
| inputs=[geotiff_input, shapefile_input], | |
| outputs=[preview_output, tif_output] | |
| ) | |
| gr.Markdown(""" | |
| --- | |
| **Notes:** | |
| - The shapefile ZIP must contain at least .shp, .dbf and .shx files | |
| - Coordinate systems will be automatically adjusted if necessary | |
| - Preview is automatically generated for visualization | |
| """) | |
| # Configure for Hugging Face Spaces | |
| if __name__ == "__main__": | |
| app.launch( | |
| server_name="0.0.0.0", # Required for Hugging Face Spaces | |
| server_port=7860, # Default Hugging Face Spaces port | |
| share=False # Don't create additional public link | |
| ) | |