Add 2D image bilateral example with astronaut figures and PSNR

Author: chrischoy

docs/user_guide/bilateral_permutohedral_filters.md

@@ -1,7 +1,7 @@
 # Bilateral and Permutohedral Filters
 
 **Created**: 2026-04-30 19:55:27 PST
-**Edited**: 2026-04-30 21:30:00 PST
+**Edited**: 2026-04-30 22:50:38 PST
 
 WarpConvNet ships three families of edge-preserving filters for point clouds
 and high-dimensional feature volumes. They differ in the underlying spatial
@@ -159,6 +159,89 @@ Sinkhorn-style bistochastization runs by default; disable via the
 underlying `bilateral_solver(grid, ..., bistochastize=False)` if extreme
 confidence values cause non-finite scaling factors.
 
+## Worked example: 2D image denoising
+
+Image denoising is the canonical bilateral demo: each pixel is a point with
+a 2D position and a 3D color, the guide is `concat(xy/sigma_xy, rgb/sigma_rgb)`, and the value being filtered is the noisy color itself.
+
+The example below applies all three filter families to the NASA `astronaut`
+test image (public domain, shipped with `scikit-image`) corrupted with
+Gaussian noise of variance 0.01. End-to-end times are on an RTX 6000 Ada at
+512×512 = 262k "points" with $\sigma_{xy} = 4$, $\sigma_{rgb} = 0.1$.
+
+### Input
+
+<table>
+<tr>
+<td align="center"><b>Original</b></td>
+<td align="center"><b>Noisy (Gaussian, var=0.01) — 20.70 dB</b></td>
+</tr>
+<tr>
+<td><img src="img/astronaut_original.jpg" alt="Original astronaut" width="100%"></td>
+<td><img src="img/astronaut_noisy.jpg" alt="Noisy astronaut" width="100%"></td>
+</tr>
+</table>
+
+### Output
+
+<table>
+<tr>
+<td align="center"><b>KNN (k=24) — 23.67 dB / ~3.3 s</b></td>
+<td align="center"><b>Grid — 23.94 dB / ~63 ms</b></td>
+<td align="center"><b>Permutohedral — 24.68 dB / ~11 ms</b></td>
+</tr>
+<tr>
+<td><img src="img/astronaut_knn.jpg" alt="KNN bilateral" width="100%"></td>
+<td><img src="img/astronaut_grid.jpg" alt="Grid bilateral" width="100%"></td>
+<td><img src="img/astronaut_permutohedral.jpg" alt="Permutohedral bilateral" width="100%"></td>
+</tr>
+</table>
+
+| Filter                         | Time   | PSNR (dB) | Notes                          |
+| ------------------------------ | ------ | --------- | ------------------------------ |
+| Noisy input (reference)        | —      | 20.70     | Gaussian noise, var = 0.01     |
+| `BilateralFilter` (KNN, k=24)  | ~3.3 s | 23.67     | Exact Gaussian, $O(N \cdot k)$ |
+| `BilateralFilterGrid`          | ~63 ms | 23.94     | $d=5$ sparse cube              |
+| `BilateralPermutohedralFilter` | ~11 ms | **24.68** | $d=5$ permutohedral lattice    |
+
+PSNR is computed against the clean original with `data_range=1.0`. The
+permutohedral lattice is both fastest *and* highest PSNR here — its
+conservative reconstruction (3-tap Gaussian on $(d{+}1)$ lattice axes)
+preserves edges slightly better than the $d$-cube grid at this bandwidth,
+and noticeably better than the limited-$K$ KNN filter where boundary
+neighbors get clipped.
+
+Reproduce with [`examples/bilateral_image_example.py`](https://github.com/NVIDIA/warpconvnet/blob/main/examples/bilateral_image_example.py):
+
+```bash
+python examples/bilateral_image_example.py \
+    --out-dir docs/user_guide/img \
+    --sigma-xy 4.0 --sigma-rgb 0.1 --noise-var 0.01
+```
+
+Sketch of the call site:
+
+```python
+from skimage import data, util
+import torch
+import warpconvnet.nn as wn
+
+img = util.img_as_float(data.astronaut())                    # (512, 512, 3)
+noisy = util.random_noise(img, mode="gaussian", var=0.01)
+h, w, _ = img.shape
+
+ys, xs = torch.meshgrid(torch.arange(h), torch.arange(w), indexing="ij")
+xy  = torch.stack([xs, ys], dim=-1).reshape(-1, 2).float().cuda()
+rgb = torch.from_numpy(noisy).reshape(-1, 3).float().cuda()
+
+filt = wn.BilateralPermutohedralFilter(sigma_xyz=4.0, sigma_feat=0.1)
+denoised = filt(xy, rgb, rgb).reshape(h, w, 3).clamp(0, 1)
+```
+
+The KNN filter is exact-Gaussian but $O(N \cdot k)$; the lattice variants
+trade ~10–20% reconstruction error for two orders of magnitude in speed and
+become the only viable option above ~$10^5$ points.
+
 ## Constraints
 
 - `PackedHashTable128` supports $D \le 7$ axes per key. Lattice

docs/user_guide/img/astronaut_grid.jpg

@@ -0,0 +1,156 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Bilateral filtering on a 2D image — denoising the NASA `astronaut` test image.
+
+Demonstrates the three bilateral families shipped in `warpconvnet.nn`:
+
+  - BilateralFilter            (KNN / radius)
+  - BilateralFilterGrid        (sparse d-cube lattice)
+  - BilateralPermutohedralFilter (permutohedral lattice)
+
+We treat each pixel as a point with 2D position (x, y) and 3D color (r, g, b).
+The bilateral guide is concat(xy/sigma_xy, rgb/sigma_rgb); the value being
+filtered is the noisy color. The astronaut image is in the public domain
+(NASA), shipped with scikit-image.
+
+Saves five PNGs to <out-dir>:
+    astronaut_original.png, astronaut_noisy.png, astronaut_knn.png,
+    astronaut_grid.png, astronaut_permutohedral.png
+
+Run:
+    python examples/bilateral_image_example.py --out-dir docs/user_guide/img
+"""
+
+import argparse
+import time
+
+import numpy as np
+import torch
+from skimage import data, util
+
+
+def _to_pixel_pointcloud(img: np.ndarray, device: torch.device):
+    """Flatten an (H, W, 3) image into (N, 2) xy + (N, 3) rgb tensors."""
+    h, w, _ = img.shape
+    ys, xs = np.meshgrid(np.arange(h), np.arange(w), indexing="ij")
+    xy = np.stack([xs, ys], axis=-1).reshape(-1, 2).astype(np.float32)
+    rgb = img.reshape(-1, 3).astype(np.float32)
+    return (
+        torch.from_numpy(xy).to(device),
+        torch.from_numpy(rgb).to(device),
+    )
+
+
+def _from_pixel_pointcloud(values: torch.Tensor, h: int, w: int) -> np.ndarray:
+    return values.detach().cpu().numpy().reshape(h, w, 3).clip(0, 1)
+
+
+def _save_image(path: str, arr: np.ndarray) -> None:
+    from PIL import Image
+
+    arr = (np.clip(arr, 0.0, 1.0) * 255.0).round().astype(np.uint8)
+    img = Image.fromarray(arr)
+    if path.lower().endswith((".jpg", ".jpeg")):
+        img.save(path, quality=92, optimize=True, progressive=True)
+    else:
+        img.save(path, optimize=True)
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--out-dir", default="docs/user_guide/img")
+    parser.add_argument("--noise-var", type=float, default=0.01)
+    parser.add_argument("--sigma-xy", type=float, default=4.0)
+    parser.add_argument("--sigma-rgb", type=float, default=0.1)
+    parser.add_argument("--knn-k", type=int, default=24)
+    args = parser.parse_args()
+    import os
+
+    os.makedirs(args.out_dir, exist_ok=True)
+
+    if not torch.cuda.is_available():
+        raise SystemExit("CUDA required for bilateral filters")
+    device = torch.device("cuda")
+
+    import warpconvnet.nn as wn
+
+    # ---- input image -------------------------------------------------------
+    img = util.img_as_float(data.astronaut())  # (512, 512, 3) in [0, 1]
+    noisy = util.random_noise(img, mode="gaussian", var=args.noise_var)
+    h, w, _ = img.shape
+
+    xy, rgb_clean = _to_pixel_pointcloud(img, device)
+    _, rgb_noisy = _to_pixel_pointcloud(noisy.astype(np.float32), device)
+
+    # ---- KNN bilateral -----------------------------------------------------
+    knn_filter = wn.BilateralFilter(
+        sigma_xyz=args.sigma_xy,
+        sigma_feat=args.sigma_rgb,
+        k=args.knn_k,
+        mode="knn",
+    )
+    torch.cuda.synchronize()
+    t0 = time.perf_counter()
+    out_knn = knn_filter(xy, rgb_noisy, rgb_noisy)
+    torch.cuda.synchronize()
+    t_knn = time.perf_counter() - t0
+
+    # ---- sparse d-cube grid -----------------------------------------------
+    grid_filter = wn.BilateralFilterGrid(
+        sigma_xyz=args.sigma_xy,
+        sigma_feat=args.sigma_rgb,
+    )
+    torch.cuda.synchronize()
+    t0 = time.perf_counter()
+    out_grid = grid_filter(xy, rgb_noisy, rgb_noisy)
+    torch.cuda.synchronize()
+    t_grid = time.perf_counter() - t0
+
+    # ---- permutohedral lattice --------------------------------------------
+    perm_filter = wn.BilateralPermutohedralFilter(
+        sigma_xyz=args.sigma_xy,
+        sigma_feat=args.sigma_rgb,
+    )
+    torch.cuda.synchronize()
+    t0 = time.perf_counter()
+    out_perm = perm_filter(xy, rgb_noisy, rgb_noisy)
+    torch.cuda.synchronize()
+    t_perm = time.perf_counter() - t0
+
+    # ---- save individual PNGs ---------------------------------------------
+    knn_img = _from_pixel_pointcloud(out_knn, h, w)
+    grid_img = _from_pixel_pointcloud(out_grid, h, w)
+    perm_img = _from_pixel_pointcloud(out_perm, h, w)
+    outputs = {
+        "astronaut_original.jpg": img,
+        "astronaut_noisy.jpg": noisy,
+        "astronaut_knn.jpg": knn_img,
+        "astronaut_grid.jpg": grid_img,
+        "astronaut_permutohedral.jpg": perm_img,
+    }
+    for name, arr in outputs.items():
+        path = os.path.join(args.out_dir, name)
+        _save_image(path, arr)
+        print(f"Saved {path}")
+
+    # ---- PSNR vs original (data_range=1.0 since img is float in [0, 1]) ---
+    from skimage.metrics import peak_signal_noise_ratio as psnr
+
+    ref = img.astype(np.float32)
+    psnr_noisy = psnr(ref, np.clip(noisy, 0, 1).astype(np.float32), data_range=1.0)
+    psnr_knn = psnr(ref, np.clip(knn_img, 0, 1).astype(np.float32), data_range=1.0)
+    psnr_grid = psnr(ref, np.clip(grid_img, 0, 1).astype(np.float32), data_range=1.0)
+    psnr_perm = psnr(ref, np.clip(perm_img, 0, 1).astype(np.float32), data_range=1.0)
+
+    print()
+    print(f"  {'Stage':<22}{'Time':>10}   PSNR (dB)")
+    print(f"  {'-' * 46}")
+    print(f"  {'Noisy input':<22}{'-':>10}   {psnr_noisy:6.2f}")
+    print(f"  {'KNN (k=' + str(args.knn_k) + ')':<22}{t_knn*1e3:>8.1f} ms   {psnr_knn:6.2f}")
+    print(f"  {'Grid':<22}{t_grid*1e3:>8.1f} ms   {psnr_grid:6.2f}")
+    print(f"  {'Permutohedral':<22}{t_perm*1e3:>8.1f} ms   {psnr_perm:6.2f}")
+
+
+if __name__ == "__main__":
+    main()