app.py

"""
DepthWizard — Hugging Face Spaces (ZeroGPU compliant)

Architecture:
  1. demo.queue().launch(prevent_thread_lock=True)
       → Gradio creates the ACTUAL served app (demo.server_app), sets up ZeroGPU,
         and starts uvicorn in a background thread. Main thread continues.
  2. We insert our routes at index 0 of demo.server_app.router.routes
       → They appear BEFORE Gradio's catch-all, so /process is always matched first.
  3. threading.Event().wait()
       → Main thread blocks forever so the Space stays alive.

ZeroGPU compliance:
  • demo.launch() is called  ✓
  • @spaces.GPU decorates both inference functions  ✓
  • predict_gpu is wired to _hidden_btn.click() inside the Gradio Blocks  ✓
"""

import os
import io
import json
import base64
import tempfile
import threading
import time
import numpy as np
from PIL import Image
import matplotlib.cm as cm
from typing import Optional, Dict, Any

import gradio as gr
import spaces

from fastapi import UploadFile, File, Form, HTTPException
from fastapi.responses import HTMLResponse, JSONResponse, FileResponse
from fastapi.middleware.cors import CORSMiddleware
from starlette.middleware import Middleware
from fastapi.staticfiles import StaticFiles
from fastapi.routing import APIRoute
from starlette.concurrency import run_in_threadpool

from process_image import process_image, export_dsm, export_slope
from depth.depth_model import load_model

# ---------------------------------------------------------------------------
# Global model cache & export store
# ---------------------------------------------------------------------------
GLOBAL_MODEL = None
GLOBAL_PROCESSOR = None
GLOBAL_DEVICE = None
EXPORT_CACHE: Dict[str, Dict[str, Any]] = {}


def get_model():
    global GLOBAL_MODEL, GLOBAL_PROCESSOR, GLOBAL_DEVICE
    if GLOBAL_MODEL is None:
        print("[API Server] Preloading Depth Anything V2 model...", flush=True)
        GLOBAL_MODEL, GLOBAL_PROCESSOR, GLOBAL_DEVICE = load_model()
        print(f"[API Server] Model ready on device: {GLOBAL_DEVICE}", flush=True)
    return GLOBAL_MODEL, GLOBAL_PROCESSOR, GLOBAL_DEVICE


# @spaces.GPU
# def run_depth_inference(
#     path: str,
#     gcps: Optional[list] = None,
#     dem_path: Optional[str] = None,
#     use_shadows: bool = True,
#     a_prior: Optional[float] = None,
#     lambda_prior: float = 0.0,
#     terrain_percentile: float = 25.0,
#     calibration_method: str = "linear"
# ) -> Dict[str, Any]:
#     """Execute ML inference inside ZeroGPU isolated worker."""
#     model, processor, device = get_model()
#     return process_image(
#         path=path,
#         gcps=gcps,
#         dem_path=dem_path,
#         use_shadows=use_shadows,
#         model=model,
#         processor=processor,
#         device=device,
#         a_prior=a_prior,
#         lambda_prior=lambda_prior,
#         terrain_percentile=terrain_percentile,
#         calibration_method=calibration_method
#     )



@spaces.GPU
def predict_gpu(
    path: str,
    gcps: Optional[list] = None,
    dem_path: Optional[str] = None,
    use_shadows: bool = True,
    a_prior: Optional[float] = None,
    lambda_prior: float = 0.0,
    terrain_percentile: float = 25.0,
    calibration_method: str = "linear"
) -> Dict[str, Any]:
    """Run DepthWizard inference through a registered ZeroGPU event."""
    model, processor, device = get_model()

    return process_image(
        path=path,
        gcps=gcps,
        dem_path=dem_path,
        use_shadows=use_shadows,
        model=model,
        processor=processor,
        device=device,
        a_prior=a_prior,
        lambda_prior=lambda_prior,
        terrain_percentile=terrain_percentile,
        calibration_method=calibration_method
    )



# ---------------------------------------------------------------------------
# JSON / array helpers
# ---------------------------------------------------------------------------

def sanitize_for_json(obj: Any) -> Any:
    if obj is None or isinstance(obj, (bool, int, float, str)):
        if isinstance(obj, float) and (np.isnan(obj) or np.isinf(obj)):
            return None
        return obj
    if isinstance(obj, (np.floating, np.float32, np.float64)):
        v = float(obj)
        return None if (np.isnan(v) or np.isinf(v)) else v
    if isinstance(obj, (np.integer, np.int32, np.int64)):
        return int(obj)
    if isinstance(obj, np.bool_):
        return bool(obj)
    if isinstance(obj, np.ndarray):
        return [sanitize_for_json(v) for v in obj.tolist()]
    if isinstance(obj, (list, tuple)):
        return [sanitize_for_json(v) for v in obj]
    if isinstance(obj, dict):
        return {str(k): sanitize_for_json(v) for k, v in obj.items()}
    return str(obj)


def downsample_array(arr: np.ndarray, target_w: int = 256, target_h: int = 256) -> np.ndarray:
    img = Image.fromarray(arr.astype(np.float32), mode="F")
    return np.array(img.resize((target_w, target_h), Image.Resampling.BILINEAR), dtype=np.float32)


def encode_rgb_to_base64_jpeg(rgb: np.ndarray, max_size: int = 1280, quality: int = 92) -> str:
    h, w = rgb.shape[:2]
    scale = min(1.0, max_size / max(h, w))
    new_w = max(1, int(round(w * scale)))
    new_h = max(1, int(round(h * scale)))
    img = Image.fromarray(rgb)
    if (new_w, new_h) != (w, h):
        img = img.resize((new_w, new_h), Image.Resampling.LANCZOS)
    buf = io.BytesIO()
    img.save(buf, format="JPEG", quality=quality)
    return f"data:image/jpeg;base64,{base64.b64encode(buf.getvalue()).decode()}"


def encode_colormap_to_base64_jpeg(
    arr: np.ndarray, cmap_name: str = "turbo", max_size: int = 1280,
    invert: bool = False, quality: int = 90,
    vmin: Optional[float] = None, vmax: Optional[float] = None,
    hillshade: Optional[np.ndarray] = None
) -> str:
    valid = np.isfinite(arr)
    if not np.any(valid):
        arr_norm = np.zeros_like(arr, dtype=np.float32)
    else:
        lo = vmin if vmin is not None else float(np.percentile(arr[valid], 1.0))
        hi = vmax if vmax is not None else float(np.percentile(arr[valid], 99.0))
        span = hi - lo
        arr_norm = np.zeros_like(arr, dtype=np.float32) if span < 1e-7 else np.clip((arr - lo) / span, 0.0, 1.0)
    if invert:
        arr_norm = 1.0 - arr_norm
    cmap = cm.get_cmap(cmap_name)
    rgb_f = (cmap(arr_norm)[:, :, :3] * 255.0)

    # Blend with hillshade for realistic 3D relief
    if hillshade is not None and hillshade.shape == arr.shape:
        hs_norm = hillshade.astype(np.float32) / 255.0
        hs_3d = np.repeat(hs_norm[:, :, np.newaxis], 3, axis=2)
        rgb_f = rgb_f * (0.55 + 0.45 * hs_3d)

    rgb = np.clip(rgb_f, 0.0, 255.0).astype(np.uint8)
    h, w = rgb.shape[:2]
    scale = min(1.0, max_size / max(h, w))
    new_w = max(1, int(round(w * scale)))
    new_h = max(1, int(round(h * scale)))
    img = Image.fromarray(rgb)
    if (new_w, new_h) != (w, h):
        img = img.resize((new_w, new_h), Image.Resampling.LANCZOS)
    buf = io.BytesIO()
    img.save(buf, format="JPEG", quality=quality)
    return f"data:image/jpeg;base64,{base64.b64encode(buf.getvalue()).decode()}"


# ---------------------------------------------------------------------------
# File paths
# ---------------------------------------------------------------------------
_DIR = os.path.dirname(os.path.abspath(__file__))
_DEMO_HTML = os.path.join(_DIR, "demo.html")
if not os.path.isfile(_DEMO_HTML):
    _DEMO_HTML = os.path.join(_DIR, "m6_dashboard.html")


# ---------------------------------------------------------------------------
# Endpoint functions (plain async functions — added via add_api_route later)
# ---------------------------------------------------------------------------

def health_check():
    return {"status": "online", "service": "DepthWizard Engine"}


def serve_demo():
    with open(_DEMO_HTML, "r", encoding="utf-8") as f:
        return HTMLResponse(content=f.read(), headers={"Cache-Control": "no-cache"})


def download_export(session_id: str):
    if session_id not in EXPORT_CACHE:
        raise HTTPException(status_code=404, detail="Export session not found.")
    item = EXPORT_CACHE[session_id]
    if not os.path.isfile(item["path"]):
        raise HTTPException(status_code=404, detail="Export file missing.")
    return FileResponse(path=item["path"], filename=item["filename"], media_type="application/octet-stream")


async def process_image_endpoint(
    file: UploadFile = File(...),
    gcps: Optional[str] = Form(None),
    dem_file: Optional[UploadFile] = File(None),
    dem_path: Optional[str] = Form(None),
    use_shadows: bool = Form(True),
    visual_size: int = Form(256),
    a_prior: Optional[float] = Form(None),
    lambda_prior: float = Form(0.0),
    terrain_percentile: float = Form(25.0),
    calibration_method: str = Form("linear")
):
    parsed_gcps = None
    if gcps:
        try:
            parsed_gcps = json.loads(gcps)
        except Exception:
            pass

    filename = file.filename or "upload.png"
    ext = os.path.splitext(filename)[1].lower() or ".png"

    with tempfile.NamedTemporaryFile(delete=False, suffix=ext) as tmp:
        tmp.write(await file.read())
        temp_path = tmp.name

    temp_dem_path = None
    if dem_file is not None and dem_file.filename:
        dem_ext = os.path.splitext(dem_file.filename)[1].lower() or ".tif"
        with tempfile.NamedTemporaryFile(delete=False, suffix=dem_ext) as tmp_dem:
            tmp_dem.write(await dem_file.read())
            temp_dem_path = tmp_dem.name
    elif dem_path and os.path.isfile(dem_path):
        temp_dem_path = dem_path

    try:
        result = await run_in_threadpool(
            predict_gpu,
            path=temp_path,
            gcps=parsed_gcps,
            dem_path=temp_dem_path,
            use_shadows=use_shadows,
            a_prior=a_prior,
            lambda_prior=lambda_prior,
            terrain_percentile=terrain_percentile,
            calibration_method=calibration_method
        )

        h_orig, w_orig = result["height_map"].shape
        aspect = float(h_orig) / max(1, float(w_orig))
        if w_orig >= h_orig:
            v_w = 320
            v_h = max(32, int(round(320 * aspect)))
        else:
            v_h = 320
            v_w = max(32, int(round(320 / aspect)))

        h_sub = downsample_array(result["height_map"], target_w=v_w, target_h=v_h)
        s_sub = downsample_array(result["slope_map"], target_w=v_w, target_h=v_h)
        c_sub = downsample_array(result["confidence_map"], target_w=v_w, target_h=v_h)

        # High-resolution previews for crisp dashboard visualization
        rgb_b64   = encode_rgb_to_base64_jpeg(result["rgb"], max_size=1280, quality=92)
        depth_b64 = encode_colormap_to_base64_jpeg(result["depth_map"], "plasma", max_size=1280)
        dsm_b64   = encode_colormap_to_base64_jpeg(result["height_map"], "turbo", max_size=1280, hillshade=result.get("hillshade"))
        slope_b64 = encode_colormap_to_base64_jpeg(result["slope_map"], "magma", max_size=1280, vmin=0.0, vmax=45.0)
        conf_b64  = encode_colormap_to_base64_jpeg(result["confidence_map"], "viridis", max_size=1280, vmin=0.0, vmax=1.0)

        error_b64 = None
        val_metrics = None
        err = result.get("error_map")
        err_type = result.get("error_type", "estimated_uncertainty")

        if err is not None:
            valid_err = err[np.isfinite(err)]
            if len(valid_err) > 0:
                abs_err = np.abs(valid_err)
                h_var = float(np.var(result["height_map"]))
                r2_val = float(max(0.0, 1.0 - (float(np.var(valid_err)) / (h_var + 1e-6)))) if h_var > 1e-6 else 0.0
                val_metrics = {
                    "mae": float(np.mean(abs_err)),
                    "rmse": float(np.sqrt(np.mean(valid_err ** 2))),
                    "max_error": float(np.max(abs_err)),
                    "r2": r2_val,
                    "bias": float(np.mean(valid_err)),
                    "count": int(len(valid_err)),
                    "type": err_type
                }
            error_b64 = encode_colormap_to_base64_jpeg(np.abs(err), "inferno", max_size=1280)

        def _stat(arr, fn, default):
            f = arr[np.isfinite(arr)]
            return float(fn(f)) if len(f) > 0 else default

        h_min  = _stat(result["height_map"], np.min, 0.0)
        h_max  = _stat(result["height_map"], np.max, 1.0)
        s_min  = _stat(result["slope_map"],  np.min, 0.0)
        s_max  = _stat(result["slope_map"],  np.max, 90.0)
        c_mean = _stat(result["confidence_map"], np.mean, 1.0)

        calibrated    = bool(result.get("calibrated", False))
        georeferenced = bool(result.get("georeferenced", False))
        mode          = str(result.get("mode", "rdsm"))
        height_unit   = str(result.get("height_unit", "m" if calibrated else "relative"))

        mid_y = v_h // 2
        profile_pts = [
            {"x": x, "elevation": (None if np.isnan(float(h_sub[mid_y, x])) else float(h_sub[mid_y, x]))}
            for x in range(v_w)
        ]

        h_flat = h_sub.flatten()
        s_flat = s_sub.flatten()
        c_flat = c_sub.flatten()
        h_list = [None if np.isnan(v) else float(v) for v in h_flat]
        s_list = [None if np.isnan(v) else float(v) for v in s_flat]
        c_list = [None if np.isnan(v) else float(v) for v in c_flat]

        import uuid
        session_id = uuid.uuid4().hex[:12]
        os.makedirs("outputs/export", exist_ok=True)
        export_ext      = ".tif" if georeferenced else ".png"
        export_filename = f"depthwizard_{session_id}_{'dsm' if calibrated else 'rdsm'}{export_ext}"
        export_filepath = os.path.join("outputs/export", export_filename)
        export_dsm(result, export_filepath)
        EXPORT_CACHE[session_id] = {"path": export_filepath, "filename": export_filename}

        payload = {
            "session_id": session_id, "mode": mode,
            "calibrated": calibrated, "georeferenced": georeferenced,
            "height_unit": height_unit,
            "width": int(w_orig), "height": int(h_orig),
            "visual_width": int(v_w), "visual_height": int(v_h),
            "height_min": h_min, "height_max": h_max,
            "height_map": h_list,
            "rgb": rgb_b64, "depth_preview": depth_b64,
            "dsm_preview": dsm_b64, "slope_preview": slope_b64,
            "confidence_preview": conf_b64, "error_preview": error_b64,
            "validation": val_metrics,
            "slope_map": s_list, "confidence_map": c_list,
            "slope_min": s_min, "slope_max": s_max, "confidence_mean": c_mean,
            "elevation_profile": profile_pts,
            "crs": str(result["crs"]) if result["crs"] is not None else None,
            "export_url": f"/export/{session_id}",
            "metadata": result["metadata"]
        }
        return JSONResponse(content=sanitize_for_json(payload))

    except Exception as exc:
        import traceback
        traceback.print_exc()
        return JSONResponse(status_code=500, content={"error": str(exc), "detail": str(exc)})
    finally:
        try:
            if os.path.exists(temp_path):
                os.remove(temp_path)
            if temp_dem_path and os.path.exists(temp_dem_path) and temp_dem_path != (dem_path or ""):
                os.remove(temp_dem_path)
        except Exception:
            pass


# ---------------------------------------------------------------------------
# Gradio Blocks — ZeroGPU requires demo.launch()
# ---------------------------------------------------------------------------

custom_css = """
body, html { margin: 0; padding: 0; height: 100vh; overflow: hidden; background: #08080a; }
.gradio-container { max-width: 100% !important; margin: 0 !important; padding: 0 !important;
                    height: 100vh !important; overflow: hidden; background: #08080a; }
footer { display: none !important; }
#dw-wrap, #dw-wrap iframe { width: 100% !important; height: 100vh !important; border: none; display: block; }
"""

with gr.Blocks(title="DepthWizard — 3D Elevation Platform", css=custom_css, fill_height=True) as demo:
    _hidden_in = gr.Textbox(visible=False)
    _hidden_gcps = gr.JSON(visible=False)
    _hidden_dem_path = gr.Textbox(visible=False)
    _hidden_use_shadows = gr.Checkbox(value=True, visible=False)
    _hidden_a_prior = gr.Number(value=None, visible=False)
    _hidden_lambda_prior = gr.Number(value=0.0, visible=False)
    _hidden_terrain_percentile = gr.Number(value=25.0, visible=False)
    _hidden_calibration_method = gr.Textbox(value="linear", visible=False)
    _hidden_out = gr.JSON(visible=False)

    _hidden_btn = gr.Button(value="process", visible=False)
    _hidden_btn.click(
        fn=predict_gpu,
        inputs=[
            _hidden_in,
            _hidden_gcps,
            _hidden_dem_path,
            _hidden_use_shadows,
            _hidden_a_prior,
            _hidden_lambda_prior,
            _hidden_terrain_percentile,
            _hidden_calibration_method,
        ],
        outputs=[_hidden_out],
        api_name="predict_gpu",
    )


    # /demo.html is served by our custom route inserted below
    gr.HTML("""
        <div id="dw-wrap">
            <iframe src="/demo.html"
                    style="width:100%;height:100vh;border:none;display:block;"
                    allow="cross-origin-isolated">
            </iframe>
        </div>
    """)


# ---------------------------------------------------------------------------
# Main entry point
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    # Launch with prevent_thread_lock=True so main thread continues.
    # Do not preload the model here: ZeroGPU inference loads it inside run_depth_inference.
    # This starts uvicorn in a background thread AND sets demo.server_app
    # to the actual FastAPI instance that's being served.
    demo.queue().launch(
        server_name="0.0.0.0",
        server_port=7860,
        prevent_thread_lock=True,
        show_api=False,
        ssr_mode=False
    )

    # demo.server_app is the real served FastAPI app — NOT demo.app (which is
    # a separate pre-launch instance that gets replaced by launch()).
    # Insert our routes at index 0, BEFORE Gradio's catch-all route, so they
    # are matched first by Starlette's order-dependent router.
    server_app = demo.server_app
    print(f"[DepthWizard] Injecting custom routes into server_app: {type(server_app)}", flush=True)

    # Permit the separately hosted Netlify frontend to call this public API.
    # No cookies or credentials are used by this API.
    # Allow the deployed Netlify frontend to call this public API.
    server_app.user_middleware.insert(0, Middleware(
        CORSMiddleware,
        allow_origins=[
            "https://depthwizard-yukti-kriti.netlify.app"
        ],
        allow_credentials=False,
        allow_methods=["GET", "POST", "OPTIONS"],
        allow_headers=["*"]
    ))

    # Serve the bundled sample assets used by the embedded demo UI.
    if os.path.isdir(os.path.join(_DIR, "data")):
        server_app.mount("/data", StaticFiles(directory=os.path.join(_DIR, "data")), name="dw-data")
    if os.path.isdir(os.path.join(_DIR, "demoImages")):
        server_app.mount("/demoImages", StaticFiles(directory=os.path.join(_DIR, "demoImages")), name="dw-demo-images")

    # Build our route objects
    custom_routes = [
        APIRoute("/health",               health_check,             methods=["GET", "POST"]),
        APIRoute("/demo.html",            serve_demo,               methods=["GET"]),
        APIRoute("/demo",                 serve_demo,               methods=["GET"]),
        APIRoute("/export/{session_id}",  download_export,          methods=["GET"]),
        APIRoute("/api/process",          process_image_endpoint,   methods=["POST"]),
        APIRoute("/process",              process_image_endpoint,   methods=["POST"]),
    ]

    # Insert in reverse order so /process ends up at index 0
    for route in reversed(custom_routes):
        server_app.router.routes.insert(0, route)

    # Rebuild the routing table so the new routes are included
    server_app.router.on_startup = server_app.router.on_startup  # no-op touch
    # Force Starlette to rebuild its compiled middleware/route stack on next request
    if hasattr(server_app, 'middleware_stack'):
        server_app.middleware_stack = None

    print("[DepthWizard] Custom routes injected. /process is at index 0.", flush=True)

    # Block main thread forever — the server runs in the background thread
    threading.Event().wait()