app.py

"""Gradio application for single-image height estimation.

The demo is a thin wrapper around two things that already exist in this repo and
are shared with the offline CLI, so the Space cannot drift from the numbers the
CLI reports:

* `infer/` + `dwdata/` — the model and the exact preprocessing contract stored
  in the checkpoint.
* `viz/` and `viewer/` — the visualisation layer: the figure primitives
  (`viz.figures`), the mesh writers (`viz.mesh`), and the standalone dual-pane
  DSM inspector (`viewer/index.html` + `viewer/app.js`), which is inlined into
  an iframe here and works identically opened straight off disk.

The viewer is *not* re-implemented for the Space.  `viewer/index.html` is read at
import time and its `<script src=...>` tags are replaced with the vendored file
contents, giving one self-contained document.  The prediction is handed to it as
`window.DW_EMBED`, the third of the viewer's three input paths (file picker,
`?result=`, embedded payload).
"""

from __future__ import annotations

import base64
import html
import io
import json
import os
import re
import shutil
import tempfile
import threading
import time
from pathlib import Path

import gradio as gr
import numpy as np
import spaces
import torch
from PIL import Image

from config import CLASS_NAMES
from viz import classes as viz_classes
from viz import figures as viz_figures
from viz import mesh as viz_mesh

MODEL_PATH_ON_BUCKET = Path("/data/best.pt")
MODEL_PATH: str | None = None
MODEL = None
MODEL_SPEC = None
MODEL_CFG = None
MODEL_LOCK = threading.Lock()

_VIEWER_ROOT = Path(__file__).resolve().parent / "viewer"
# The iframe payload is capped independently of the prediction: a 1600 px scene
# inlined as base64 is a ~7 MB HTML string on every run, and the mesh vertex
# budget decimates it back down again anyway.  Downloads stay full resolution.
_VIEWER_MAX_SIDE = 1024
_PAYLOAD_SLOT = "__DW_EMBED_JSON__"


def _load_model_once():
    """Build the model **on the CPU**, once, outside any GPU window.

    On ZeroGPU the GPU is attached only for the duration of a `@spaces.GPU`
    call and the whole reservation is billed against the caller's daily quota,
    so reading a checkpoint off /data and pulling the DINOv3 encoder from the
    hub must not happen in there.  CUDA also cannot be initialised in the main
    process, hence `torch.device("cpu")` here and `.to("cuda")` inside the
    worker.
    """
    global MODEL_PATH, MODEL, MODEL_SPEC, MODEL_CFG
    from infer.predict import load_model

    with MODEL_LOCK:
        if MODEL is None:
            if not MODEL_PATH_ON_BUCKET.is_file():
                raise RuntimeError(
                    f"Model checkpoint not found at {MODEL_PATH_ON_BUCKET}. "
                    "Mount the Hugging Face bucket containing best.pt at /data."
                )
            hf_token = os.environ.get("HF_TOKEN", "").strip()
            if not hf_token:
                raise RuntimeError(
                    "HF_TOKEN is not configured. Add it as a Space secret for "
                    "the DINOv3 encoder download."
                )
            MODEL_PATH = str(MODEL_PATH_ON_BUCKET)
            MODEL, MODEL_SPEC, MODEL_CFG = load_model(
                MODEL_PATH, torch.device("cpu"), hf_token
            )
    return MODEL, MODEL_SPEC, MODEL_CFG


# ---------------------------------------------------------------------
# the viewer document
# ---------------------------------------------------------------------
def _build_viewer_shell() -> str:
    """`viewer/index.html` with its vendored scripts inlined, once, at import.

    The result carries `_PAYLOAD_SLOT` where the per-prediction JSON goes, so a
    request only substitutes a string instead of re-reading 660 KB of JavaScript.
    """
    document = (_VIEWER_ROOT / "index.html").read_text(encoding="utf-8")

    # the payload is read by app.js at DOMContentLoaded; keep it ahead of the
    # script that reads it so the document also reads sensibly to a human
    document = document.replace(
        '<script src="app.js"></script>',
        f'<script type="application/json" id="dw-embed">{_PAYLOAD_SLOT}</script>\n'
        '<script src="app.js"></script>',
    )

    def inline(match: "re.Match[str]") -> str:
        src = match.group(1)
        js = (_VIEWER_ROOT / src).read_text(encoding="utf-8")
        # a literal </script anywhere in the source would close this tag early
        return "<script>\n" + js.replace("</script", "<\\/script") + "\n</script>"

    document, n = re.subn(r'<script src="([^"]+)"></script>', inline, document)
    if n < 4:  # three.js, OrbitControls, PointerLockControls, app.js
        raise RuntimeError(f"viewer/index.html: inlined only {n} of 4 scripts")
    if _PAYLOAD_SLOT not in document:
        raise RuntimeError("viewer/index.html: no slot for the prediction payload")
    return document


_VIEWER_SHELL = _build_viewer_shell()


def _png_data_url(image: Image.Image) -> str:
    buffer = io.BytesIO()
    image.save(buffer, format="PNG", optimize=True)
    encoded = base64.b64encode(buffer.getvalue()).decode("ascii")
    return f"data:image/png;base64,{encoded}"


def _viewer(payload: dict) -> str:
    # `<` never appears outside a JSON string, so escaping it is enough to make
    # the blob inert inside the <script> element that carries it
    blob = json.dumps(payload, separators=(",", ":"), allow_nan=False)
    blob = blob.replace("<", "\\u003c")
    document = _VIEWER_SHELL.replace(_PAYLOAD_SLOT, blob)
    escaped = html.escape(document, quote=True)
    # No `sandbox`: an opaque-origin frame is one browser quirk away from
    # tainting the canvas that `decodeHeightRG16` reads its bytes out of, which
    # would take the whole viewer down.  The document is entirely
    # self-generated, so there is nothing here to sandbox from.  `allow` is the
    # permissions policy first-person mode needs to capture the pointer.
    return (
        '<iframe title="DSM viewer and 3D flythrough" '
        'allow="fullscreen; pointer-lock" allowfullscreen '
        'style="width:100%;height:min(760px,78vh);min-height:420px;border:0;'
        'border-radius:10px;background:#0d1013" '
        f'srcdoc="{escaped}"></iframe>'
    )


# A placeholder shown before the first run and whenever a run fails, so the
# 560px slot the CSS reserves never reads as a broken/blank panel.
_EMPTY_VIEWER_HTML = (
    '<div style="height:min(760px,78vh);min-height:420px;border:1px dashed #3a4048;'
    'border-radius:10px;background:#0d1013;color:#8a949e;display:flex;'
    'flex-direction:column;align-items:center;justify-content:center;gap:10px;'
    'text-align:center;padding:24px;font-family:system-ui,-apple-system,'
    '\'Segoe UI\',sans-serif">'
    '<div style="font-size:16px;color:#c7ced6">Your 3D terrain will appear here</div>'
    '<div style="font-size:13px;max-width:420px;line-height:1.5">'
    'Upload an aerial or satellite image, set its resolution (m/pixel), and '
    'click <b>Estimate height</b>.</div>'
    '</div>'
)


# ---------------------------------------------------------------------
# encodings and artefacts
# ---------------------------------------------------------------------
def _sanitise(height: np.ndarray) -> np.ndarray:
    """Replace any non-finite pixel with the scene minimum.

    Nothing downstream tolerates a NaN: it casts to garbage in the uint16
    encoding, puts a hole in the mesh, and makes `json.dumps` refuse the
    metadata.  A fully non-finite prediction is a real failure, so it is raised
    rather than patched — the caller turns this into a status message instead
    of a raw traceback.
    """
    finite = np.isfinite(height)
    if finite.all():
        return height
    if not finite.any():
        raise ValueError("The model returned no finite heights for this image.")
    out = height.copy()
    out[~finite] = float(height[finite].min())
    return out


def _height_bounds(height: np.ndarray) -> tuple[float, float]:
    """Full range, not a robust percentile.

    16 bits over a 0-60 m scene is a 0.9 mm step, so there is nothing to gain by
    clipping the tail — and clipping the tail is exactly how a tall building
    loses its roof between the model and the render.
    """
    low = float(np.nanmin(height))
    high = float(np.nanmax(height))
    return low, max(high, low + 1e-3)


def _encode_rg16(height: np.ndarray, low: float, high: float) -> Image.Image:
    """16-bit height split across R and G of an 8-bit PNG.

    A 2-D canvas read gives those bytes back exactly, unlike a true 16-bit PNG
    which the browser down-converts to 8 bits; `viewer/app.js` decodes this with
    `decodeHeightRG16`.
    """
    norm = np.clip((height - low) / (high - low), 0.0, 1.0)
    q16 = (norm * 65535 + 0.5).astype(np.uint16)
    return Image.fromarray(
        np.stack(
            [
                (q16 >> 8).astype(np.uint8),
                (q16 & 255).astype(np.uint8),
                np.zeros_like(q16, dtype=np.uint8),
            ],
            axis=-1,
        ),
        mode="RGB",
    )


def _fit(array: np.ndarray, max_side: int, order: str) -> tuple[np.ndarray, float]:
    """Downsample so the long side is `max_side`.  Returns (array, row scale)."""
    from dwdata.preprocess import resize

    side = max(array.shape[:2])
    if not max_side or side <= max_side:
        return array, 1.0
    scale = max_side / side
    out_hw = (max(1, int(array.shape[0] * scale)), max(1, int(array.shape[1] * scale)))
    return resize(array, out_hw, order), out_hw[0] / array.shape[0]


def _with_colorbar(colormap: np.ndarray, low: float, high: float) -> Image.Image:
    """Height colormap + a vertical turbo scale down its right edge.

    Colour cannot be converted back to metres by eye without this — the turbo
    ramp is the app's primary output and shipped with no legend otherwise.
    """
    from PIL import ImageDraw

    h, w = colormap.shape[:2]
    bar_w = max(28, w // 22)
    pad = 8
    canvas = Image.new("RGB", (w + bar_w + pad + 54, h), (13, 16, 19))
    canvas.paste(Image.fromarray(colormap), (0, 0))
    ramp = np.array(
        viz_figures.colorize(np.linspace(high, low, h).reshape(-1, 1), low, high),
        dtype=np.uint8,
    )
    canvas.paste(Image.fromarray(np.repeat(ramp, bar_w, axis=1)), (w + pad, 0))
    draw = ImageDraw.Draw(canvas)
    for frac, value in ((0.0, high), (0.5, (low + high) / 2), (1.0, low)):
        y = int(frac * (h - 1))
        y = min(max(y, 6), h - 8)
        draw.text((w + pad + bar_w + 4, y - 6), f"{value:.1f} m", fill=(232, 236, 241))
    return canvas


def _write_artifacts(
    out_dir: Path, rgb: np.ndarray, height: np.ndarray, meta: dict, gsd_m: float,
    seg: np.ndarray | None = None, objects: dict | None = None,
    uncertainty: np.ndarray | None = None,
) -> list[str]:
    """Everything a user can take away, written by the same code as the CLI."""
    out_dir.mkdir(parents=True, exist_ok=True)
    files: list[str] = []

    if seg is not None:
        # class ids as an 8-bit greyscale PNG on the height grid; the legend is
        # `meta.json` `classes`.  Lossless, ~50x smaller than an .npy.
        Image.fromarray(seg.astype(np.uint8), mode="L").save(out_dir / "seg.png")
        viz_classes.class_preview(rgb, seg, CLASS_NAMES).save(out_dir / "classes_preview.png")
        files += ["seg.png", "classes_preview.png"]
    if objects is not None:
        (out_dir / "objects.json").write_text(
            json.dumps(objects, separators=(",", ":")), encoding="utf-8"
        )
        viz_classes.objects_overlay(rgb, objects).save(out_dir / "objects_overlay.png")
        files += ["objects.json", "objects_overlay.png"]

    Image.fromarray(rgb).save(out_dir / "rgb.png")
    np.save(out_dir / "ndsm_m.npy", height.astype(np.float32))
    if uncertainty is not None and uncertainty.shape == height.shape:
        np.save(out_dir / "ndsm_std_m.npy", uncertainty.astype(np.float32))
        files.append("ndsm_std_m.npy")
    png16, affine = viz_mesh.encode_height_png16(height)
    (out_dir / "ndsm16.png").write_bytes(png16)
    files += ["ndsm_m.npy", "ndsm16.png", "rgb.png"]

    # the two pure-numpy figure primitives; neither needs matplotlib
    vmax = float(max(np.nanmax(height), 1.0))
    _with_colorbar(viz_figures.colorize(height, 0.0, vmax), 0.0, vmax).save(
        out_dir / "height_colormap.png"
    )
    Image.fromarray(viz_figures.hillshade(height, gsd_m), mode="L").save(
        out_dir / "hillshade.png"
    )
    files += ["height_colormap.png", "hillshade.png"]

    files += viz_mesh.write_mesh(out_dir, height, rgb, gsd_m, max_verts=400_000)

    payload = dict(meta)
    payload["ndsm16_encode"] = affine["encode"]
    payload["files"] = files + ["meta.json"]
    (out_dir / "meta.json").write_text(json.dumps(payload, indent=2), encoding="utf-8")
    files.append("meta.json")
    return files


def _bundle(out_dir: Path, files: list[str], name_stem: str) -> list[str]:
    """One zip plus the headline files — a mesh without its texture is junk, so
    the OBJ set is only offered inside the archive.

    `name_stem` comes from the uploaded file, not the temp directory, so the
    download reads as e.g. `rooftop_result.zip` instead of `dw_a1b2c3_result.zip`.
    """
    archive = shutil.make_archive(
        str(out_dir.parent / f"{name_stem}_result"), "zip", out_dir
    )
    # the frontend finds its inputs by name in this list, so a file it reads
    # (seg.png, objects.json) must be here and not only inside the zip
    keep = ["ndsm_m.npy", "meta.json", "terrain.glb", "height_colormap.png",
            "hillshade.png", "seg.png", "objects.json", "ndsm_std_m.npy"]
    return [archive] + [str(out_dir / f) for f in keep if (out_dir / f).is_file()]


def _safe_stem(image_path: str) -> str:
    stem = Path(image_path).stem or "scene"
    stem = re.sub(r"[^A-Za-z0-9_-]+", "_", stem).strip("_")
    return stem[:40] or "scene"


# ---------------------------------------------------------------------
# ZeroGPU bookkeeping
# ---------------------------------------------------------------------
# A `@spaces.GPU(duration=N)` call reserves N seconds against the caller's
# daily GPU allowance *up front*: if their remaining quota is below N the
# request is refused with "GPU time exhausted" no matter how many GPUs are
# idle.  The old code asked for 180 s around the whole handler — checkpoint
# load, hub download, meshing, PNG encoding and zipping included — which put
# the demo out of reach of anonymous visitors (~60 s) and burned a signed-in
# user's quota on work that never touches the GPU.
#
# So: only the forward pass is decorated, the durations are sized to it, and
# TTA (~4x the work) gets its own longer window instead of everyone paying for
# the worst case.
#
# Those two are floors.  A large scene is no longer shrunk to 1600 px — it is
# cut into 512 px windows at full resolution and batched through the GPU — so
# the window grows with the number of tiles (see `_gpu_duration`).
_GPU_DURATION = 45
_GPU_DURATION_TTA = 120
_GPU_DURATION_MAX = int(os.environ.get("DW_GPU_MAX_DURATION", "300"))
# Seconds of GPU per tile forward (8 of them per tile under TTA), padded well
# above what a batched ViT-L/512 pass costs on an H200 or an RTX Blackwell.
_GPU_SECONDS_PER_VIEW = 0.02

# Long side above which a scene is downsampled before inference.  This used to
# be 1600, which threw away most of a large GeoTIFF's resolution; the tiler has
# no size limit of its own, so the cap is now only a host-memory guard.  It
# matches the frontend's in-browser GeoTIFF decode limit.
_MAX_SIDE = int(os.environ.get("DW_MAX_SIDE", "8192"))
# Tiles per forward pass; 0 sizes it from free GPU memory.
_TILE_BATCH = int(os.environ.get("DW_TILE_BATCH", "0"))
_OVERLAP = 0.5


def _gpu_duration(rgb, gsd_m, tta: bool, input_scale: float = 1.0) -> int:
    """ZeroGPU window for this scene, sized to its tile count."""
    from infer.engine import count_tiles
    from models.tta import n_views

    floor = _GPU_DURATION_TTA if tta else _GPU_DURATION
    if MODEL_SPEC is None:
        return floor
    tiles = count_tiles(rgb.shape[0], rgb.shape[1], gsd_m, MODEL_SPEC, _OVERLAP)
    need = 20 + tiles * (n_views() if tta else 1) * _GPU_SECONDS_PER_VIEW * input_scale**2
    return int(min(_GPU_DURATION_MAX, max(floor, need)))


def _run_model(rgb, gsd_m, tta: bool, input_scale: float = 1.0):
    """The only GPU-bound step.  Runs inside a ZeroGPU reservation."""
    from infer.engine import predict_scene

    worker_start = time.perf_counter()
    model, spec, cfg = _load_model_once()
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = model.to(device).eval()
    amp_dtype = (
        (torch.bfloat16 if cfg.amp_dtype == "bf16" else torch.float16)
        if device.type == "cuda"
        else None
    )
    try:
        # the class map rides along on the same forward pass (under TTA it is
        # read off the identity view, so it costs no extra pass either)
        height, seg, uncertainty = predict_scene(
            model, rgb, gsd_m, spec, device,
            tta=bool(tta), tta_scales=(input_scale,), amp_dtype=amp_dtype, overlap=_OVERLAP,
            batch_tiles=_TILE_BATCH, want_seg=True, want_std=True,
        )
    finally:
        # hand the card back clean; the next reservation is a fresh attach
        model.to("cpu")
        if device.type == "cuda":
            torch.cuda.empty_cache()
        print("[space-timing] " + json.dumps({
            "d4": bool(tta), "scale": input_scale, "size_px": list(rgb.shape[:2]),
            "worker_seconds": time.perf_counter() - worker_start,
        }), flush=True)
    return height, seg.astype(np.uint8), uncertainty


# `duration` may be a callable taking the same arguments as the function
_run_model_gpu = spaces.GPU(duration=_gpu_duration)(_run_model)


@spaces.GPU(duration=180)
def _benchmark_gpu(rgb, gsd_m):
    from infer.benchmark import benchmark_model

    model, spec, cfg = _load_model_once()
    return benchmark_model(model, spec, cfg, rgb, gsd_m, _OVERLAP, _TILE_BATCH)


def benchmark_inference(image_path: str, gsd: float):
    from dwdata.preprocess import read_scene

    if not image_path:
        raise ValueError("Upload a representative scene for the benchmark.")
    _, spec, _ = _load_model_once()
    rgb, meta = read_scene(image_path, user_gsd_m=float(gsd or 0),
                           assumed_gsd_m=spec.canonical_gsd_m, max_side=1024)
    return _benchmark_gpu(rgb, meta.gsd_m)


# Outputs of `predict`, in order — kept as a tuple factory so every return
# path (success, validation failure, GPU failure) fills exactly these six
# slots and a failed run never leaves a stale result on screen.
def _result(gallery=None, height16=None, viewer_html=None, downloads=None,
            status="", warning=""):
    return gallery, height16, viewer_html or _EMPTY_VIEWER_HTML, downloads, status, warning


def predict(image_path: str | None, gsd: float, tta: bool = False,
            input_scale: float = 1.0):
    if not image_path:
        return _result(status="Upload an image first.")
    try:
        from dwdata.preprocess import read_scene

        # CPU work, deliberately outside the reservation
        _, spec, _ = _load_model_once()
        orig_w, orig_h = Image.open(image_path).size
        rgb, scene_meta = read_scene(
            image_path,
            user_gsd_m=float(gsd or 0),
            assumed_gsd_m=spec.canonical_gsd_m,
            max_side=_MAX_SIDE,
        )

        height, seg, uncertainty = _run_model_gpu(rgb, scene_meta.gsd_m, bool(tta), input_scale)

        height = _sanitise(height)
        # Objects are an extra, never a reason to lose the height map: a failure
        # here is reported in the status and the run carries on without them.
        objects, objects_error = None, None
        try:
            from infer.objects import extract_objects
            from infer.objects import summarise as summarise_objects

            objects = extract_objects(seg, height, scene_meta.gsd_m)
        except Exception as exc:  # noqa: BLE001
            objects_error = str(exc) or type(exc).__name__

        low, high = _height_bounds(height)
        height16 = _encode_rg16(height, low, high)

        meta = {
            "inference": {"d4": bool(tta), "input_scale": input_scale,
                          "checkpoint_epoch": MODEL_CFG.checkpoint_epoch},
            "product": "nDSM" if scene_meta.georeferenced else "rDSM",
            "units": "metres_above_ground",
            "height_min_m": low,
            "height_max_m": high,
            "height_mean_m": float(np.nanmean(height)),
            "height_median_m": float(np.nanmedian(height)),
            "uncertainty": {
                "definition": "standard deviation of the fused adaptive-bin height distribution",
                "std_median_m": float(np.nanmedian(uncertainty)),
                "std_mean_m": float(np.nanmean(uncertainty)),
            },
            "frac_below_1m": float((height < 1.0).mean()),
            "size_px": [int(height.shape[0]), int(height.shape[1])],
            "scene": {
                **scene_meta.summary(),
                # read_scene rescales gsd_m when it downsamples a huge scene but
                # not the affine, so publishing the affine here would place the
                # result wrongly.  Omitted rather than approximated.
                "transform": None,
                "transform_note": "omitted: the demo may resample the scene",
            },
            "preproc": spec.to_dict(),
            # `seg.png` legend: pixel value -> class name
            "classes": {str(i): name for i, name in enumerate(CLASS_NAMES)},
            "class_fractions": {
                name: float((seg == i).mean()) for i, name in enumerate(CLASS_NAMES)
            },
            "objects": (
                {"file": "objects.json", "version": objects["version"],
                 "counts": objects["counts"], "truncated": objects["truncated"]}
                if objects is not None else {"error": objects_error}
            ),
        }

        stem = _safe_stem(image_path)
        out_dir = Path(tempfile.mkdtemp(prefix="dw_"))
        files = _write_artifacts(
            out_dir, rgb, height, meta, scene_meta.gsd_m, seg=seg, objects=objects,
            uncertainty=uncertainty,
        )

        # --- the iframe payload, capped and re-stated in its own metres ------
        rgb_small, _ = _fit(rgb, _VIEWER_MAX_SIDE, "bilinear")
        height_small, scale = _fit(height, _VIEWER_MAX_SIDE, "bilinear")
        view_meta = dict(meta)
        view_meta["size_px"] = [int(height_small.shape[0]), int(height_small.shape[1])]
        view_meta["scene"] = {**meta["scene"], "gsd_m": scene_meta.gsd_m / scale}
        viewer_html = _viewer({
            "height": _png_data_url(_encode_rg16(height_small, low, high)),
            "rgb": _png_data_url(Image.fromarray(rgb_small)),
            "meta": view_meta,
        })

        info = [
            f"**Height** {low:.2f} – {high:.2f} m · mean "
            f"{meta['height_mean_m']:.2f} m · median {meta['height_median_m']:.2f} m",
            f"**GSD** {scene_meta.gsd_m:.3f} m/px (effective, {scene_meta.gsd_source}) · "
            f"**scene** {height.shape[1]} × {height.shape[0]} px",
            f"**{meta['frac_below_1m'] * 100:.1f} %** of pixels below 1 m "
            "(a nadir urban scene is usually 40–70 %)",
            (f"**Objects** {summarise_objects(objects)}" if objects is not None
             else f"**Objects** not extracted: {objects_error}"),
            f"Product `{meta['product']}` · {len(files)} artefacts written",
        ]
        if max(orig_w, orig_h) != max(scene_meta.width, scene_meta.height):
            info.append(
                f"Scene resampled {max(orig_w, orig_h)} → "
                f"{max(scene_meta.width, scene_meta.height)} px "
                f"(GSD adjusted to {scene_meta.gsd_m:.3f} m/px to match)."
            )

        warning = ""
        if meta["frac_below_1m"] < 0.15:
            warning = (
                "⚠️ Very little flat ground in this result — the model may be "
                "reading texture as terrain. Double-check the declared GSD."
            )

        gallery = [
            (str(out_dir / "height_colormap.png"), "Height (turbo, 0 – max, with scale)"),
            (str(out_dir / "hillshade.png"), "Hillshade (315° / 45°)"),
            (str(out_dir / "classes_preview.png"), "Object classes (seg head)"),
        ]
        if objects is not None:
            gallery.append((str(out_dir / "objects_overlay.png"),
                            "Detected trees, buildings and water"))
        return _result(
            gallery, height16, viewer_html, _bundle(out_dir, files, stem),
            "\n\n".join(info), warning,
        )
    except Exception as exc:
        return _result(status=_failure_text(exc))


def _is_quota_error(exc: Exception) -> bool:
    text = str(exc).lower()
    return "quota" in text or "gpu task aborted" in text or "gpu time" in text


def _failure_text(exc: Exception) -> str:
    if _is_quota_error(exc):
        # not a model failure: the caller's ZeroGPU allowance ran out
        return (
            "Your ZeroGPU quota for today is used up. Sign in to Hugging "
            "Face for a larger allowance, turn TTA off (it asks for at "
            f"least {_GPU_DURATION_TTA} s instead of {_GPU_DURATION} s, "
            "more for a large scene), or retry after the quota resets."
        )
    return f"Inference failed: {exc}"


def predict_scaled(image_path: str | None, gsd: float, mode: str):
    modes = {"baseline": (False, 1.0), "plain_1.5": (False, 1.5),
             "d4_native": (True, 1.0), "d4_1.5": (True, 1.5)}
    if mode not in modes:
        raise ValueError("Unknown inference mode.")
    tta_mode, scale = modes[mode]
    return predict(image_path, gsd, tta_mode, scale)


# ---------------------------------------------------------------------
# area mode: a map bbox -> repository tiles -> georeferenced heights
# ---------------------------------------------------------------------
# The GeoLibre plugin's API.  Instead of an upload it takes a WGS84 bbox,
# stitches the high-resolution tiles under it from the repository at
# `DW_TILES_DIR` (see `tiles/`), and — because it built the grid itself — can
# publish an affine, so every output lands on the map.  A partly covered bbox is
# processed over its covered part only; `coverage` says how much that was.
def _area_status(status: str, message: str, **extra) -> tuple[dict, None]:
    return {"ok": False, "status": status, "message": message, **extra}, None


def predict_area(bbox_json: str, tta: bool = False):
    from tiles import outputs as geo
    from tiles.catalog import TILES_DIR, parse_bbox, query
    from tiles.mosaic import AreaNotAvailable, AreaTooLarge, build_area

    try:
        bbox = parse_bbox(bbox_json)
    except ValueError as exc:
        return _area_status("invalid", str(exc))
    try:
        # CPU work, deliberately outside the reservation
        _, spec, _ = _load_model_once()
        tiles = query(bbox)
        if not tiles:
            raise AreaNotAvailable("Area not available")
        area = build_area(bbox, tiles, TILES_DIR, gsd_m=spec.canonical_gsd_m)
    except AreaNotAvailable as exc:
        return _area_status("not_available", str(exc), coverage=0.0)
    except AreaTooLarge as exc:
        return _area_status("too_large", str(exc))
    except Exception as exc:  # noqa: BLE001
        return _area_status("error", f"Reading the tiles failed: {exc}")

    try:
        height, seg, uncertainty = _run_model_gpu(area.rgb, area.gsd_m, bool(tta))
        height = _sanitise(height)
        # nothing is predicted where there was no imagery: ground class, no
        # objects, and NaN in every georeferenced output
        seg = np.where(area.valid, seg, CLASS_NAMES.index("ground")).astype(np.uint8)
        height = np.where(area.valid, height, 0.0).astype(np.float32)
        objects, objects_error = None, None
        try:
            from infer.objects import extract_objects

            objects = extract_objects(seg, height, area.gsd_m)
        except Exception as exc:  # noqa: BLE001
            objects_error = str(exc) or type(exc).__name__

        valid_h = height[area.valid]
        low = float(valid_h.min())
        high = max(float(valid_h.max()), low + 1e-3)
        H, W = height.shape
        corners = geo.corners_wgs84(area.transform, area.crs, W, H)
        result = {
            "ok": True,
            "status": "partial" if area.coverage < 0.98 else "ok",
            "message": "",
            "coverage": round(area.coverage, 4),
            "tiles": area.tile_ids,
            "bounds_wgs84": geo.bounds_wgs84(corners),
            "corners_wgs84": corners,
            "crs": area.crs.to_string(),
            "transform": list(area.transform)[:6],
            "gsd_m": area.gsd_m,
            "size_px": [H, W],
            "height_min_m": low,
            "height_max_m": high,
            "height_mean_m": float(valid_h.mean()),
            "counts": objects["counts"] if objects else {},
            "objects_error": objects_error,
        }
        meta = {
            "product": "nDSM",
            "units": "metres_above_ground",
            **{k: result[k] for k in ("height_min_m", "height_max_m", "height_mean_m",
                                      "coverage", "tiles", "crs", "transform",
                                      "gsd_m", "size_px", "corners_wgs84")},
            "uncertainty": (
                {"std_median_m": float(np.nanmedian(uncertainty[area.valid]))}
                if uncertainty is not None and uncertainty.shape == height.shape else None
            ),
            "classes": {str(i): name for i, name in enumerate(CLASS_NAMES)},
            "objects": (
                {"file": "objects.json", "counts": objects["counts"],
                 "truncated": objects["truncated"]}
                if objects is not None else {"error": objects_error}
            ),
        }

        out_dir = Path(tempfile.mkdtemp(prefix="dw_area_"))
        files = _write_artifacts(out_dir, area.rgb, height, meta, area.gsd_m, seg=seg,
                                 objects=objects, uncertainty=uncertainty)
        masked = np.where(area.valid, height, np.nan)
        geo.write_height_tif(out_dir / "ndsm_m.tif", masked, area.transform, area.crs)
        geo.write_overlay_png(out_dir / "height_overlay.png", masked, max(high, 1.0))
        geo.write_objects_geojson(out_dir / "objects.geojson", objects, area.transform, area.crs)
        # the GeoLibre plugin's 3D surface: heights + texture on the same grid
        geo.write_heights_png(out_dir / "ndsm_rgb.png", masked)
        geo.write_imagery_jpg(out_dir / "imagery.jpg", area.rgb)
        files += ["ndsm_m.tif", "height_overlay.png", "objects.geojson", "ndsm_rgb.png",
                  "imagery.jpg"]

        archive = shutil.make_archive(str(out_dir.parent / f"{out_dir.name}_area"), "zip", out_dir)
        keep = ["height_overlay.png", "objects.geojson", "ndsm_m.tif", "meta.json", "terrain.glb",
                "ndsm_rgb.png", "imagery.jpg"]
        downloads = [archive] + [str(out_dir / f) for f in keep if (out_dir / f).is_file()]
        return result, downloads
    except Exception as exc:  # noqa: BLE001
        status = "quota" if _is_quota_error(exc) else "error"
        return _area_status(status, _failure_text(exc))


def catalog():
    """The tile repository's footprints, for the plugin's coverage layer."""
    from tiles.catalog import load_catalog

    return load_catalog()


_ABOUT = """# Single View Height Estimator
Upload an aerial or satellite image — get a per-pixel height map and a 3D view.
A first run typically takes ~15-45 s (~1-2 min with **Higher quality**); nothing
touches a GPU until you click **Estimate height**.
"""

# the viewer's controls, kept out of the way until someone has a result to read
_VIEWER_HELP = """
Drag to orbit, scroll to zoom. The toolbar switches between the orbit,
first-person and drone cameras, drapes the optical image / height / hillshade
over the surface, exaggerates the vertical scale and moves the sun. Click the
surface to probe a height or slope; drag across it for an elevation profile.
"""

_DOWNLOADS_LEGEND = (
    "`.npy` = raw heights in metres (float32) · `.png` = colourised previews · "
    "`.glb` = textured 3D mesh · `meta.json` = encoding + scene metadata · "
    "`seg.png` = per-pixel class ids (legend in `meta.json`) · "
    "`objects.json` = trees, buildings and water as 3D objects · "
    "the `.zip` bundles everything, named after your upload."
)

# The HTML viewer has no intrinsic height, so it collapses to a strip before the
# first run.  One rule, so the empty state still looks like a stage. The value
# tracks the iframe's own `min-height` (`_viewer`/`_EMPTY_VIEWER_HTML` above).
_CSS = "#viewer, #viewer > div { min-height: 420px; }"

# The viewer's palette is hard-coded dark; keep the host shell in step so a
# result does not read as a black slab dropped into a white page.
_FORCE_DARK_JS = "() => { document.body.classList.add('dark'); }"

with gr.Blocks(title="Single View Height Estimator", css=_CSS, js=_FORCE_DARK_JS) as demo:
    gr.Markdown(_ABOUT)
    with gr.Row():
        with gr.Column(scale=1):
            image = gr.Image(type="filepath", label="Image", sources=["upload"])
            gsd = gr.Number(
                value=0.5, minimum=0,
                label="Resolution (m/pixel)",
                info="How many metres one pixel covers. Wrong value = wrong "
                "heights. Leave at 0 to let the model guess from its training "
                "resolution — the effective value used is always echoed below.",
            )
            tta = gr.Checkbox(
                value=False, label="Higher quality (slower)",
                info=f"8 views per tile, all run in one GPU batch; reserves at "
                f"least {_GPU_DURATION_TTA} s of GPU quota instead of "
                f"{_GPU_DURATION} s. The usual cause of a quota error — leave "
                "off unless you need it.",
            )
            with gr.Row():
                run = gr.Button("Estimate height", variant="primary")
                clear = gr.Button("Reset")
            status = gr.Markdown()
            warning = gr.Markdown(visible=False, elem_id="lowground-warning")
            downloads = gr.File(label="Downloads", file_count="multiple")
            gr.Markdown(_DOWNLOADS_LEGEND)
            gr.Examples(
                examples=[[str(_VIEWER_ROOT / "sample" / "rgb.png"), 0.5, False]],
                inputs=[image, gsd, tta],
                label="No imagery to hand? Try a sample scene",
            )
        with gr.Column(scale=3, min_width=640):
            gr.Markdown("### 3D view")
            viewer = gr.HTML(value=_EMPTY_VIEWER_HTML, elem_id="viewer")
            with gr.Accordion("More outputs", open=False):
                gallery = gr.Gallery(
                    label="Height & hillshade", columns=2, height=320,
                    object_fit="contain",
                )
                height_output = gr.Image(label="Height (16-bit)", type="pil")
                gr.Markdown(
                    "Height (16-bit) decodes as "
                    "`min + (R*256 + G) / 65535 * (max - min)`, with `min` and "
                    "`max` from `meta.json`."
                )
    with gr.Accordion("How to read the result", open=False):
        gr.Markdown(_VIEWER_HELP)

    def _mark_warning(text: str):
        if not text:
            return gr.update(value="", visible=False)
        html_text = (
            '<div style="background:#3a1f0a;border:1px solid #7a4a1e;'
            'border-left:4px solid #e5793a;border-radius:6px;padding:10px 14px;'
            f'color:#f0c9a0;font-size:13px">{text}</div>'
        )
        return gr.update(value=html_text, visible=True)

    run.click(
        lambda: (
            gr.update(interactive=False), [], None, _EMPTY_VIEWER_HTML, None,
            "Running…", gr.update(value="", visible=False),
        ),
        outputs=[run, gallery, height_output, viewer, downloads, status, warning],
    ).then(
        predict,
        inputs=[image, gsd, tta],
        outputs=[gallery, height_output, viewer, downloads, status, warning],
        show_progress="full",
        concurrency_limit=1,
        concurrency_id="model",
    ).then(
        _mark_warning, inputs=warning, outputs=warning,
    ).then(
        lambda: gr.update(interactive=True), outputs=run,
    )

    clear.click(
        lambda: (
            None, 0.5, False, [], None, _EMPTY_VIEWER_HTML, None, "",
            gr.update(value="", visible=False),
        ),
        outputs=[image, gsd, tta, gallery, height_output, viewer, downloads,
                  status, warning],
    )

    # Private Space API for repeatable measurements; not part of the UI flow.
    benchmark_button = gr.Button(visible=False)
    benchmark_output = gr.JSON(visible=False)
    benchmark_button.click(
        benchmark_inference, inputs=[image, gsd], outputs=benchmark_output,
        api_name="benchmark_inference", concurrency_limit=1, concurrency_id="model",
    )
    scaled_button = gr.Button(visible=False)
    scaled_mode = gr.Dropdown(["baseline", "plain_1.5", "d4_native", "d4_1.5"],
                              value="baseline", visible=False)
    scaled_button.click(
        predict_scaled, inputs=[image, gsd, scaled_mode],
        outputs=[gallery, height_output, viewer, downloads, status, warning],
        api_name="predict_scaled", concurrency_limit=1, concurrency_id="model",
    )

    # Area mode for the GeoLibre plugin: bbox in, georeferenced results out.
    area_button = gr.Button(visible=False)
    area_bbox = gr.Textbox(visible=False)
    area_tta = gr.Checkbox(value=False, visible=False)
    area_result = gr.JSON(visible=False)
    area_files = gr.File(visible=False, file_count="multiple")
    area_button.click(
        predict_area, inputs=[area_bbox, area_tta], outputs=[area_result, area_files],
        api_name="predict_area", concurrency_limit=1, concurrency_id="model",
    )
    catalog_button = gr.Button(visible=False)
    catalog_output = gr.JSON(visible=False)
    catalog_button.click(
        catalog, outputs=catalog_output, api_name="catalog",
        concurrency_limit=8, concurrency_id="catalog",
    )


# Build the model on the CPU while the Space boots, so the first visitor's GPU
# reservation pays for the forward pass only and not for the checkpoint read
# and the encoder download.  A failure here (no bucket, no HF_TOKEN) is
# reported per-request as before, so it must not take the Space down.
try:
    _load_model_once()
except Exception as _warm_exc:  # pragma: no cover - startup convenience only
    print(f"[app] model warm-up deferred: {_warm_exc}")


if __name__ == "__main__":
    demo.launch()