"""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"""