Add a median-instead-of-mean toggle with a difference image

app.py

@@ -15,7 +15,7 @@ import torch
 from huggingface_hub import hf_hub_download
 from PIL import Image, ImageOps
 
-from ensemble import tta_average, unit_normals
+from ensemble import mean_of, median_of, tta_passes, unit_normals
 from framing import crop_to_box, display_size, letterbox, resize_maps
 
 try:                                   # present on ZeroGPU Spaces only
@@ -83,7 +83,7 @@ def to_img(x01: torch.Tensor) -> Image.Image:
     return Image.fromarray(arr.astype("uint8"))
 
 
-def predict(image, enhanced, progress=gr.Progress()):
+def predict(image, enhanced, median, progress=gr.Progress()):
     if image is None:
         raise gr.Error("Upload a photo first.")
     photo = ImageOps.exif_transpose(image).convert("RGB")
@@ -91,29 +91,39 @@ def predict(image, enhanced, progress=gr.Progress()):
     # mirrored from the photo's own edge) and cut the padding off the maps again.
     square, box = letterbox(photo, SIZE)
     x = to_tensor(square)
+    ow, oh = display_size(*photo.size)
+    use_median = bool(enhanced and median)
+    diff = None
     t0 = time.time()
     with torch.no_grad():
         if enhanced:
-            maps = tta_average(
+            passes = tta_passes(
                 refine, x, flip=True, rot_deg=ROT_DEG,
                 progress=lambda i, n: progress((i, n), desc="forward pass"))
+            mean_maps = mean_of(passes)
+            maps = median_of(passes) if use_median else mean_maps
         else:
             maps = refine(x)
-        maps = crop_to_box(maps.float(), box)
-        ow, oh = display_size(*photo.size)
-        maps = resize_maps(maps, ow, oh)     # photo's own shape, unit normals
+        maps = resize_maps(crop_to_box(unit_normals(maps.float()), box), ow, oh)
+        if use_median:                     # where the two reductions disagree
+            ref = resize_maps(crop_to_box(unit_normals(mean_maps.float()), box), ow, oh)
+            d = ((maps - ref).abs().mean(dim=1, keepdim=True) * 8).clamp(0, 1)
+            diff = to_img(d)
     dt = time.time() - t0
     albedo = to_img((maps[:, 0:3] + 1) / 2)
     depth = to_img((maps[:, 3:4] + 1) / 2)
     normal = to_img((maps[:, 4:7] + 1) / 2)
-    passes = 6 if enhanced else 1
-    note = (f"{passes} forward pass{'es' if passes > 1 else ''} in {dt:.1f} s "
+    passes_n = 6 if enhanced else 1
+    how = ("median" if use_median else "mean") if enhanced else "single pass"
+    note = (f"{passes_n} forward pass{'es' if passes_n > 1 else ''} ({how}) in {dt:.1f} s "
             f"on CPU; the whole {photo.size[0]}x{photo.size[1]} photo is kept "
             f"(model input {SIZE}x{SIZE}, long side scaled to {SIZE}).")
-    return photo.resize((ow, oh), Image.BICUBIC), albedo, depth, normal, note
+    if median and not enhanced:
+        note += " Median needs Enhanced: there is nothing to take a median of in one pass."
+    return photo.resize((ow, oh), Image.BICUBIC), albedo, depth, normal, diff, note
 
 
-EXAMPLES = [[os.path.join("examples", f), False]
+EXAMPLES = [[os.path.join("examples", f), False, False]
             for f in sorted(os.listdir("examples"))
             if f.lower().endswith((".jpg", ".png"))]
 
@@ -126,7 +136,10 @@ with gr.Blocks(title="ShadeNet-5") as demo:
         "x right, y up, z toward the viewer.\n\n"
         "**Fast** is one forward pass. **Enhanced** averages six passes "
         "(mirror and rotations of +-4 degrees) and gives better normals and "
-        "depth, but takes about six times as long. "
+        "depth, but takes about six times as long. The model card's recipe "
+        "takes the per-pixel **mean** of the six; tick **Median** to take the "
+        "median instead (it was marginally worse in the measurements on the "
+        "card) and see the difference image. "
         f"[Model card and files](https://huggingface.co/{MODEL_REPO}). "
         "The example photos are Flickr8k photos shown for illustration only; "
         "see the model card for their licence status.")
@@ -134,6 +147,8 @@ with gr.Blocks(title="ShadeNet-5") as demo:
         with gr.Column():
             inp = gr.Image(type="pil", label="Photo", sources=["upload", "clipboard"])
             enh = gr.Checkbox(False, label="Enhanced (6 passes, slower)")
+            med = gr.Checkbox(False, label="Median instead of mean (needs Enhanced)",
+                              interactive=False)
             btn = gr.Button("Decompose", variant="primary")
             msg = gr.Markdown()
         with gr.Column():
@@ -143,9 +158,11 @@ with gr.Blocks(title="ShadeNet-5") as demo:
             with gr.Row():
                 o_dep = gr.Image(label="Depth (lighter = farther)")
                 o_nor = gr.Image(label="Normal")
-    btn.click(predict, [inp, enh], [o_in, o_alb, o_dep, o_nor, msg])
-    gr.Examples(EXAMPLES, [inp, enh], [o_in, o_alb, o_dep, o_nor, msg],
-                fn=predict, cache_examples=False)
+            o_diff = gr.Image(label="Median minus mean (x8, all maps; black = identical)")
+    enh.change(lambda e: gr.update(interactive=bool(e), value=False), enh, med)
+    outs = [o_in, o_alb, o_dep, o_nor, o_diff, msg]
+    btn.click(predict, [inp, enh, med], outs)
+    gr.Examples(EXAMPLES, [inp, enh, med], outs, fn=predict, cache_examples=False)
 
 if __name__ == "__main__":
     demo.queue(max_size=8).launch()

ensemble.py

@@ -22,17 +22,18 @@ def unit_normals(maps, eps=1e-6):
 
 
 @torch.no_grad()
-def tta_average(refine_fn, x, flip=False, rot_deg=0.0, progress=None):
-    """Per-pixel mean over the mirror (optional) x rotations (0, -rot, +rot).
+def tta_passes(refine_fn, x, flip=False, rot_deg=0.0, progress=None):
+    """The passes of the mirror (optional) x rotations (0, -rot, +rot) recipe,
+    each brought back to the photo's frame, as a list of (maps, valid).
 
     Each pass: mirror -> rotate the photo (reflected border) -> predict ->
-    rotate the maps back, reorienting the normals -> un-mirror. A pass only
-    counts at the pixels its rotated canvas covers.
+    rotate the maps back, reorienting the normals -> un-mirror. `valid` marks
+    the pixels the pass covers (a rotated canvas loses the corners).
     """
     angles = (0.0, -float(rot_deg), float(rot_deg)) if rot_deg else (0.0,)
     flips = (False, True) if flip else (False,)
     passes = list(itertools.product(flips, angles))
-    num = den = None
+    out = []
     for i, (f, a) in enumerate(passes):
         if progress is not None:
             progress(i, len(passes))
@@ -48,6 +49,30 @@ def tta_average(refine_fn, x, flip=False, rot_deg=0.0, progress=None):
             p = p.flip(-1).clone()
             p[:, NORMAL_X] = -p[:, NORMAL_X]
             valid = valid.flip(-1)
+        out.append((p, valid))
+    return out
+
+
+def mean_of(passes):
+    """Per-pixel mean over the passes that cover the pixel (the card's recipe)."""
+    num = den = None
+    for p, valid in passes:
         num = p * valid if num is None else num + p * valid
         den = valid if den is None else den + valid
     return num / den.clamp_min(1.0)
+
+
+def median_of(passes):
+    """Per-pixel median over the passes that cover the pixel (the average of the
+    two middle values when an even number cover it). The unrotated pass covers
+    everything, so every pixel has at least one value."""
+    P = torch.stack([p for p, _ in passes])                   # [n, B, C, H, W]
+    V = torch.stack([v for _, v in passes]) > 0.5             # [n, B, 1, H, W]
+    P = torch.where(V.expand_as(P), P, torch.full_like(P, float("nan")))
+    return torch.nanquantile(P, 0.5, dim=0)
+
+
+@torch.no_grad()
+def tta_average(refine_fn, x, flip=False, rot_deg=0.0, progress=None):
+    """Per-pixel mean over the recipe's passes (see `tta_passes`)."""
+    return mean_of(tta_passes(refine_fn, x, flip, rot_deg, progress))

test_space.py

@@ -15,7 +15,8 @@ REPO = os.path.dirname(HERE)
 sys.path.insert(0, HERE)
 os.chdir(HERE)
 
-import app  # noqa: E402  (downloads config.json + the MIX ONNX from the Hub)
+import app  # noqa: E402
+import ensemble  # noqa: E402  (downloads config.json + the MIX ONNX from the Hub)
 
 sys.path.insert(0, REPO)
 from detail_diffuser.inference.core import tta_average as ref_tta, unit_normals as ref_unit  # noqa: E402
@@ -40,25 +41,49 @@ def stand_in(c):
     return torch.cat([c * 0.7, c.mean(1, keepdim=True) * 0.5, c.flip(-1) * 0.3], 1)
 
 
-mine = app.unit_normals(app.tta_average(stand_in, x, flip=True, rot_deg=app.ROT_DEG))
+mine = app.unit_normals(ensemble.tta_average(stand_in, x, flip=True, rot_deg=app.ROT_DEG))
 ref = ref_unit(ref_tta(lambda c, nz=None: stand_in(c), x, flip=True, rot_deg=app.ROT_DEG))
 d = float((mine - ref).abs().max())
 print("averaging code, max |space - repo| =", d)
 assert d == 0.0, d
 
-e1 = app.unit_normals(app.tta_average(app.refine, x, flip=True, rot_deg=app.ROT_DEG))
+e1 = app.unit_normals(ensemble.tta_average(app.refine, x, flip=True, rot_deg=app.ROT_DEG))
 assert e1.shape == (1, 7, app.SIZE, app.SIZE) and torch.isfinite(e1).all()
 n = e1[:, 4:7].norm(dim=1)
 assert float((n - 1).abs().max()) < 1e-4, "normals are not unit length"
 print("real model, 6 passes ok, mean |maps| =", float(e1.abs().mean()))
 
-out = app.predict(img, True)
-assert len(out) == 5 and out[0].size == out[1].size == img.size, (out[0].size, img.size)  # whole photo, own shape
-print("predict ok:", out[4])
+# median_of against NumPy's nanmedian on real stand-in passes (corners are covered by fewer passes)
+import numpy as np  # noqa: E402
+
+passes = app.tta_passes(lambda c: stand_in(c) + 0.05 * torch.randn(1, 7, app.SIZE, app.SIZE),
+                        x, flip=True, rot_deg=app.ROT_DEG)
+assert len(passes) == 6 and any(float(v.min()) < 0.5 for _, v in passes), "rotated passes should have uncovered corners"
+P = np.stack([p.numpy() for p, _ in passes]).copy()
+V = np.stack([v.numpy() for _, v in passes]) > 0.5
+P[~np.broadcast_to(V, P.shape)] = np.nan
+want = np.nanmedian(P, axis=0)
+got = app.median_of(passes).numpy()
+dm = float(np.abs(got - want).max())
+print("median_of vs numpy nanmedian, max diff =", dm)
+assert dm < 1e-5, dm
+assert float((app.median_of(passes) - app.mean_of(passes)).abs().mean()) > 1e-4, "median should differ from mean on noisy passes"
+one = [passes[0]]
+assert torch.equal(app.median_of(one), app.mean_of(one)), "one pass: median == mean"
+
+out = app.predict(img, True, False)
+assert len(out) == 6 and out[0].size == out[1].size == img.size, (out[0].size, img.size)  # whole photo, own shape
+assert out[4] is None and "(mean)" in out[5]
+print("predict ok:", out[5])
+out = app.predict(img, True, True)
+assert out[4] is not None and out[4].size == img.size and "(median)" in out[5], out[5]
+print("median ok:", out[5])
+out = app.predict(img, False, True)                 # median without Enhanced: single pass, note says so
+assert out[4] is None and "single pass" in out[5] and "needs Enhanced" in out[5], out[5]
 # odd shapes: a panorama (padding longer than the photo), tall, square, tiny, big
 for size in [(2000, 300), (300, 2000), (448, 448), (64, 48), (1500, 1100)]:
     im = Image.effect_noise(size, 40).convert("RGB")
-    out = app.predict(im, False)
+    out = app.predict(im, False, False)
     ow, oh = app.display_size(*size)
     assert out[0].size == out[1].size == out[2].size == out[3].size == (ow, oh), (size, out[0].size, (ow, oh))
 print("odd shapes ok")