app.py

"""DBFT Pavement Strain Predictor — Gradio app for Hugging Face Spaces.

Predicts the two critical pavement strains directly from an FWD deflection
basin + layer thicknesses, using the best model from the paper (combined
loss lambda_f = 1.0, extended surrogate), with a local SHAP explanation
for every prediction.
"""

from pathlib import Path

try:
    import spaces  # ZeroGPU: must be imported before torch
except ImportError:  # local run / CPU Space without the spaces package
    spaces = None

import gradio as gr
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import shap
import torch

from fwd_fusion_transformer import DBFT, basin_indices

HERE = Path(__file__).resolve().parent
ASSETS = HERE / "assets"
TAG = "combined_lam1.0_ext"
FEATURES = ["D0", "D200", "D300", "D450", "D600", "D900", "D1200", "D1500",
            "D1800", "h_AC", "h_Base", "h_Subbase"]
NSAMPLES = 200

# ---------------- model ----------------
sc = np.load(ASSETS / f"dbft_{TAG}_scalers.npz")
model = DBFT()
model.load_state_dict(torch.load(ASSETS / f"dbft_{TAG}.pt", map_location="cpu"))
model.eval()


def f(X):
    """X: (n, 12) raw [D0..D1800 um, h_AC..h_Subbase mm] -> (n, 2) strains ue."""
    X = np.asarray(X, np.float32)
    D, H = X[:, :9], X[:, 9:12]
    I = basin_indices(D).astype(np.float32)
    t = lambda a, k: torch.tensor((a - sc[f"{k}_mean"]) / sc[f"{k}_std"],
                                  dtype=torch.float32)
    with torch.no_grad():
        p = model(t(D, "D"), t(I, "I"), t(H, "H")).numpy()
    return p * sc["Y_std"] + sc["Y_mean"]


bg = np.load(ASSETS / "bg_cache.npz")["bg"]
explainer = shap.KernelExplainer(f, bg)
BASE = np.asarray(explainer.expected_value, float)

RED, BLUE, GREEN, AMBER = "#d64545", "#3b7dd8", "#10a37f", "#d69e2e"


def severity(v, lo, hi):
    return (("low", GREEN) if v < lo else
            ("moderate", AMBER) if v < hi else ("high", RED))


def shap_figure(sv_ac, sv_sg, pred):
    """Two-panel horizontal bar chart of local SHAP values (ue)."""
    fig, axes = plt.subplots(1, 2, figsize=(11, 4.2), dpi=140)
    titles = [(r"$\varepsilon_t$ — AC tensile", sv_ac, BASE[0], pred[0]),
              (r"$\varepsilon_c$ — subgrade compressive", sv_sg, BASE[1],
               pred[1])]
    for ax, (title, sv, base, p) in zip(axes, titles):
        order = np.argsort(np.abs(sv))
        names = [FEATURES[i] for i in order]
        vals = sv[order]
        colors = [RED if v >= 0 else BLUE for v in vals]
        ax.barh(range(len(vals)), vals, color=colors, alpha=0.85)
        ax.axvline(0, color="#999", lw=1)
        ax.set_yticks(range(len(vals)))
        ax.set_yticklabels(names, fontsize=9)
        ax.set_xlabel("SHAP value (με)", fontsize=9)
        ax.set_title(f"{title}\nbaseline {base:.0f} με → prediction "
                     f"{p:.0f} με", fontsize=10)
        ax.grid(alpha=0.3, axis="x")
        for s in ["top", "right"]:
            ax.spines[s].set_visible(False)
    fig.suptitle("Local SHAP — red pushes strain up, blue pushes it down",
                 fontsize=10.5, y=1.02)
    fig.tight_layout()
    return fig


def _gpu(fn):
    """ZeroGPU hardware refuses to start without a @spaces.GPU function.
    Inference itself runs on CPU in <1 s (SHAP ~3 s), so the short duration
    just satisfies the check while keeping queue priority high."""
    return spaces.GPU(duration=30)(fn) if spaces is not None else fn


@_gpu
def predict(d0, d200, d300, d450, d600, d900, d1200, d1500, d1800,
            h_ac, h_base, h_subbase, explain):
    d = [d0, d200, d300, d450, d600, d900, d1200, d1500, d1800]
    h = [h_ac, h_base, h_subbase]
    if any(v is None for v in d + h):
        raise gr.Error("Please fill in all 12 inputs (use 0 only for "
                       "h_Subbase when there is no subbase).")
    d, h = np.array(d, np.float32), np.array(h, np.float32)
    if np.any(d <= 0):
        raise gr.Error("Deflections must be positive (μm at 707 kPa).")
    if h[0] <= 0 or np.any(h < 0):
        raise gr.Error("Thicknesses must be ≥ 0 mm with h_AC > 0.")
    if d[0] < d[-1]:
        raise gr.Error("D0 should exceed D1800 — check the basin order.")

    x = np.concatenate([d, h])[None, :]
    eps_ac, eps_sg = map(float, f(x)[0])

    s_ac, c_ac = severity(eps_ac, 150, 400)
    s_sg, c_sg = severity(eps_sg, 300, 600)
    html = f"""
    <div style="display:flex;gap:14px;flex-wrap:wrap;font-family:system-ui">
      <div style="flex:1;min-width:230px;border:1px solid #e3e3e6;
                  border-radius:12px;padding:14px;background:#fafbfc">
        <div style="font-size:13px;color:#6e6e80">ε<sub>t</sub> — AC tensile
          strain (fatigue criterion)</div>
        <div style="font-size:30px;font-weight:700">{eps_ac:.1f}
          <span style="font-size:15px;color:#6e6e80">με</span>
          <span style="font-size:14px;color:{c_ac}">· {s_ac}</span></div>
        <div style="font-size:12px;color:#6e6e80">bottom of the asphalt layer
        </div>
      </div>
      <div style="flex:1;min-width:230px;border:1px solid #e3e3e6;
                  border-radius:12px;padding:14px;background:#fafbfc">
        <div style="font-size:13px;color:#6e6e80">ε<sub>c</sub> — subgrade
          compressive strain (rutting criterion)</div>
        <div style="font-size:30px;font-weight:700">{eps_sg:.1f}
          <span style="font-size:15px;color:#6e6e80">με</span>
          <span style="font-size:14px;color:{c_sg}">· {s_sg}</span></div>
        <div style="font-size:12px;color:#6e6e80">top of the subgrade</div>
      </div>
    </div>"""

    fig = None
    if explain:
        sv = explainer.shap_values(x, nsamples=NSAMPLES, silent=True)
        sv = np.array(sv)
        if sv.ndim == 3 and sv.shape[-1] == 2:
            sv = np.moveaxis(sv, -1, 0)
        fig = shap_figure(sv[0, 0], sv[1, 0], (eps_ac, eps_sg))
    return html, fig


EXAMPLES = [
    # Thai DOH route 23 (thin structure, unbound base)
    [171.0, 154.5, 145.1, 89.6, 62.8, 47.4, 32.9, 22.6, 17.9,
     100, 150, 300, True],
    # factorial-like softer structure
    [1126.8, 940.1, 815.5, 663.5, 548.0, 390.7, 293.9, 231.7, 189.9,
     100, 200, 300, True],
    # stiff bound-base section
    [120.0, 100.0, 90.0, 75.0, 62.0, 45.0, 33.0, 25.0, 20.0,
     50, 200, 0, True],
]

with gr.Blocks(theme=gr.themes.Soft(primary_hue="emerald"),
               title="DBFT — Pavement Strain Predictor") as demo:
    gr.Markdown(
        "# 🛣️ DBFT — Pavement Strain Predictor\n"
        "Predicts the two critical pavement strains **directly** from a "
        "Falling Weight Deflectometer test — no backcalculation step. "
        "Model: Deflection-Basin Fusion Transformer (~150k parameters), "
        "combined loss λ = 1.0 on an extended layered-elastic surrogate + "
        "7,651 Thai DOH field points. Held-out-route field accuracy: "
        "**R² = 0.96 (AC) / 0.88 (subgrade)**."
    )
    gr.Image(str(ASSETS / "fwd_diagram.png"), show_label=False,
             container=False, interactive=False, show_download_button=False,
             show_fullscreen_button=False)

    gr.Markdown("### Deflection basin — μm, normalized to 707 kPa")
    with gr.Row():
        d_in = [gr.Number(label=lab, precision=1) for lab in
                ["D0", "D200", "D300", "D450", "D600", "D900", "D1200",
                 "D1500", "D1800"]]
    gr.Markdown("### Layer thicknesses — mm (h_Subbase = 0 if no subbase)")
    with gr.Row():
        h_in = [gr.Number(label="h_AC (mm)", precision=0),
                gr.Number(label="h_Base (mm)", precision=0),
                gr.Number(label="h_Subbase (mm)", precision=0)]
        explain_in = gr.Checkbox(value=True,
                                 label="Explain with local SHAP (~3 s)")

    btn = gr.Button("Predict", variant="primary")
    out_html = gr.HTML()
    out_plot = gr.Plot(label="Why? — local SHAP attribution")

    btn.click(predict, inputs=d_in + h_in + [explain_in],
              outputs=[out_html, out_plot])
    gr.Examples(examples=EXAMPLES, inputs=d_in + h_in + [explain_in],
                label="Examples (click a row, then Predict)")
    gr.Markdown(
        "<small>Research prototype — trained on Thai DOH FWD data and "
        "layered-elastic theory at 707 kPa / 150 mm plate. SHAP: "
        "KernelExplainer over the 12 measurable inputs; red bars push the "
        "prediction up, blue bars down.</small>"
    )

if __name__ == "__main__":
    # ssr_mode=False: the experimental SSR (Node.js) path is flaky on
    # HF Spaces — "Stopping Node.js server..." restart loops.
    demo.launch(ssr_mode=False)