app.py

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
    )