Update app.py

app.py

@@ -1,34 +1,56 @@
 """
 Hugging Face Space Entry Point for Project Vaayu
-Serves both an interactive Gradio interface and REST API endpoints for Vercel.
+
+Serves:
+1. Gradio interactive interface
+2. FastAPI REST API endpoints for the Vercel frontend
+3. ZeroGPU-compatible model inference
 """
 
+# ============================================================
+# IMPORTANT:
+# Import `spaces` before other ML/GPU-related imports.
+# ============================================================
+
+import spaces
+
 import os
 import io
 import base64
 import json
-import numpy as np
 from pathlib import Path
+
+import numpy as np
 from PIL import Image
+
 import gradio as gr
+
 from fastapi import FastAPI, File, UploadFile, Form
 from fastapi.middleware.cors import CORSMiddleware
-from fastapi.responses import JSONResponse, StreamingResponse
-from fastapi.staticfiles import StaticFiles
+from fastapi.responses import JSONResponse
 
-try:
-    import spaces
-    has_spaces = True
-except ImportError:
-    has_spaces = False
+# ============================================================
+# Existing Vaayu inference engine
+# ============================================================
 
-# Import existing inference engine
 from inference import VaayuSegmentationModel
 
-# Initialize FastAPI App
-app = FastAPI(title="Vaayu Geospatial AI API")
 
-# Enable CORS for Vercel & localhost
+# ============================================================
+# FASTAPI APPLICATION
+# ============================================================
+
+app = FastAPI(
+    title="Vaayu Geospatial AI API",
+    description="AI-powered geospatial feature extraction using UNet++",
+    version="1.0.0",
+)
+
+
+# ============================================================
+# CORS
+# ============================================================
+
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -37,91 +59,346 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# Initialize Model Instance
+
+# ============================================================
+# MODEL
+# ============================================================
+
 model_instance = None
 
+
 def get_model():
+    """
+    Lazily initialize the Vaayu segmentation model.
+
+    The model is initialized only when inference is actually requested.
+    """
     global model_instance
+
     if model_instance is None:
-        model_instance = VaayuSegmentationModel(models_dir="models")
+        print("Loading Vaayu segmentation model...")
+
+        model_instance = VaayuSegmentationModel(
+            models_dir="models"
+        )
+
+        print("Vaayu segmentation model loaded successfully.")
+
     return model_instance
 
+
+# ============================================================
+# UTILITY FUNCTIONS
+# ============================================================
+
 def numpy_to_base64(arr, fmt="PNG"):
+    """
+    Convert a NumPy image array to a base64 encoded image.
+    """
+
+    # Make sure array is uint8
+    if arr.dtype != np.uint8:
+        arr = np.clip(arr, 0, 255).astype(np.uint8)
+
     img = Image.fromarray(arr)
-    buf = io.BytesIO()
-    img.save(buf, format=fmt)
-    return base64.b64encode(buf.getvalue()).decode("utf-8")
 
-# --- REST API Endpoints for Vercel Frontend ---
+    buffer = io.BytesIO()
+
+    img.save(
+        buffer,
+        format=fmt
+    )
+
+    return base64.b64encode(
+        buffer.getvalue()
+    ).decode("utf-8")
+
+
+# ============================================================
+# HEALTH CHECK
+# ============================================================
 
 @app.get("/api/health")
 def health():
-    return {"status": "ok", "message": "Vaayu API on Hugging Face is running"}
+
+    return {
+        "status": "ok",
+        "message": "Vaayu API on Hugging Face is running"
+    }
+
+
+# ============================================================
+# DATASET STATS
+# ============================================================
 
 @app.get("/api/dataset-stats")
 def dataset_stats():
+
+    # --------------------------------------------------------
+    # Default fallback statistics
+    # --------------------------------------------------------
+
     stats = {
         "village": "Lalpur, Gujarat",
         "ecw_file": "ortho_lalpur(511638)_3857.ecw",
         "ecw_size_mb": 21.9,
         "resolution": "50 cm per pixel",
+
         "total_buildings": 550,
         "total_roads": 49,
+
         "roof_types": {
             "RCC / Concrete": 312,
             "Tiled / Terracotta": 98,
             "Tin / Metal Sheet": 85,
             "Others": 55
         },
+
         "avg_structure_area_m2": 87.3,
     }
-    
+
+    # --------------------------------------------------------
+    # Try reading actual shapefiles
+    # --------------------------------------------------------
+
     try:
+
         import shapefile
-        bldg_path = Path("sample_data/ShapeFiles/LalPur/Gujarat_Build_Up_Area_Type.shp")
-        road_path = Path("sample_data/ShapeFiles/LalPur/Gujarat_Road.shp")
-        
+
+        bldg_path = Path(
+            "sample_data/ShapeFiles/LalPur/Gujarat_Build_Up_Area_Type.shp"
+        )
+
+        road_path = Path(
+            "sample_data/ShapeFiles/LalPur/Gujarat_Road.shp"
+        )
+
         if bldg_path.exists():
-            sf_b = shapefile.Reader(str(bldg_path))
-            sf_r = shapefile.Reader(str(road_path))
+
+            sf_b = shapefile.Reader(
+                str(bldg_path)
+            )
+
+            if road_path.exists():
+                sf_r = shapefile.Reader(
+                    str(road_path)
+                )
+            else:
+                sf_r = None
+
             recs = sf_b.records()
-            roof_counts = {"RCC / Concrete": 0, "Tiled / Terracotta": 0, "Tin / Metal Sheet": 0, "Others": 0}
+
+            roof_counts = {
+                "RCC / Concrete": 0,
+                "Tiled / Terracotta": 0,
+                "Tin / Metal Sheet": 0,
+                "Others": 0
+            }
+
             areas = []
+
             table_data = []
-            
-            for idx, r in enumerate(recs):
-                d = r.as_dict()
-                rt = d.get("roof_type", 1)
-                area = float(d.get("area_sqm", 0.0) or d.get("area_geo", 0.0) or 0.0)
+
+            # ------------------------------------------------
+            # Process building records
+            # ------------------------------------------------
+
+            for idx, record in enumerate(recs):
+
+                try:
+                    d = record.as_dict()
+                except Exception:
+                    d = {}
+
+                rt = d.get(
+                    "roof_type",
+                    1
+                )
+
+                area = float(
+                    d.get(
+                        "area_sqm",
+                        0.0
+                    )
+                    or
+                    d.get(
+                        "area_geo",
+                        0.0
+                    )
+                    or 0.0
+                )
+
                 areas.append(area)
-                
+
+                # --------------------------------------------
+                # Roof classification
+                # --------------------------------------------
+
                 if rt == 1:
-                    roof_counts["RCC / Concrete"] += 1
+
+                    roof_counts[
+                        "RCC / Concrete"
+                    ] += 1
+
+                    roof_type_name = "RCC / Concrete"
+
                 elif rt == 2:
-                    roof_counts["Tiled / Terracotta"] += 1
+
+                    roof_counts[
+                        "Tiled / Terracotta"
+                    ] += 1
+
+                    roof_type_name = "Tiled / Terracotta"
+
                 elif rt == 3:
-                    roof_counts["Tin / Metal Sheet"] += 1
+
+                    roof_counts[
+                        "Tin / Metal Sheet"
+                    ] += 1
+
+                    roof_type_name = "Tin / Metal Sheet"
+
                 else:
-                    roof_counts["Others"] += 1
-                
+
+                    roof_counts[
+                        "Others"
+                    ] += 1
+
+                    roof_type_name = "Others"
+
+                # --------------------------------------------
+                # First 20 records for frontend table
+                # --------------------------------------------
+
                 if idx < 20:
-                    table_data.append({
-                        "objectId": d.get("objectid", idx),
-                        "roofType": "RCC / Concrete" if rt == 1 else "Tiled / Terracotta" if rt == 2 else "Tin / Metal Sheet" if rt == 3 else "Others",
-                        "areaSqm": round(area, 2),
-                        "propertyId": d.get("property_i") or d.get("uniq_id") or f"LAL-PROP-{d.get('objectid', idx)}",
-                        "village": d.get("village_na") or "Lalpur"
-                    })
-            
-            stats["total_buildings"] = len(recs)
-            stats["total_roads"] = len(sf_r)
-            stats["roof_types"] = roof_counts
-            stats["avg_structure_area_m2"] = round(float(np.mean(areas)), 1) if areas else 0
-            stats["sample_records"] = table_data
+
+                    table_data.append(
+                        {
+                            "objectId": d.get(
+                                "objectid",
+                                idx
+                            ),
+
+                            "roofType": roof_type_name,
+
+                            "areaSqm": round(
+                                area,
+                                2
+                            ),
+
+                            "propertyId": (
+                                d.get("property_i")
+                                or
+                                d.get("uniq_id")
+                                or
+                                f"LAL-PROP-{d.get('objectid', idx)}"
+                            ),
+
+                            "village": (
+                                d.get("village_na")
+                                or
+                                "Lalpur"
+                            )
+                        }
+                    )
+
+            # ------------------------------------------------
+            # Update statistics
+            # ------------------------------------------------
+
+            stats[
+                "total_buildings"
+            ] = len(recs)
+
+            if sf_r is not None:
+
+                stats[
+                    "total_roads"
+                ] = len(sf_r)
+
+            stats[
+                "roof_types"
+            ] = roof_counts
+
+            stats[
+                "avg_structure_area_m2"
+            ] = (
+                round(
+                    float(np.mean(areas)),
+                    1
+                )
+                if areas
+                else 0
+            )
+
+            stats[
+                "sample_records"
+            ] = table_data
+
     except Exception as e:
-        stats["shapefile_note"] = f"Fallback stats: {e}"
-    
-    return JSONResponse(content=stats)
+
+        stats[
+            "shapefile_note"
+        ] = f"Fallback stats: {str(e)}"
+
+    return JSONResponse(
+        content=stats
+    )
+
+
+# ============================================================
+# GPU INFERENCE FUNCTION
+# ============================================================
+#
+# IMPORTANT:
+# This is the critical ZeroGPU change.
+#
+# Do NOT do:
+#
+#     gradio_predict = spaces.GPU(function)
+#
+# Instead use:
+#
+#     @spaces.GPU(...)
+#     def function(...):
+#
+# ============================================================
+
+
+@spaces.GPU(duration=120)
+def run_prediction(
+    image_path,
+    feature="buildings",
+    threshold=0.5
+):
+    """
+    GPU-enabled inference function.
+
+    This function is shared by:
+    - Gradio
+    - FastAPI REST API
+    """
+
+    print(
+        f"Running Vaayu inference: "
+        f"feature={feature}, "
+        f"threshold={threshold}"
+    )
+
+    model = get_model()
+
+    result = model.predict(
+        image_path,
+        feature=feature,
+        threshold=threshold
+    )
+
+    return result
+
+
+# ============================================================
+# REST API PREDICT
+# ============================================================
 
 @app.post("/api/predict")
 async def predict_api(
@@ -130,97 +407,499 @@ async def predict_api(
     sample_tile: str = Form(None),
     image: UploadFile = File(None)
 ):
-    m = get_model()
-    
-    if sample_tile:
-        image_path = str(Path("sample_data/tiles") / sample_tile)
-        if not Path(image_path).exists():
-            return JSONResponse(status_code=404, content={"error": f"Tile not found: {sample_tile}"})
-    elif image:
-        temp_path = Path("temp_hf_upload.png")
-        contents = await image.read()
-        with open(temp_path, "wb") as f:
-            f.write(contents)
-        image_path = str(temp_path)
-    else:
-        return JSONResponse(status_code=400, content={"error": "No image or sample_tile provided"})
+
+    # --------------------------------------------------------
+    # Validate feature
+    # --------------------------------------------------------
+
+    allowed_features = {
+        "buildings",
+        "roads",
+        "water_bodies"
+    }
+
+    if feature not in allowed_features:
+
+        return JSONResponse(
+            status_code=400,
+            content={
+                "error": (
+                    f"Invalid feature '{feature}'. "
+                    f"Allowed values: "
+                    f"{', '.join(allowed_features)}"
+                )
+            }
+        )
+
+    # --------------------------------------------------------
+    # Validate threshold
+    # --------------------------------------------------------
+
+    if threshold < 0.0 or threshold > 1.0:
+
+        return JSONResponse(
+            status_code=400,
+            content={
+                "error": "Threshold must be between 0 and 1."
+            }
+        )
+
+    # --------------------------------------------------------
+    # Determine image source
+    # --------------------------------------------------------
+
+    temp_path = None
 
     try:
-        res = m.predict(image_path, feature=feature, threshold=threshold)
-        
+
+        if sample_tile:
+
+            image_path = (
+                Path("sample_data/tiles")
+                / sample_tile
+            )
+
+            if not image_path.exists():
+
+                return JSONResponse(
+                    status_code=404,
+                    content={
+                        "error": (
+                            f"Tile not found: "
+                            f"{sample_tile}"
+                        )
+                    }
+                )
+
+            image_path = str(
+                image_path
+            )
+
+        elif image:
+
+            # ------------------------------------------------
+            # Save uploaded image
+            # ------------------------------------------------
+
+            extension = (
+                Path(image.filename).suffix
+                if image.filename
+                else ".png"
+            )
+
+            temp_path = Path(
+                f"temp_hf_upload{extension}"
+            )
+
+            contents = await image.read()
+
+            with open(
+                temp_path,
+                "wb"
+            ) as f:
+
+                f.write(contents)
+
+            image_path = str(
+                temp_path
+            )
+
+        else:
+
+            return JSONResponse(
+                status_code=400,
+                content={
+                    "error": (
+                        "No image or sample_tile provided"
+                    )
+                }
+            )
+
+        # ----------------------------------------------------
+        # GPU inference
+        # ----------------------------------------------------
+
+        result = run_prediction(
+            image_path,
+            feature=feature,
+            threshold=threshold
+        )
+
+        # ----------------------------------------------------
+        # Build API response
+        # ----------------------------------------------------
+
         response_data = {
-            "feature": res["feature"],
-            "detected_ratio": float(res["detected_ratio"]),
-            "detected_pixels": int(res["detected_pixels"]),
-            "total_pixels": int(res["total_pixels"]),
-            "original_image": numpy_to_base64(res["original_image"]),
-            "mask": numpy_to_base64(res["mask"]),
-            "overlay": numpy_to_base64(res["overlay"]),
+
+            "feature": result["feature"],
+
+            "detected_ratio": float(
+                result["detected_ratio"]
+            ),
+
+            "detected_pixels": int(
+                result["detected_pixels"]
+            ),
+
+            "total_pixels": int(
+                result["total_pixels"]
+            ),
+
+            "original_image": numpy_to_base64(
+                result["original_image"]
+            ),
+
+            "mask": numpy_to_base64(
+                result["mask"]
+            ),
+
+            "overlay": numpy_to_base64(
+                result["overlay"]
+            )
         }
-        
-        if res.get("rooftop_analysis"):
-            ra = res["rooftop_analysis"]
-            response_data["rooftop_analysis"] = {
-                "rooftop_counts": ra["rooftop_counts"],
-                "total_buildings": ra["total_buildings"],
-                "total_builtup_m2": ra["total_builtup_m2"],
-                "rcc_area_m2": ra["rcc_area_m2"],
-                "annual_solar_kwh": ra["annual_solar_kwh"],
-                "annual_property_tax_inr": ra["annual_property_tax_inr"],
-                "classification_map": numpy_to_base64(ra["classification_map"]),
-                "building_details": ra["building_details"][:50],
+
+        # ----------------------------------------------------
+        # Rooftop analysis
+        # ----------------------------------------------------
+
+        if result.get("rooftop_analysis"):
+
+            rooftop = result[
+                "rooftop_analysis"
+            ]
+
+            response_data[
+                "rooftop_analysis"
+            ] = {
+
+                "rooftop_counts": rooftop[
+                    "rooftop_counts"
+                ],
+
+                "total_buildings": rooftop[
+                    "total_buildings"
+                ],
+
+                "total_builtup_m2": rooftop[
+                    "total_builtup_m2"
+                ],
+
+                "rcc_area_m2": rooftop[
+                    "rcc_area_m2"
+                ],
+
+                "annual_solar_kwh": rooftop[
+                    "annual_solar_kwh"
+                ],
+
+                "annual_property_tax_inr": rooftop[
+                    "annual_property_tax_inr"
+                ],
+
+                "classification_map":
+                    numpy_to_base64(
+                        rooftop[
+                            "classification_map"
+                        ]
+                    ),
+
+                "building_details":
+                    rooftop[
+                        "building_details"
+                    ][:50]
             }
-        
-        return JSONResponse(content=response_data)
+
+        return JSONResponse(
+            content=response_data
+        )
+
     except Exception as e:
-        return JSONResponse(status_code=500, content={"error": str(e)})
 
-# --- Gradio UI Component ---
+        print(
+            f"Prediction error: {str(e)}"
+        )
+
+        return JSONResponse(
+            status_code=500,
+            content={
+                "error": str(e)
+            }
+        )
+
+    finally:
+
+        # ----------------------------------------------------
+        # Delete temporary uploaded image
+        # ----------------------------------------------------
+
+        if temp_path is not None:
+
+            try:
+
+                if temp_path.exists():
+                    temp_path.unlink()
+
+            except Exception as cleanup_error:
+
+                print(
+                    f"Temporary file cleanup failed: "
+                    f"{cleanup_error}"
+                )
+
+
+# ============================================================
+# GRADIO PREDICTION
+# ============================================================
+
+def _raw_gradio_predict(
+    input_image,
+    feature_type,
+    confidence_threshold
+):
 
-def _raw_gradio_predict(input_image, feature_type, confidence_threshold):
     if input_image is None:
-        return None, None, "Please upload an orthophoto image."
-    
-    # Save temp image for inference
-    temp_p = "temp_gradio_input.png"
-    Image.fromarray(input_image).save(temp_p)
-    
-    feature_key = "buildings" if "Building" in feature_type else "roads" if "Road" in feature_type else "water_bodies"
-    
-    m = get_model()
-    res = m.predict(temp_p, feature=feature_key, threshold=confidence_threshold)
-    
-    summary = f"Detected Pixels: {res['detected_pixels']:,} ({res['detected_ratio']*100:.1f}% coverage)\n"
-    if res.get("rooftop_analysis"):
-        ra = res["rooftop_analysis"]
-        summary += f"Total Identified Parcels: {ra['total_buildings']}\n"
-        summary += f"Total Built-up Area: {ra['total_builtup_m2']:,.1f} m²\n"
-        summary += f"Est. Solar Potential: {ra['annual_solar_kwh']:,.0f} kWh/yr\n"
-        summary += f"Est. Property Tax: ₹{ra['annual_property_tax_inr']:,.0f}/yr\n"
-
-    return res["overlay"], res["mask"], summary
-
-if has_spaces:
-    gradio_predict = spaces.GPU(_raw_gradio_predict)
-else:
-    gradio_predict = _raw_gradio_predict
-
-# Build Gradio Interface
+
+        return (
+            None,
+            None,
+            "Please upload an orthophoto image."
+        )
+
+    # --------------------------------------------------------
+    # Save uploaded Gradio image
+    # --------------------------------------------------------
+
+    temp_path = Path(
+        "temp_gradio_input.png"
+    )
+
+    try:
+
+        # Gradio can return either a NumPy array
+        # or PIL image depending on configuration.
+
+        if isinstance(
+            input_image,
+            Image.Image
+        ):
+
+            image = input_image
+
+        else:
+
+            image = Image.fromarray(
+                input_image
+            )
+
+        image.save(
+            temp_path
+        )
+
+        # ----------------------------------------------------
+        # Convert UI feature name to model feature key
+        # ----------------------------------------------------
+
+        if "Building" in feature_type:
+
+            feature_key = "buildings"
+
+        elif "Road" in feature_type:
+
+            feature_key = "roads"
+
+        else:
+
+            feature_key = "water_bodies"
+
+        # ----------------------------------------------------
+        # GPU inference
+        # ----------------------------------------------------
+
+        result = run_prediction(
+            str(temp_path),
+            feature=feature_key,
+            threshold=confidence_threshold
+        )
+
+        # ----------------------------------------------------
+        # Analytics summary
+        # ----------------------------------------------------
+
+        summary = (
+            f"Detected Pixels: "
+            f"{result['detected_pixels']:,} "
+            f"({result['detected_ratio'] * 100:.1f}% coverage)\n"
+        )
+
+        if result.get("rooftop_analysis"):
+
+            rooftop = result[
+                "rooftop_analysis"
+            ]
+
+            summary += (
+                f"Total Identified Parcels: "
+                f"{rooftop['total_buildings']}\n"
+            )
+
+            summary += (
+                f"Total Built-up Area: "
+                f"{rooftop['total_builtup_m2']:,.1f} m²\n"
+            )
+
+            summary += (
+                f"Est. Solar Potential: "
+                f"{rooftop['annual_solar_kwh']:,.0f} kWh/yr\n"
+            )
+
+            summary += (
+                f"Est. Property Tax: "
+                f"₹{rooftop['annual_property_tax_inr']:,.0f}/yr\n"
+            )
+
+        return (
+            result["overlay"],
+            result["mask"],
+            summary
+        )
+
+    except Exception as e:
+
+        print(
+            f"Gradio prediction error: {str(e)}"
+        )
+
+        return (
+            None,
+            None,
+            f"Inference error: {str(e)}"
+        )
+
+    finally:
+
+        try:
+
+            if temp_path.exists():
+                temp_path.unlink()
+
+        except Exception:
+            pass
+
+
+# ============================================================
+# ZERO GPU GRADIO FUNCTION
+# ============================================================
+#
+# This is the important part that fixes:
+#
+# "No @spaces.GPU function detected during startup"
+#
+# ============================================================
+
+@spaces.GPU(duration=120)
+def gradio_predict(
+    input_image,
+    feature_type,
+    confidence_threshold
+):
+
+    return _raw_gradio_predict(
+        input_image,
+        feature_type,
+        confidence_threshold
+    )
+
+
+# ============================================================
+# GRADIO INTERFACE
+# ============================================================
+
 demo = gr.Interface(
+
     fn=gradio_predict,
+
     inputs=[
-        gr.Image(label="Upload Orthophoto Tile (PNG / JPG / GeoTIFF)"),
-        gr.Radio(["Building Footprints & Roofs", "Road Infrastructure", "Waterbodies & Tanks"], label="Geospatial Feature", value="Building Footprints & Roofs"),
-        gr.Slider(minimum=0.1, maximum=0.9, value=0.5, step=0.05, label="Confidence Threshold"),
+
+        gr.Image(
+            label="Upload Orthophoto Tile (PNG / JPG / GeoTIFF)",
+            type="numpy"
+        ),
+
+        gr.Radio(
+            [
+                "Building Footprints & Roofs",
+                "Road Infrastructure",
+                "Waterbodies & Tanks"
+            ],
+
+            label="Geospatial Feature",
+
+            value="Building Footprints & Roofs"
+        ),
+
+        gr.Slider(
+            minimum=0.1,
+            maximum=0.9,
+            value=0.5,
+            step=0.05,
+            label="Confidence Threshold"
+        )
     ],
+
     outputs=[
-        gr.Image(label="AI Demarcation Overlay"),
-        gr.Image(label="Binary Segmentation Mask"),
-        gr.Textbox(label="Analytics & Governance Telemetry"),
+
+        gr.Image(
+            label="AI Demarcation Overlay"
+        ),
+
+        gr.Image(
+            label="Binary Segmentation Mask"
+        ),
+
+        gr.Textbox(
+            label="Analytics & Governance Telemetry"
+        )
     ],
-    title="Project Vaayu — Geospatial AI Intelligence",
-    description="Deep neural feature extraction from 50cm drone orthophotos using UNet++ ResNet-34.",
+
+    title=(
+        "Project Vaayu — "
+        "Geospatial AI Intelligence"
+    ),
+
+    description=(
+        "Deep neural feature extraction "
+        "from 50cm drone orthophotos "
+        "using UNet++ ResNet-34."
+    )
+)
+
+
+# ============================================================
+# MOUNT GRADIO INTO FASTAPI
+# ============================================================
+
+app = gr.mount_gradio_app(
+    app,
+    demo,
+    path="/"
 )
 
-# Mount Gradio UI inside FastAPI App
-app = gr.mount_gradio_app(app, demo, path="/")
+
+# ============================================================
+# LOCAL DEVELOPMENT
+# ============================================================
+
+if __name__ == "__main__":
+
+    import uvicorn
+
+    uvicorn.run(
+        app,
+        host="0.0.0.0",
+        port=int(
+            os.environ.get(
+                "PORT",
+                7860
+            )
+        )
+    )