v4:Two-stage cascade (WNet)

.gitignore

@@ -8,4 +8,5 @@ __pycache__/
 .DS_Store
 *.onnx
 checkpoints_diag/
-checkpoints_v2/
+checkpoints_v2/
+checkpoints_v3/

checkpoints_v4/history.csv

@@ -0,0 +1,21 @@
+epoch,train_mse,val_mse
+1,0.000914,0.000824
+2,0.000637,0.000939
+3,0.000570,0.000672
+4,0.000504,0.000629
+5,0.000455,0.000595
+6,0.000416,0.000640
+7,0.000391,0.000564
+8,0.000365,0.000544
+9,0.000344,0.000527
+10,0.000327,0.000515
+11,0.000309,0.000534
+12,0.000297,0.000504
+13,0.000284,0.000517
+14,0.000271,0.000462
+15,0.000263,0.000468
+16,0.000250,0.000450
+17,0.000243,0.000460
+18,0.000234,0.000448
+19,0.000227,0.000446
+20,0.000220,0.000443

src/export_onnx.py

@@ -1,45 +1,56 @@
 """
-export_onnx.py — export the trained U-Net to ONNX and verify it matches PyTorch.
+export_onnx.py — export a trained model to ONNX, auto-detecting config from the checkpoint.
-Run from the project root:  python src\\export_onnx.py
+
+  python src\\export_onnx.py --ckpt checkpoints\\best.pt        --out web\\radio_unet_v1.onnx
+  python src\\export_onnx.py --ckpt checkpoints_v2\\best.pt     --out web\\radio_unet_v2.onnx
+  python src\\export_onnx.py --ckpt checkpoints_v3\\best.pt     --out web\\radio_unet_v3.onnx
+  python src\\export_onnx.py --ckpt checkpoints_v4\\best.pt     --out web\\radio_unet_v4.onnx
 """
+import argparse
 import numpy as np
 import torch
 
-from model import UNet
+from model import UNet, WNet
 from dataset import RadioMapSeerDataset
 
-CKPT = "checkpoints/best.pt"
-OUT = "web/radio_unet.onnx"
 
-# --- load the trained model ---
+def main():
-ckpt = torch.load(CKPT, map_location="cpu", weights_only=False)
+    p = argparse.ArgumentParser()
-base = ckpt.get("args", {}).get("base", 64)
+    p.add_argument("--ckpt", default="checkpoints/best.pt")
-model = UNet(2, 1, base=base)
+    p.add_argument("--out", default="web/radio_unet.onnx")
-model.load_state_dict(ckpt["model_state"])
+    args = p.parse_args()
-model.eval()
 
-# --- export to ONNX (fixed 1x2x256x256 input) ---
+    ckpt = torch.load(args.ckpt, map_location="cpu", weights_only=False)
-dummy = torch.randn(1, 2, 256, 256)
+    cargs = ckpt.get("args", {})
-torch.onnx.export(
+    base = cargs.get("base", 64)
-    model, dummy, OUT,
+    add_dist = bool(cargs.get("distance", False))
-    input_names=["input"], output_names=["coverage"],
+    up_mode = cargs.get("up_mode", "deconv")
-    opset_version=17,
+    arch = cargs.get("arch", "unet")
-)
+    in_ch = 3 if add_dist else 2
-print(f"Exported {OUT}")
 
-# --- verify ONNX output matches PyTorch on a real test sample ---
+    Net = WNet if arch == "wnet" else UNet
-import onnxruntime as ort
+    model = Net(in_channels=in_ch, out_channels=1, base=base, up_mode=up_mode)
+    model.load_state_dict(ckpt["model_state"])
+    model.eval()
+    print(f"Loaded {args.ckpt} | arch={arch} | in_channels={in_ch} | "
+          f"distance={add_dist} | up_mode={up_mode}")
 
-ds = RadioMapSeerDataset(split="test", img_size=256)
+    dummy = torch.randn(1, in_ch, 256, 256)
-x, _ = ds[0]
+    torch.onnx.export(model, dummy, args.out,
-xb = x.unsqueeze(0).numpy().astype(np.float32)
+                      input_names=["input"], output_names=["coverage"],
+                      opset_version=17)
+    print(f"Exported {args.out}")
 
-with torch.no_grad():
+    import onnxruntime as ort
-    torch_out = model(x.unsqueeze(0)).numpy()
+    ds = RadioMapSeerDataset(split="test", img_size=256, add_distance=add_dist)
+    x, _ = ds[0]
+    xb = x.unsqueeze(0).numpy().astype(np.float32)
+    with torch.no_grad():
+        t_out = model(x.unsqueeze(0)).numpy()
+    sess = ort.InferenceSession(args.out, providers=["CPUExecutionProvider"])
+    o_out = sess.run(["coverage"], {"input": xb})[0]
+    print(f"Max |PyTorch - ONNX|: {np.max(np.abs(t_out - o_out)):.2e}")
 
-sess = ort.InferenceSession(OUT, providers=["CPUExecutionProvider"])
-onnx_out = sess.run(["coverage"], {"input": xb})[0]
 
-max_diff = float(np.max(np.abs(torch_out - onnx_out)))
+if __name__ == "__main__":
-print(f"Max |PyTorch - ONNX| on a test sample: {max_diff:.2e}")
+    main()
-print("OK — outputs match" if max_diff < 1e-3 else "WARNING: large mismatch, stop and check")

src/model.py

@@ -42,19 +42,26 @@ class Down(nn.Module):
 
 class Up(nn.Module):
     """Upsample, concatenate the skip connection, then DoubleConv."""
-    def __init__(self, in_ch, skip_ch, out_ch):
+    def __init__(self, in_ch, skip_ch, out_ch, up_mode="deconv"):
         super().__init__()
-        self.up = nn.ConvTranspose2d(in_ch, in_ch // 2, kernel_size=2, stride=2)
+        if up_mode == "resize":
+            # bilinear upsample + 1x1 conv — uniform, so no checkerboard artifacts
+            self.up = nn.Sequential(
+                nn.Upsample(scale_factor=2, mode="bilinear", align_corners=False),
+                nn.Conv2d(in_ch, in_ch // 2, kernel_size=1),
+            )
+        else:  # "deconv" (v1/v2): transposed convolution
+            self.up = nn.ConvTranspose2d(in_ch, in_ch // 2, kernel_size=2, stride=2)
         self.conv = DoubleConv(in_ch // 2 + skip_ch, out_ch)
 
     def forward(self, x, skip):
         x = self.up(x)
-        x = torch.cat([skip, x], dim=1)   # skip connection
+        x = torch.cat([skip, x], dim=1)
         return self.conv(x)
 
 
 class UNet(nn.Module):
-    def __init__(self, in_channels=2, out_channels=1, base=64):
+    def __init__(self, in_channels=2, out_channels=1, base=64, up_mode="deconv"):
         super().__init__()
         self.inc   = DoubleConv(in_channels, base)   # 256
         self.down1 = Down(base,     base * 2)        # 128
@@ -62,10 +69,10 @@ class UNet(nn.Module):
         self.down3 = Down(base * 4, base * 8)        # 32
         self.down4 = Down(base * 8, base * 8)        # 16  (bottleneck)
 
-        self.up1 = Up(base * 8, base * 8, base * 4)
+        self.up1 = Up(base * 8, base * 8, base * 4, up_mode)
-        self.up2 = Up(base * 4, base * 4, base * 2)
+        self.up2 = Up(base * 4, base * 4, base * 2, up_mode)
-        self.up3 = Up(base * 2, base * 2, base)
+        self.up3 = Up(base * 2, base * 2, base,     up_mode)
-        self.up4 = Up(base,     base,     base)
+        self.up4 = Up(base,     base,     base,     up_mode)
 
         self.outc = nn.Conv2d(base, out_channels, kernel_size=1)
 
@@ -82,17 +89,44 @@ class UNet(nn.Module):
         x = self.up4(x, x1)
         return torch.sigmoid(self.outc(x))
 
+class WNet(nn.Module):
+    """Two cascaded U-Nets (RadioUNet-style): a coarse predictor + a refiner.
+
+    unet1: input channels         -> coarse coverage map (1 ch)
+    unet2: input channels + coarse -> refined coverage map (1 ch)
+
+    The refiner sees a global coverage estimate as an extra input channel,
+    so every pixel has long-range context -> larger effective receptive field.
+    forward() returns the refined map by default (so eval/ONNX stay single-output);
+    pass return_coarse=True during training for deep supervision.
+    """
+    def __init__(self, in_channels=3, out_channels=1, base=64, up_mode="resize"):
+        super().__init__()
+        self.unet1 = UNet(in_channels,     out_channels, base=base, up_mode=up_mode)
+        self.unet2 = UNet(in_channels + 1, out_channels, base=base, up_mode=up_mode)
+
+    def forward(self, x, return_coarse=False):
+        coarse = self.unet1(x)
+        refined = self.unet2(torch.cat([x, coarse], dim=1))
+        if return_coarse:
+            return refined, coarse
+        return refined
 
 if __name__ == "__main__":
     device = "cuda" if torch.cuda.is_available() else "cpu"
-    model = UNet(in_channels=2, out_channels=1, base=64).to(device)
-    n_params = sum(p.numel() for p in model.parameters())
     print(f"Device: {device}")
-    print(f"Parameters: {n_params:,}")
 
-    x = torch.randn(4, 2, 256, 256, device=device)
+    unet = UNet(in_channels=3, out_channels=1, base=64, up_mode="resize").to(device)
-    y = model(x)
+    print(f"UNet params: {sum(p.numel() for p in unet.parameters()):,}")
-    print("Input :", tuple(x.shape))
+
-    print("Output:", tuple(y.shape), "| range:", round(float(y.min()), 3), "-", round(float(y.max()), 3))
+    wnet = WNet(in_channels=3, out_channels=1, base=64, up_mode="resize").to(device)
+    print(f"WNet params: {sum(p.numel() for p in wnet.parameters()):,}")
+
+    x = torch.randn(2, 3, 256, 256, device=device)
+    refined, coarse = wnet(x, return_coarse=True)
+    print("Input  :", tuple(x.shape))
+    print("Refined:", tuple(refined.shape),
+          "| range:", round(float(refined.min()), 3), "-", round(float(refined.max()), 3))
+    print("Coarse :", tuple(coarse.shape))
     if device == "cuda":
-        print(f"Peak VRAM (forward, batch 4): {torch.cuda.max_memory_allocated()/1e9:.2f} GB")
+        print(f"Peak VRAM (WNet forward, batch 2): {torch.cuda.max_memory_allocated()/1e9:.2f} GB")

src/train.py

@@ -1,14 +1,14 @@
 """
-train.py — train the U-Net on RadioMapSeer coverage maps.
+train.py — train the radio-coverage models on RadioMapSeer.
 
-v1: 2-channel input (buildings, transmitter).
+v1: 2-channel UNet.  v2: --distance (3-channel).  v3: --up-mode resize.
-v2: add --distance for a 3rd channel (distance-to-Tx) that fixes the square cutoff.
+v4: --arch wnet  (two cascaded U-Nets: coarse + refine, deep supervision).
 
 Examples:
-  # v1
+  # v3 (single U-Net, resize-conv)
-  python src\\train.py --epochs 30 --batch-size 32 --num-workers 8
+  python src\\train.py --distance --up-mode resize --epochs 30 --batch-size 32 --num-workers 8 --ckpt-dir checkpoints_v3
-  # v2 (separate checkpoint dir so v1 stays intact)
+  # v4 (WNet), warm-starting stage 1 from the trained v3
-  python src\\train.py --distance --epochs 30 --batch-size 32 --num-workers 8 --ckpt-dir checkpoints_v2
+  python src\\train.py --arch wnet --distance --up-mode resize --init-unet1-from checkpoints_v3\\best.pt --epochs 25 --batch-size 16 --num-workers 8 --ckpt-dir checkpoints_v4
 """
 
 import argparse
@@ -21,12 +21,13 @@ import torch.nn as nn
 from torch.utils.data import DataLoader
 
 from dataset import RadioMapSeerDataset
-from model import UNet
+from model import UNet, WNet
 
 torch.backends.cudnn.benchmark = True   # autotune convs (input size is fixed)
 
 
 def evaluate(net, loader, criterion, device, amp):
+    """Validation MSE on the final output (WNet returns the refined map by default)."""
     net.eval()
     total, n = 0.0, 0
     with torch.no_grad():
@@ -50,14 +51,25 @@ def main():
     p.add_argument("--num-workers", type=int, default=4)
     p.add_argument("--ckpt-dir", default="checkpoints")
     p.add_argument("--distance", action="store_true",
-                   help="v2: add a distance-to-transmitter input channel")
+                   help="v2+: add a distance-to-transmitter input channel")
     p.add_argument("--compile", action="store_true",
                    help="wrap model in torch.compile (optional; flaky on Windows)")
+    p.add_argument("--up-mode", default="deconv", choices=["deconv", "resize"],
+                   help="v3+: 'resize' = bilinear upsample + conv (removes checkerboard)")
+    p.add_argument("--arch", default="unet", choices=["unet", "wnet"],
+                   help="v4: 'wnet' = two cascaded U-Nets (coarse + refine)")
+    p.add_argument("--aux-weight", type=float, default=0.4,
+                   help="WNet deep-supervision weight on the coarse output")
+    p.add_argument("--init-from", default=None,
+                   help="warm-start: load full model weights from a checkpoint")
+    p.add_argument("--init-unet1-from", default=None,
+                   help="WNet: warm-start stage-1 (unet1) from a single-UNet checkpoint (e.g. v3)")
     args = p.parse_args()
 
     device = "cuda" if torch.cuda.is_available() else "cpu"
     amp = (device == "cuda")
-    print(f"Device: {device} | AMP: {amp} | distance: {args.distance} | compile: {args.compile}")
+    print(f"Device: {device} | AMP: {amp} | arch: {args.arch} | distance: {args.distance} | "
+          f"up_mode: {args.up_mode} | compile: {args.compile}")
 
     train_ds = RadioMapSeerDataset(split="train", img_size=args.img_size,
                                    max_samples=args.max_samples, add_distance=args.distance)
@@ -72,8 +84,26 @@ def main():
     val_loader = DataLoader(val_ds, shuffle=False, **dl_kwargs)
 
     in_ch = 3 if args.distance else 2
-    model = UNet(in_channels=in_ch, out_channels=1, base=args.base).to(device)
+    if args.arch == "wnet":
+        model = WNet(in_channels=in_ch, out_channels=1, base=args.base, up_mode=args.up_mode).to(device)
+    else:
+        model = UNet(in_channels=in_ch, out_channels=1, base=args.base, up_mode=args.up_mode).to(device)
+
+    # WNet: warm-start stage 1 from an existing single-UNet checkpoint (same in_ch + up_mode)
+    if args.arch == "wnet" and args.init_unet1_from:
+        ck1 = torch.load(args.init_unet1_from, map_location=device, weights_only=False)
+        model.unet1.load_state_dict(ck1["model_state"])
+        print(f"Warm-started unet1 from {args.init_unet1_from} "
+              f"(epoch {ck1.get('epoch','?')}, val {ck1.get('val_loss','?')})")
+
     net = torch.compile(model) if args.compile else model
+
+    if args.init_from:
+        ckpt0 = torch.load(args.init_from, map_location=device, weights_only=False)
+        model.load_state_dict(ckpt0["model_state"])
+        print(f"Warm-started from {args.init_from} "
+              f"(epoch {ckpt0.get('epoch','?')}, val {ckpt0.get('val_loss','?')})")
+
     criterion = nn.MSELoss()
     optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
     scaler = torch.amp.GradScaler(enabled=amp)
@@ -94,20 +124,26 @@ def main():
             x, y = x.to(device), y.to(device)
             optimizer.zero_grad(set_to_none=True)
             with torch.autocast(device_type=device, dtype=torch.float16, enabled=amp):
-                loss = criterion(net(x), y)
+                if args.arch == "wnet":
+                    refined, coarse = net(x, return_coarse=True)
+                    main_loss = criterion(refined, y)                       # the real metric
+                    loss = main_loss + args.aux_weight * criterion(coarse, y)  # + deep supervision
+                else:
+                    main_loss = criterion(net(x), y)
+                    loss = main_loss
             scaler.scale(loss).backward()
             scaler.step(optimizer)
             scaler.update()
-            running += loss.item() * x.size(0)
+            running += main_loss.item() * x.size(0)   # log refined MSE (comparable to val)
             n += x.size(0)
             if i % 50 == 0:
-                print(f"  epoch {epoch} step {i}/{len(train_loader)} | loss {running/n:.5f}")
+                print(f"  epoch {epoch} step {i}/{len(train_loader)} | loss {running/n:.6f}")
 
         train_loss = running / n
         val_loss = evaluate(net, val_loader, criterion, device, amp)
         dt = time.time() - t0
-        print(f"[epoch {epoch}/{args.epochs}] train MSE {train_loss:.5f} | "
+        print(f"[epoch {epoch}/{args.epochs}] train MSE {train_loss:.6f} | "
-              f"val MSE {val_loss:.5f} | val RMSE {val_loss**0.5:.5f} | {dt:.0f}s")
+              f"val MSE {val_loss:.6f} | val RMSE {val_loss**0.5:.5f} | {dt:.0f}s")
         with open(history_path, "a", newline="") as f:
             csv.writer(f).writerow([epoch, f"{train_loss:.6f}", f"{val_loss:.6f}"])
 
@@ -116,13 +152,13 @@ def main():
             torch.save({"epoch": epoch, "model_state": model.state_dict(),
                         "val_loss": val_loss, "args": vars(args)},
                        ckpt_dir / "best.pt")
-            print(f"           -> saved best.pt (val MSE {val_loss:.5f})")
+            print(f"           -> saved best.pt (val MSE {val_loss:.6f})")
 
         if device == "cuda":
             print(f"           peak VRAM: {torch.cuda.max_memory_allocated()/1e9:.2f} GB")
             torch.cuda.reset_peak_memory_stats()
 
-    print(f"\nDone. Best val MSE: {best_val:.5f}  (RMSE {best_val**0.5:.5f})")
+    print(f"\nDone. Best val MSE: {best_val:.6f}  (RMSE {best_val**0.5:.5f})")
 
 
 if __name__ == "__main__":

src/visualize.py

@@ -16,7 +16,8 @@ import torch
 import matplotlib.pyplot as plt
 
 from dataset import RadioMapSeerDataset
-from model import UNet
+from model import UNet, WNet
+
 
 
 def main():
@@ -34,8 +35,15 @@ def main():
     cargs = ckpt.get("args", {})
     base = cargs.get("base", 64)
     add_dist = bool(cargs.get("distance", False))
+    up_mode = cargs.get("up_mode", "deconv")
     in_ch = 3 if add_dist else 2
-    model = UNet(in_channels=in_ch, out_channels=1, base=base).to(device)
+    arch = cargs.get("arch", "unet")
+    if arch == "wnet":
+        model = WNet(in_channels=in_ch, out_channels=1, base=base, up_mode=up_mode).to(device)
+    else:
+        arch = cargs.get("arch", "unet")
+        Net = WNet if arch == "wnet" else UNet
+        model = Net(in_channels=in_ch, out_channels=1, base=base, up_mode=up_mode).to(device)
     model.load_state_dict(ckpt["model_state"])
     model.eval()
     print(f"Loaded {args.ckpt} | in_channels={in_ch} | distance={add_dist} | "

web/index.html

@@ -6,43 +6,40 @@
 <title>Radio Coverage Prediction — Deep Learning Demo</title>
 <script src="https://cdn.jsdelivr.net/npm/onnxruntime-web/dist/ort.min.js"></script>
 <style>
-  :root { --bg:#0f1115; --panel:#1a1d24; --line:#2a2e37; --fg:#e6e8eb; --muted:#9aa0a6; --accent:#3ddc84; }
+  :root { --bg:#0f1115; --panel:#1a1d24; --line:#2a2e37; --fg:#e6e8eb; --muted:#9aa0a6; --accent:#3ddc84; --warn:#f0a93d; }
   * { box-sizing:border-box; }
   body { margin:0; background:var(--bg); color:var(--fg);
          font-family: system-ui,-apple-system,"Segoe UI","PingFang SC","Microsoft YaHei",Roboto,sans-serif; line-height:1.6; }
   .wrap { max-width:920px; margin:0 auto; padding:28px 20px 72px; }
-  .topbar { display:flex; justify-content:flex-end; gap:6px; margin-bottom:18px; }
+  .topbar { display:flex; justify-content:flex-end; gap:6px; margin-bottom:16px; }
-  .lang { cursor:pointer; border:1px solid var(--line); background:transparent; color:var(--muted);
+  .lang { cursor:pointer; border:1px solid var(--line); background:transparent; color:var(--muted); padding:5px 12px; border-radius:20px; font-size:.85rem; }
-          padding:5px 12px; border-radius:20px; font-size:.85rem; }
+  .lang.active { background:var(--accent); color:#0f1115; border-color:var(--accent); font-weight:500; }
-  .lang.active { background:var(--accent); color:#0f1115; border-color:var(--accent); font-weight:600; }
   h1 { font-size:1.7rem; margin:0 0 6px; }
-  .sub { color:var(--muted); margin:0 0 28px; font-size:1.05rem; }
+  .sub { color:var(--muted); margin:0 0 24px; }
-  h2 { font-size:1.15rem; margin:36px 0 10px; }
+  h2 { font-size:1.15rem; margin:34px 0 10px; }
-  p { margin:0 0 12px; }
+  .vers { display:flex; gap:10px; flex-wrap:wrap; margin:0 0 14px; }
-  .metrics { display:flex; gap:14px; flex-wrap:wrap; margin:18px 0 8px; }
+  .ver-btn { cursor:pointer; border:1px solid var(--line); background:var(--panel); color:var(--fg);
-  .metric { background:var(--panel); border-radius:10px; padding:12px 16px; flex:1 1 150px; }
+             padding:10px 18px; border-radius:10px; font-size:.95rem; font-weight:500; }
-  .metric .v { font-size:1.5rem; font-weight:700; color:var(--accent); }
+  .ver-btn.active { background:var(--accent); color:#0f1115; border-color:var(--accent); }
-  .metric .l { color:var(--muted); font-size:.85rem; }
+  .ver-btn:disabled { opacity:.4; cursor:not-allowed; }
-  .demo { background:var(--panel); border-radius:14px; padding:20px; margin-top:14px; }
+  .notice { background:var(--panel); border-left:3px solid var(--accent); border-radius:0 8px 8px 0;
+            padding:14px 16px; margin:0 0 18px; font-size:.95rem; }
+  .metricbar { display:flex; gap:14px; flex-wrap:wrap; margin:0 0 8px; }
+  .metric { background:var(--panel); border-radius:10px; padding:10px 16px; }
+  .metric .v { font-size:1.4rem; font-weight:500; color:var(--accent); }
+  .metric .l { color:var(--muted); font-size:.82rem; }
+  .demo { background:var(--panel); border-radius:14px; padding:20px; }
   .row { display:flex; gap:24px; flex-wrap:wrap; align-items:flex-start; }
   .col { flex:1 1 256px; }
-  canvas { width:100%; max-width:340px; height:auto; image-rendering:pixelated;
+  canvas { width:100%; max-width:340px; height:auto; image-rendering:pixelated; border-radius:8px; background:#000; cursor:crosshair; display:block; }
-           border-radius:8px; background:#000; cursor:crosshair; display:block; }
+  .col h3 { font-size:.95rem; margin:0 0 8px; font-weight:500; }
-  .col h3 { font-size:.95rem; margin:0 0 8px; font-weight:600; }
   .controls { margin:0 0 18px; display:flex; gap:12px; align-items:center; flex-wrap:wrap; }
   select { background:#0f1115; color:var(--fg); border:1px solid var(--line); border-radius:8px; padding:8px 10px; }
   .status { color:var(--muted); font-size:.9rem; margin-top:14px; min-height:1.2em; }
   .hint { color:var(--muted); font-size:.85rem; }
-  .legend { height:12px; border-radius:6px; margin-top:6px;
+  .legend { height:12px; border-radius:6px; margin-top:6px; background:linear-gradient(90deg,#440154,#3b528b,#21918c,#5ec962,#fde725); }
-            background:linear-gradient(90deg,#440154,#3b528b,#21918c,#5ec962,#fde725); }
   .legend-labels { display:flex; justify-content:space-between; color:var(--muted); font-size:.75rem; }
-  .specs { border-top:1px solid var(--line); margin-top:8px; }
+  footer { margin-top:36px; color:var(--muted); font-size:.8rem; border-top:1px solid var(--line); padding-top:16px; }
-  .spec { display:flex; gap:16px; padding:11px 0; border-bottom:1px solid var(--line); }
-  .spec .k { flex:0 0 130px; color:var(--accent); font-size:.9rem; }
-  .spec .val { color:var(--fg); font-size:.95rem; }
-  .tags { display:flex; gap:8px; flex-wrap:wrap; margin-top:10px; }
-  .tag { background:var(--panel); border:1px solid var(--line); border-radius:20px; padding:5px 12px; font-size:.85rem; color:var(--muted); }
-  footer { margin-top:40px; color:var(--muted); font-size:.8rem; border-top:1px solid var(--line); padding-top:16px; }
 </style>
 </head>
 <body>
@@ -55,21 +52,19 @@
   <h1 data-i18n="title"></h1>
   <p class="sub" data-i18n="subtitle"></p>
 
-  <div class="metrics">
+  <h2 data-i18n="ver_h"></h2>
-    <div class="metric"><div class="v">4.1 dB</div><div class="l" data-i18n="m_rmse"></div></div>
+  <div class="vers">
-    <div class="metric"><div class="v">71%</div><div class="l" data-i18n="m_improve"></div></div>
+    <button class="ver-btn" data-v="v1" data-i18n="btn_v1"></button>
-    <div class="metric"><div class="v">~13M</div><div class="l" data-i18n="m_params"></div></div>
+    <button class="ver-btn" data-v="v2" data-i18n="btn_v2"></button>
-    <div class="metric"><div class="v">&lt;2 s</div><div class="l" data-i18n="m_speed"></div></div>
+    <button class="ver-btn" data-v="v3" data-i18n="btn_v3"></button>
+    <button class="ver-btn" data-v="v4" data-i18n="btn_v4"></button>
+  </div>
+  <div class="notice" id="verNotice"></div>
+  <div class="metricbar">
+    <div class="metric"><div class="v" id="verMetric"></div><div class="l" data-i18n="m_error"></div></div>
   </div>
 
-  <h2 data-i18n="problem_h"></h2>
-  <p data-i18n="problem_p"></p>
-
-  <h2 data-i18n="how_h"></h2>
-  <p data-i18n="how_p"></p>
-
   <h2 data-i18n="demo_h"></h2>
-  <p data-i18n="demo_p"></p>
   <div class="demo">
     <div class="controls">
       <label for="mapSel" data-i18n="select_label"></label>
@@ -91,23 +86,8 @@
     <div class="status" id="status"></div>
   </div>
 
-  <h2 data-i18n="results_h"></h2>
+  <h2 data-i18n="how_h"></h2>
-  <p data-i18n="results_p"></p>
+  <p data-i18n="how_p"></p>
-
-  <h2 data-i18n="hood_h"></h2>
-  <div class="specs">
-    <div class="spec"><div class="k" data-i18n="s_data_k"></div><div class="val" data-i18n="s_data_v"></div></div>
-    <div class="spec"><div class="k" data-i18n="s_eval_k"></div><div class="val" data-i18n="s_eval_v"></div></div>
-    <div class="spec"><div class="k" data-i18n="s_model_k"></div><div class="val" data-i18n="s_model_v"></div></div>
-    <div class="spec"><div class="k" data-i18n="s_train_k"></div><div class="val" data-i18n="s_train_v"></div></div>
-    <div class="spec"><div class="k" data-i18n="s_deploy_k"></div><div class="val" data-i18n="s_deploy_v"></div></div>
-  </div>
-
-  <h2 data-i18n="stack_h"></h2>
-  <div class="tags">
-    <span class="tag">Python</span><span class="tag">PyTorch</span><span class="tag">U-Net / CNN</span>
-    <span class="tag">ONNX</span><span class="tag">ONNX Runtime Web</span><span class="tag">NumPy</span>
-  </div>
 
   <footer data-i18n="footer"></footer>
 </div>
@@ -116,88 +96,63 @@
 const T = {
   en: {
     title: "Predicting Wireless Coverage with Deep Learning",
-    subtitle: "A U-Net neural network predicts a base station's downlink signal coverage across a city — running live in your browser.",
+    subtitle: "A U-Net predicts a base station's signal coverage across a city — running live in your browser. Switch between model versions to see how it improved.",
-    m_rmse: "RMSE on unseen cities", m_improve: "lower error vs. classical model",
+    ver_h: "Model version",
-    m_params: "parameters", m_speed: "in-browser inference",
+    btn_v1: "v1 · baseline", btn_v2: "v2 · distance fix", btn_v3: "v3 · checkerboard fix", btn_v4: "v4 · WNet refine",
-    problem_h: "The problem",
+    m_error: "RMSE on held-out cities",
-    problem_p: "Before building a base station, mobile operators need to know how its signal will spread across a city — where coverage is strong, and where buildings cast \"shadows\" that block it. Traditionally this is done either with fast but crude empirical formulas (like the Hata model) that ignore building geometry entirely, or with slow, physically accurate ray-tracing simulations that can take minutes per map. This project asks: can a neural network match the accuracy of the slow simulation, but in milliseconds?",
+    notice_v1: "Baseline (2-channel input: buildings + transmitter). Known issue: coverage cuts off in a square around the transmitter — the far field stays blank because the network can't propagate the source that far.",
-    how_h: "How it works",
+    notice_v2: "Adds a distance-to-transmitter input channel (3-channel). Fixed: the square cutoff — coverage now spans the whole map, with far-field rays and distant building shadows. Known issue: faint checkerboard texture in smooth areas, from transposed-convolution upsampling.",
-    how_p: "The model is a U-Net — a convolutional encoder-decoder with skip connections. It takes two inputs as image channels (the city's building layout and the transmitter's location) and outputs one image: the predicted coverage map, where bright means strong signal and dark means weak. The network was never given any propagation equations — it learned signal falloff with distance and shadowing behind buildings entirely from examples.",
+    notice_v3: "Replaces transposed-convolution upsampling with resize-convolution. Fixed: the checkerboard texture — smooth regions are clean. Known issue: along a long unobstructed corridor the far-field beam fades before reaching the edge — a residual receptive-field limit.",
+    notice_v4: "Two cascaded U-Nets — a coarse predictor feeds a refiner. Because the refiner's input already contains a global coverage estimate, every pixel gets long-range context (a larger effective receptive field). Fixed: the fading corridor — coverage stays bright to the edge. Lowest error of all four versions.",
     demo_h: "Try it yourself",
-    demo_p: "Pick a city layout and click anywhere to place a transmitter. The model runs entirely in your browser and renders the predicted coverage in a fraction of a second. Place it near buildings and watch the signal shadows form.",
     select_label: "City layout:", click_hint: "Click the map to place a transmitter.",
     input_title: "City map — click to place transmitter", output_title: "Predicted coverage",
     legend_weak: "weak", legend_strong: "strong",
-    results_h: "Results",
+    how_h: "How it works",
-    results_p: "Evaluated on test cities the model never saw during training, it reaches 4.1 dB RMSE — a 71% reduction in error compared to a classical log-distance path-loss model (14.3 dB), which can only model distance and is blind to buildings. That gap is exactly the value the network adds: it learned to account for the building shadowing the classical model cannot represent.",
+    how_p: "The model takes the building layout and the transmitter location as image channels and outputs a coverage heatmap, learned from the public RadioMapSeer dataset (5.9 GHz urban propagation). It was given no propagation equations — it learned signal falloff and building shadowing from data. The model runs entirely in your browser via ONNX Runtime Web; nothing is sent to a server.",
-    hood_h: "Under the hood",
+    footer: "Dataset: RadioMapSeer (Yapar et al., 2022), CC BY 4.0. Independent portfolio project. Errors are on held-out cities the model never trained on.",
-    s_data_k: "Dataset", s_data_v: "RadioMapSeer — 56,000+ ray-traced coverage maps from real cities (Berlin, London, Tel Aviv…) at 5.9 GHz.",
+    st_loading: "Loading model…", st_ready: "Pick a layout and click to place a transmitter.",
-    s_eval_k: "Evaluation", s_eval_v: "Train/test split by city, so the model is measured on building layouts it has never seen.",
-    s_model_k: "Model", s_model_v: "U-Net CNN, ~13M parameters, image-to-image regression with MSE loss and a sigmoid output.",
-    s_train_k: "Training", s_train_v: "PyTorch on a single NVIDIA RTX 2080, automatic mixed precision (AMP), Adam optimizer.",
-    s_deploy_k: "Deployment", s_deploy_v: "Exported to ONNX and run fully client-side via ONNX Runtime Web — no server required.",
-    stack_h: "Tech stack",
-    footer: "Dataset: RadioMapSeer (Yapar et al., 2022), CC BY 4.0. Independent portfolio project.",
-    st_loading: "Loading model… (downloads once, ~50 MB)",
-    st_ready: "Model ready. Pick a layout and click to place a transmitter.",
     st_running: "Running inference…",
-    st_done: (ms) => "Coverage predicted in " + ms + " ms. Click again to move the transmitter.",
+    st_done: (ms) => "Predicted in " + ms + " ms. Click again to move the transmitter.",
     st_error: (m) => "Error: " + m
   },
   zh: {
     title: "基于深度学习的无线信号覆盖预测",
-    subtitle: "一个 U-Net 神经网络直接预测基站在城市中的下行信号覆盖——并在你的浏览器中实时运行。",
+    subtitle: "一个 U-Net 在你的浏览器中实时预测基站的信号覆盖。切换不同模型版本,看看它是如何一步步改进的。",
-    m_rmse: "未见城市上的 RMSE", m_improve: "较经典模型误差降低",
+    ver_h: "模型版本",
-    m_params: "参数量", m_speed: "浏览器端推理",
+    btn_v1: "v1 · 基线", btn_v2: "v2 · 距离修复", btn_v3: "v3 · 棋盘格修复", btn_v4: "v4 · WNet 精修",
-    problem_h: "问题背景",
+    m_error: "留出城市上的 RMSE",
-    problem_p: "在建设基站之前，移动运营商需要预估信号在城市中的传播情况——哪里覆盖良好，哪里因建筑遮挡形成信号「阴影」。传统方法要么使用快速但粗糙的经验公式（如 Hata 模型），完全忽略建筑几何结构；要么使用物理精确但缓慢的射线追踪仿真，每张地图可能耗时数分钟。本项目要回答的问题是：神经网络能否在毫秒级时间内，达到接近慢速仿真的预测精度？",
+    notice_v1: "基线模型(2 通道输入:建筑 + 发射机)。已知问题:覆盖在发射机周围呈方形截断——远场为空白,因为网络无法将信号源的影响传播到远处。",
-    how_h: "实现原理",
+    notice_v2: "新增「到发射机的距离」输入通道(3 通道)。已修复:方形截断——覆盖现已贯穿整张地图,远场射线与远处建筑阴影都能呈现。已知问题:平滑区域出现轻微棋盘格纹理,来自转置卷积上采样。",
-    how_p: "模型是一个 U-Net——带跳跃连接的卷积编码器-解码器。它以两个图像通道作为输入（城市建筑布局与发射机位置），输出一张图像：预测的信号覆盖热力图，越亮表示信号越强，越暗表示越弱。网络从未被告知任何传播公式，它完全从样本数据中自行学会了信号随距离衰减、以及在建筑后方形成阴影的规律。",
+    notice_v3: "用 resize 卷积(双线性上采样 + 卷积)替换转置卷积。已修复:棋盘格纹理消失,平滑区域变得干净。已知问题:在无遮挡的长走廊中,远场光束在到达边缘前逐渐变暗——残留的感受野限制。",
+    notice_v4: "两个级联 U-Net——粗预测网络的输出送入精修网络。由于精修网络的输入已包含一张全局覆盖估计,每个像素都获得了长程上下文(更大的有效感受野)。已修复:长走廊衰减,覆盖一直明亮延伸到边缘。四个版本中误差最低。",
     demo_h: "在线体验",
-    demo_p: "选择一个城市布局，点击任意位置放置发射机。模型完全在你的浏览器中运行，不到一秒即可渲染出预测覆盖。把发射机放在建筑附近，观察信号阴影是如何形成的。",
+    select_label: "城市布局:", click_hint: "点击地图以放置发射机。",
-    select_label: "城市布局：", click_hint: "点击地图以放置发射机。",
     input_title: "城市地图——点击放置发射机", output_title: "预测覆盖",
     legend_weak: "弱", legend_strong: "强",
-    results_h: "实验结果",
+    how_h: "实现原理",
-    results_p: "在训练中从未见过的测试城市上评估，模型达到 4.1 dB 的均方根误差（RMSE）——相比只能建模距离、对建筑「视而不见」的经典对数距离路径损耗模型（14.3 dB），误差降低了 71%。这一差距正是神经网络的价值所在：它学会了经典模型无法表达的建筑遮蔽效应。",
+    how_p: "模型以建筑布局和发射机位置作为图像通道输入,输出覆盖热力图,训练自公开的 RadioMapSeer 数据集(5.9 GHz 城市传播)。它没有被告知任何传播公式,完全从数据中学会了信号衰减与建筑遮蔽。模型通过 ONNX Runtime Web 完全在你的浏览器中运行,不向服务器发送任何数据。",
-    hood_h: "技术细节",
+    footer: "数据集:RadioMapSeer(Yapar 等,2022),CC BY 4.0 许可。独立作品集项目。误差基于模型从未训练过的留出城市。",
-    s_data_k: "数据集", s_data_v: "RadioMapSeer——基于真实城市（柏林、伦敦、特拉维夫等）、5.9 GHz、56,000+ 张射线追踪仿真覆盖图。",
+    st_loading: "正在加载模型…", st_ready: "选择布局并点击放置发射机。",
-    s_eval_k: "评估方式", s_eval_v: "按城市划分训练/测试集，模型在从未见过的建筑布局上接受评估。",
-    s_model_k: "模型", s_model_v: "U-Net 卷积网络，约 1300 万参数，图像到图像回归，MSE 损失 + Sigmoid 输出。",
-    s_train_k: "训练", s_train_v: "PyTorch，单张 NVIDIA RTX 2080，混合精度（AMP），Adam 优化器。",
-    s_deploy_k: "部署", s_deploy_v: "导出为 ONNX，通过 ONNX Runtime Web 完全在浏览器端运行，无需服务器。",
-    stack_h: "技术栈",
-    footer: "数据集：RadioMapSeer（Yapar 等，2022），CC BY 4.0 许可。独立作品集项目。",
-    st_loading: "正在加载模型…（仅首次下载，约 50 MB）",
-    st_ready: "模型已就绪。选择布局并点击放置发射机。",
     st_running: "正在推理…",
-    st_done: (ms) => "推理完成，用时 " + ms + " 毫秒。再次点击可移动发射机。",
+    st_done: (ms) => "推理完成,用时 " + ms + " 毫秒。再次点击可移动发射机。",
-    st_error: (m) => "错误：" + m
+    st_error: (m) => "错误:" + m
   }
 };
 
-let currentLang = (navigator.language || "en").toLowerCase().startsWith("zh") ? "zh" : "en";
+const VERSIONS = {
+  v1: { model: "radio_unet_v1.onnx", channels: 2, available: true },
+  v2: { model: "radio_unet_v2.onnx", channels: 3, available: true },
+  v3: { model: "radio_unet_v3.onnx", channels: 3, available: true },
+  v4: { model: "radio_unet_v4.onnx", channels: 3, available: true },
+};
+const METRICS = { v1: "0.052", v2: "0.034", v3: "0.030", v4: "0.021" };
 
-function applyLang(lang){
-  currentLang = lang;
-  document.documentElement.lang = (lang === "zh") ? "zh" : "en";
-  document.querySelectorAll("[data-i18n]").forEach(el=>{
-    const v = T[lang][el.getAttribute("data-i18n")];
-    if (typeof v === "string") el.textContent = v;
-  });
-  document.querySelectorAll(".lang").forEach(b=>
-    b.classList.toggle("active", b.getAttribute("data-lang") === lang));
-  // refresh the status line in the new language
-  if (!session) statusEl.textContent = T[lang].st_loading;
-  else if (!txRC) statusEl.textContent = T[lang].st_ready;
-}
-document.querySelectorAll(".lang").forEach(b=>
-  b.addEventListener("click", ()=> applyLang(b.getAttribute("data-lang"))));
-
-// ---------- demo logic ----------
-const MAP_IDS = [600, 615, 630, 645, 660, 675];
 const SIZE = 256;
-let session = null, buildings = null, txRC = null;
+let currentLang = (navigator.language || "en").toLowerCase().startsWith("zh") ? "zh" : "en";
+let activeVersion = "v4";
+const sessions = {};
+let buildings = null, txRC = null;
 
 const statusEl = document.getElementById("status");
 const inCanvas = document.getElementById("inputCanvas");
@@ -215,22 +170,37 @@ function viridis(t){
   return [a[0]+(b[0]-a[0])*f, a[1]+(b[1]-a[1])*f, a[2]+(b[2]-a[2])*f];
 }
 
-async function init(){
+function applyLang(lang){
-  applyLang(currentLang);
+  currentLang = lang;
-  MAP_IDS.forEach((id,i)=>{
+  document.documentElement.lang = (lang === "zh") ? "zh" : "en";
-    const o = document.createElement("option");
+  document.querySelectorAll("[data-i18n]").forEach(el=>{
-    o.value = id; o.textContent = (currentLang==="zh"?"布局 ":"Layout ")+(i+1);
+    const v = T[lang][el.getAttribute("data-i18n")];
-    mapSel.appendChild(o);
+    if (typeof v === "string") el.textContent = v;
   });
-  mapSel.addEventListener("change", ()=> loadMap(mapSel.value));
+  document.querySelectorAll(".lang").forEach(b=> b.classList.toggle("active", b.dataset.lang === lang));
+  updateVersionInfo();
+}
 
-  ort.env.wasm.wasmPaths = "https://cdn.jsdelivr.net/npm/onnxruntime-web/dist/";
+function updateVersionInfo(){
-  ort.env.wasm.numThreads = 1;
+  document.getElementById("verNotice").textContent = T[currentLang]["notice_" + activeVersion];
+  document.getElementById("verMetric").textContent = METRICS[activeVersion];
+  document.querySelectorAll(".ver-btn").forEach(b=> b.classList.toggle("active", b.dataset.v === activeVersion));
+}
 
+async function loadSession(v){
+  if (sessions[v]) return sessions[v];
   statusEl.textContent = T[currentLang].st_loading;
-  session = await ort.InferenceSession.create("radio_unet.onnx", { executionProviders: ["wasm"] });
+  sessions[v] = await ort.InferenceSession.create(VERSIONS[v].model, { executionProviders: ["wasm"] });
-  statusEl.textContent = T[currentLang].st_ready;
+  return sessions[v];
-  await loadMap(MAP_IDS[0]);
+}
+
+async function setVersion(v){
+  if (!VERSIONS[v].available) return;
+  activeVersion = v;
+  updateVersionInfo();
+  await loadSession(v);
+  if (txRC) await runPrediction();
+  else statusEl.textContent = T[currentLang].st_ready;
 }
 
 function loadMap(id){
@@ -249,35 +219,35 @@ function loadMap(id){
   });
 }
 
-inCanvas.addEventListener("click", async (e)=>{
+function buildInput(){
-  if(!session || !buildings) return;
+  const ch = VERSIONS[activeVersion].channels;
-  const rect = inCanvas.getBoundingClientRect();
+  const input = new Float32Array(ch*SIZE*SIZE);
-  const col = Math.floor((e.clientX-rect.left) / rect.width  * SIZE);
+  input.set(buildings, 0);                         // channel 0: buildings
-  const row = Math.floor((e.clientY-rect.top ) / rect.height * SIZE);
-  txRC = [row, col];
-  await predict();
-  drawTxMarker();
-});
-
-async function predict(){
-  statusEl.textContent = T[currentLang].st_running;
-  const input = new Float32Array(2*SIZE*SIZE);
-  input.set(buildings, 0);
   const [r,c] = txRC;
-  input[SIZE*SIZE + r*SIZE + c] = 1.0;
+  input[SIZE*SIZE + r*SIZE + c] = 1.0;             // channel 1: transmitter
-  const tensor = new ort.Tensor("float32", input, [1,2,SIZE,SIZE]);
+  if (ch === 3){                                   // channel 2: distance to Tx (v2/v3/v4)
+    const off = 2*SIZE*SIZE, norm = Math.SQRT2 * SIZE;
+    for(let yy=0; yy<SIZE; yy++)
+      for(let xx=0; xx<SIZE; xx++)
+        input[off + yy*SIZE + xx] = Math.sqrt((yy-r)*(yy-r)+(xx-c)*(xx-c)) / norm;
+  }
+  return new ort.Tensor("float32", input, [1, ch, SIZE, SIZE]);
+}
 
+async function runPrediction(){
+  statusEl.textContent = T[currentLang].st_running;
+  const session = await loadSession(activeVersion);
   const t0 = performance.now();
-  const out = await session.run({ input: tensor });
+  const out = await session.run({ input: buildInput() });
   const dt = performance.now()-t0;
   const cov = out.coverage.data;
-
   const imgData = outCtx.createImageData(SIZE,SIZE);
   for(let p=0; p<SIZE*SIZE; p++){
     const [rr,gg,bb] = viridis(cov[p]);
     imgData.data[p*4]=rr; imgData.data[p*4+1]=gg; imgData.data[p*4+2]=bb; imgData.data[p*4+3]=255;
   }
   outCtx.putImageData(imgData, 0, 0);
+  drawTxMarker();
   statusEl.textContent = T[currentLang].st_done(dt.toFixed(0));
 }
 
@@ -289,6 +259,33 @@ function drawTxMarker(){
   outCtx.lineWidth = 1; outCtx.strokeStyle = "#fff"; outCtx.stroke();
 }
 
+inCanvas.addEventListener("click", async (e)=>{
+  if(!buildings) return;
+  const rect = inCanvas.getBoundingClientRect();
+  txRC = [Math.floor((e.clientY-rect.top)/rect.height*SIZE),
+          Math.floor((e.clientX-rect.left)/rect.width*SIZE)];
+  await runPrediction();
+});
+
+async function init(){
+  [600,615,630,645,660,675].forEach((id,i)=>{
+    const o = document.createElement("option");
+    o.value = id; o.textContent = (currentLang==="zh"?"布局 ":"Layout ")+(i+1);
+    mapSel.appendChild(o);
+  });
+  mapSel.addEventListener("change", ()=> loadMap(mapSel.value));
+  document.querySelectorAll(".lang").forEach(b=> b.addEventListener("click", ()=> applyLang(b.dataset.lang)));
+  document.querySelectorAll(".ver-btn").forEach(b=> b.addEventListener("click", ()=> setVersion(b.dataset.v)));
+
+  ort.env.wasm.wasmPaths = "https://cdn.jsdelivr.net/npm/onnxruntime-web/dist/";
+  ort.env.wasm.numThreads = 1;
+
+  applyLang(currentLang);
+  await loadMap(600);
+  await loadSession(activeVersion);
+  statusEl.textContent = T[currentLang].st_ready;
+}
+
 init().catch(err=>{ statusEl.textContent = T[currentLang].st_error(err.message); console.error(err); });
 </script>
 </body>