""" 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("""
""") # --------------------------------------------------------------------------- # 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()