Expose adaptive-bin uncertainty outputs

Derive per-pixel height uncertainty from the adaptive-bin distribution, propagate it through tiled inference and TTA, and publish ndsm_std_m.npy for the DepthWizard frontend.

app.py

@@ -267,6 +267,7 @@ def _with_colorbar(colormap: np.ndarray, low: float, high: float) -> Image.Image
 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)
@@ -287,6 +288,9 @@ def _write_artifacts(
 
     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"]
@@ -324,7 +328,7 @@ def _bundle(out_dir: Path, files: list[str], name_stem: str) -> list[str]:
     # 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"]
+            "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()]
 
 
@@ -397,17 +401,17 @@ def _run_model(rgb, gsd_m, tta: bool):
     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 = predict_scene(
+        height, seg, uncertainty = predict_scene(
             model, rgb, gsd_m, spec, device,
             tta=bool(tta), amp_dtype=amp_dtype, overlap=_OVERLAP,
-            batch_tiles=_TILE_BATCH, want_seg=True,
+            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()
-    return height, seg.astype(np.uint8)
+    return height, seg.astype(np.uint8), uncertainty
 
 
 # `duration` may be a callable taking the same arguments as the function
@@ -438,7 +442,7 @@ def predict(image_path: str | None, gsd: float, tta: bool = False):
             max_side=_MAX_SIDE,
         )
 
-        height, seg = _run_model_gpu(rgb, scene_meta.gsd_m, bool(tta))
+        height, seg, uncertainty = _run_model_gpu(rgb, scene_meta.gsd_m, bool(tta))
 
         height = _sanitise(height)
         # Objects are an extra, never a reason to lose the height map: a failure
@@ -462,6 +466,11 @@ def predict(image_path: str | None, gsd: float, tta: bool = False):
             "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": {
@@ -488,7 +497,8 @@ def predict(image_path: str | None, gsd: float, tta: bool = False):
         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
+            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 ------

infer/engine.py

@@ -57,12 +57,12 @@ class _Forward:
     start at the size that worked instead of rediscovering the limit each time.
     """
 
-    def __init__(self, model, device, amp_dtype, max_batch: int, want_seg: bool):
+    def __init__(self, model, device, amp_dtype, max_batch: int, want_seg: bool, want_std: bool):
         self.model, self.device, self.amp_dtype = model, device, amp_dtype
-        self.max_batch, self.want_seg = max(1, int(max_batch)), want_seg
+        self.max_batch, self.want_seg, self.want_std = max(1, int(max_batch)), want_seg, want_std
 
     def __call__(self, x):
-        preds, segs, i = [], [], 0
+        preds, stds, segs, i = [], [], [], 0
         while i < x.shape[0]:
             n = self.max_batch
             try:
@@ -77,12 +77,15 @@ class _Forward:
                 print(f"[engine] OOM at {n} tiles/forward; retrying at {self.max_batch}")
                 continue
             preds.append(out["fused"].float())
+            if self.want_std:
+                stds.append(out["b_std"].float())
             if self.want_seg:
                 # the ignore channel is a training target, never a prediction
                 segs.append(out["seg"][:, :SEG_IGNORE_INDEX].float().argmax(1))
             del out
             i += n
-        return torch.cat(preds), (torch.cat(segs) if self.want_seg else None)
+        return (torch.cat(preds), torch.cat(stds) if self.want_std else None,
+                torch.cat(segs) if self.want_seg else None)
 
 
 @torch.no_grad()
@@ -90,6 +93,7 @@ def predict_canonical(
     model, rgb_canon_u8: np.ndarray, spec: PreprocSpec, device,
     *, tta: bool = False, tta_scales=(1.0,), amp_dtype=None,
     overlap: float = 0.25, batch_tiles: int = 0, want_seg: bool = False,
+    want_std: bool = False,
     progress=None,
 ):
     """Sliding-window prediction over an image already at the canonical GSD.
@@ -135,6 +139,7 @@ def predict_canonical(
     std = torch.tensor(spec.std, dtype=torch.float32, device=device).view(1, 3, 1, 1)
     win = torch.from_numpy(hann2d(S)).to(canvas_dev)
     acc = torch.zeros((Hp, Wp), dtype=torch.float32, device=canvas_dev)
+    acc_m2 = torch.zeros_like(acc) if want_std else None
     wsum = torch.zeros((Hp, Wp), dtype=torch.float32, device=canvas_dev)
     # Labels cannot be averaged, so each pixel keeps the label from the tile it
     # sits most centrally in (largest Hann weight).  Last-write-wins put a
@@ -148,7 +153,7 @@ def predict_canonical(
     fwd = _Forward(
         model, device, amp_dtype,
         batch_tiles * views if batch_tiles > 0 else auto_forward_batch(device, S),
-        want_seg,
+        want_seg, want_std,
     )
     step = max(1, fwd.max_batch // views)  # tiles per step; TTA multiplies them
 
@@ -177,15 +182,19 @@ def predict_canonical(
                 tiles = cut(chunk)
             t = (tiles.permute(0, 3, 1, 2).float().div_(255.0) - mean) / std
             if tta:
-                pred, seg = tta_predict_parallel(fwd, t, tuple(tta_scales), patch=spec.patch)
+                pred, pred_std, seg = tta_predict_parallel(fwd, t, tuple(tta_scales), patch=spec.patch)
             else:
-                pred, seg = fwd(t)
+                pred, pred_std, seg = fwd(t)
             pred = pred[:, 0].to(canvas_dev)
+            if pred_std is not None:
+                pred_std = pred_std[:, 0].to(canvas_dev)
             if seg is not None:
                 seg = seg.to(device=canvas_dev, dtype=torch.int16)
 
             for i, (y, x) in enumerate(chunk):
                 acc[y:y + S, x:x + S] += pred[i] * win
+                if acc_m2 is not None:
+                    acc_m2[y:y + S, x:x + S] += (pred_std[i].square() + pred[i].square()) * win
                 wsum[y:y + S, x:x + S] += win
                 if seg is not None:
                     best = seg_best[y:y + S, x:x + S]
@@ -201,8 +210,12 @@ def predict_canonical(
             pool.shutdown(wait=True)
 
     height = (acc / wsum.clamp_min(1e-6))[:H, :W].cpu().numpy()
+    std_out = None
+    if acc_m2 is not None:
+        variance = (acc_m2 / wsum.clamp_min(1e-6) - (acc / wsum.clamp_min(1e-6)).square()).clamp_min(0)
+        std_out = variance[:H, :W].sqrt().cpu().numpy().astype(np.float32)
     seg_out = seg_acc[:H, :W].cpu().numpy() if seg_acc is not None else None
-    return height.astype(np.float32), seg_out
+    return (height.astype(np.float32), seg_out, std_out) if want_std else (height.astype(np.float32), seg_out)
 
 
 def _autocast(device, amp_dtype):
@@ -226,7 +239,7 @@ def count_tiles(h: int, w: int, src_gsd_m: float, spec: PreprocSpec,
 def predict_scene(
     model, rgb_u8: np.ndarray, src_gsd_m: float, spec: PreprocSpec, device,
     *, tta: bool = False, tta_scales=(1.0,), amp_dtype=None, overlap: float = 0.25,
-    batch_tiles: int = 0, want_seg: bool = False, progress=None,
+    batch_tiles: int = 0, want_seg: bool = False, want_std: bool = False, progress=None,
 ):
     """RGB at any GSD -> nDSM in metres on the input grid (+ optional class map).
 
@@ -237,11 +250,15 @@ def predict_scene(
     canon_hw = gsd_to_shape(H, W, src_gsd_m, spec.canonical_gsd_m)
     rgb_c = resize(rgb_s, canon_hw, "bilinear")
 
-    height_c, seg_c = predict_canonical(
+    prediction = predict_canonical(
         model, rgb_c, spec, device, tta=tta, tta_scales=tta_scales,
         amp_dtype=amp_dtype, overlap=overlap, batch_tiles=batch_tiles,
-        want_seg=want_seg, progress=progress,
+        want_seg=want_seg, want_std=want_std, progress=progress,
     )
+    height_c, seg_c = prediction[:2]
     height = resize(height_c, (H, W), "bilinear").astype(np.float32)
     seg = resize(seg_c.astype(np.int32), (H, W), "nearest") if seg_c is not None else None
+    if want_std:
+        std = resize(prediction[2], (H, W), "bilinear").astype(np.float32)
+        return height, seg, std
     return height, seg

models/heads.py

@@ -98,11 +98,21 @@ class DepthWizardNetV3(nn.Module):
         b = self.head_b(x)
         fused, alpha = self.fusion(x, a, b["height"])
 
+        # Head B is a per-pixel distribution over adaptive height bins.  Its
+        # standard deviation is the predictive spread; scale by the fusion
+        # gate because only Head B contributes uncertainty to the fused height.
+        probs = F.softmax(b["logits"].float(), dim=1)
+        centres = b["centres"].float()[:, :, None, None]
+        b_height = b["height"].float()
+        variance = (probs * (centres - b_height) ** 2).sum(dim=1, keepdim=True)
+        b_std = variance.clamp_min(0).sqrt() * (1.0 - alpha.float())
+
         up = lambda t: F.interpolate(t, size=hw, mode="bilinear", align_corners=False)  # noqa: E731
         return {
             "a": up(a),
             "b": up(b["height"]),
             "fused": up(fused),
+            "b_std": up(b_std),
             "alpha": up(alpha),
             "seg": up(self.head_c(x)),
             # kept at half res — the bin loss is computed against a downsampled

models/tta.py

@@ -66,12 +66,12 @@ def tta_predict_parallel(forward, image, scales=(1.0,), patch: int = 16):
     views stack along the batch axis and `forward` sees one (8B, 3, S, S) batch
     per scale; it is free to split that into whatever fits in memory.
 
-    `forward(x) -> (pred (N, 1, h, w) float32, seg (N, h, w) long | None)`.
-    Returns the view-averaged height and the class map of the identity view, so
-    the class map no longer costs its own extra pass.
+    `forward(x) -> (pred, std, seg)`. Average first and second moments across
+    views so the returned spread includes both bin uncertainty and TTA disagreement.
     """
     b, _, h, w = image.shape
     acc = torch.zeros((b, 1, h, w), device=image.device, dtype=torch.float32)
+    acc_m2 = torch.zeros_like(acc)
     seg_id, n = None, 0
     for s in scales:
         if abs(s - 1.0) < 1e-6:
@@ -81,14 +81,23 @@ def tta_predict_parallel(forward, image, scales=(1.0,), patch: int = 16):
             nw = max(patch, int(round(w * s / patch)) * patch)
             img_s = F.interpolate(image, size=(nh, nw), mode="bilinear", align_corners=False)
         views = torch.cat([_apply(img_s, k, flip) for k, flip in _D4])
-        pred, seg = forward(views)
+        pred, pred_std, seg = forward(views)
         for i, (k, flip) in enumerate(_D4):
             p = _invert(pred[i * b:(i + 1) * b], k, flip)
+            p_slice = pred[i * b:(i + 1) * b]
+            m2_slice = p_slice.square()
+            if pred_std is not None:
+                m2_slice = m2_slice + pred_std[i * b:(i + 1) * b].square()
+            m2 = _invert(m2_slice, k, flip)
             if p.shape[-2:] != (h, w):
                 p = F.interpolate(p, size=(h, w), mode="bilinear", align_corners=False)
+                m2 = F.interpolate(m2, size=(h, w), mode="bilinear", align_corners=False)
             acc += p
+            acc_m2 += m2
             n += 1
         # _D4[0] is the identity; only the full-resolution pass supplies labels
         if seg is not None and seg_id is None and img_s is image:
             seg_id = seg[:b]
-    return acc / max(1, n), seg_id
+    mean = acc / max(1, n)
+    std = (acc_m2 / max(1, n) - mean.square()).clamp_min(0).sqrt()
+    return mean, std, seg_id