v1: U-Net radio coverage prediction (2-channel input, baseline)

.gitignore

@@ -0,0 +1,11 @@
+data/
+checkpoints/
+outputs/
+*.zip
+__pycache__/
+*.pyc
+.ipynb_checkpoints/
+.DS_Store
+*.onnx
+checkpoints_diag/
+checkpoints_v2/

src/dataset.py

@@ -0,0 +1,102 @@
+"""
+dataset.py — RadioMapSeer loader for radio coverage prediction.
+
+Input  (2 channels): city building map + transmitter location
+Target (1 channel) : radio coverage / path-loss map ("gain")
+
+Expected layout under data_root:
+    png/buildings_complete/{map}.png       buildings vs background
+    png/antennas/{map}_{tx}.png            transmitter marker
+    gain/DPM/{map}_{tx}.png                coverage heatmap (target)
+
+Naming: map id 0..699, tx id 0..79  ->  "{map}_{tx}.png"
+"""
+
+from pathlib import Path
+
+import numpy as np
+import torch
+from torch.utils.data import Dataset
+from PIL import Image
+
+# dataset.py lives in src/, so the project's data/ folder is one level up.
+DEFAULT_DATA_ROOT = Path(r"D:\dataset\RadioMapSeer")
+
+
+def _load_gray(path: Path, img_size: int | None = None) -> np.ndarray:
+    """Load a PNG as float32 in [0, 1]; optionally resize to img_size."""
+    img = Image.open(path).convert("L")
+    if img_size is not None and img.size != (img_size, img_size):
+        img = img.resize((img_size, img_size), Image.BILINEAR)
+    return np.asarray(img, dtype=np.float32) / 255.0
+
+
+class RadioMapSeerDataset(Dataset):
+    """Returns (input[2,H,W], target[1,H,W]) tensors in [0, 1]."""
+
+    SPLITS = {
+        "train": range(0, 500),
+        "val":   range(500, 600),
+        "test":  range(600, 700),
+    }
+
+    def __init__(self, data_root=DEFAULT_DATA_ROOT, split="train",
+                 sim="DPM", img_size=256, max_samples=None):
+        self.data_root = Path(data_root)
+        self.img_size = img_size
+
+        self.gain_dir  = self.data_root / "gain" / sim
+        self.build_dir = self.data_root / "png" / "buildings_complete"
+        self.ant_dir   = self.data_root / "png" / "antennas"
+
+        for d in (self.gain_dir, self.build_dir, self.ant_dir):
+            if not d.is_dir():
+                raise FileNotFoundError(f"Expected folder not found: {d}")
+
+        if split not in self.SPLITS:
+            raise ValueError(f"split must be train/val/test, got '{split}'")
+        allowed = set(self.SPLITS[split])
+
+        # Build the sample index by scanning the gain folder for "{map}_{tx}.png".
+        self.samples = []
+        for f in sorted(self.gain_dir.glob("*.png")):
+            try:
+                map_id, tx_id = (int(x) for x in f.stem.split("_"))
+            except ValueError:
+                continue
+            if map_id in allowed:
+                self.samples.append((map_id, tx_id))
+
+        if not self.samples:
+            raise RuntimeError(f"No samples found in {self.gain_dir} — check naming.")
+
+        if max_samples is not None:
+            self.samples = self.samples[:max_samples]
+
+    def __len__(self):
+        return len(self.samples)
+
+    def __getitem__(self, idx):
+        map_id, tx_id = self.samples[idx]
+
+        buildings = _load_gray(self.build_dir / f"{map_id}.png", self.img_size)
+        antenna   = _load_gray(self.ant_dir   / f"{map_id}_{tx_id}.png", self.img_size)
+        gain      = _load_gray(self.gain_dir  / f"{map_id}_{tx_id}.png", self.img_size)
+
+        x = np.stack([buildings, antenna], axis=0)   # [2, H, W]
+        y = gain[None, :, :]                          # [1, H, W]
+        return torch.from_numpy(x), torch.from_numpy(y)
+
+
+if __name__ == "__main__":
+    # Self-test: confirms paths, naming, shapes, and value ranges.
+    print("Data root:", DEFAULT_DATA_ROOT)
+    ds = RadioMapSeerDataset(split="train", img_size=256)
+    print(f"Train samples: {len(ds)}")
+
+    x, y = ds[0]
+    print("Input shape :", tuple(x.shape), "| min/max:", float(x.min()), float(x.max()))
+    print("Target shape:", tuple(y.shape), "| min/max:", float(y.min()), float(y.max()))
+
+    for s in ("val", "test"):
+        print(f"{s} samples:", len(RadioMapSeerDataset(split=s, img_size=256)))

src/evaluate.py

@@ -0,0 +1,111 @@
+"""
+evaluate.py — quantitative evaluation on held-out TEST maps (unseen cities).
+
+Reports U-Net RMSE (gray + dB), a log-distance path-loss baseline, and the
+relative improvement. dB conversion: RadioMapSeer scales an ~80 dB pathloss
+range into [0,1], so RMSE_dB = RMSE_gray * 80 (Levie et al.; Yapar et al.).
+"""
+
+import argparse
+import numpy as np
+import torch
+from torch.utils.data import DataLoader
+
+from dataset import RadioMapSeerDataset
+from model import UNet
+
+DB_DYNAMIC_RANGE = 80.0  # dB spanned by the [0,1] gray scale
+
+
+def tx_location(antenna_img):
+    return np.unravel_index(int(np.argmax(antenna_img)), antenna_img.shape)
+
+
+def fit_log_distance(dataset, n_fit=200, seed=0):
+    """Fit gray ~ a + b*log10(distance) by least squares (distance-only baseline)."""
+    rng = np.random.default_rng(seed)
+    idxs = rng.choice(len(dataset), size=min(n_fit, len(dataset)), replace=False)
+    H = W = dataset.img_size
+    yy, xx = np.mgrid[0:H, 0:W]
+    logd_all, gray_all = [], []
+    for i in idxs:
+        x, y = dataset[int(i)]
+        r, c = tx_location(x[1].numpy())
+        dist = np.sqrt((yy - r) ** 2 + (xx - c) ** 2) + 1.0
+        logd_all.append(np.log10(dist).ravel())
+        gray_all.append(y.squeeze().numpy().ravel())
+    logd, gray = np.concatenate(logd_all), np.concatenate(gray_all)
+    A = np.vstack([np.ones_like(logd), logd]).T
+    (a, b), *_ = np.linalg.lstsq(A, gray, rcond=None)
+    return float(a), float(b)
+
+
+def baseline_predict(antenna_img, a, b):
+    H, W = antenna_img.shape
+    yy, xx = np.mgrid[0:H, 0:W]
+    r, c = tx_location(antenna_img)
+    dist = np.sqrt((yy - r) ** 2 + (xx - c) ** 2) + 1.0
+    return np.clip(a + b * np.log10(dist), 0.0, 1.0)
+
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument("--ckpt", default="checkpoints/best.pt")
+    p.add_argument("--img-size", type=int, default=256)
+    p.add_argument("--batch-size", type=int, default=32)
+    p.add_argument("--num-workers", type=int, default=4)
+    p.add_argument("--max-test", type=int, default=None, help="limit test samples for speed")
+    args = p.parse_args()
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    ckpt = torch.load(args.ckpt, map_location=device, weights_only=False)
+    base = ckpt.get("args", {}).get("base", 64)
+    model = UNet(2, 1, base=base).to(device)
+    model.load_state_dict(ckpt["model_state"])
+    model.eval()
+    print(f"Loaded {args.ckpt} (epoch {ckpt.get('epoch', '?')})")
+
+    test_ds = RadioMapSeerDataset(split="test", img_size=args.img_size,
+                                  max_samples=args.max_test)
+    print(f"Test samples: {len(test_ds)}")
+
+    # --- U-Net error over the test set ---
+    loader = DataLoader(test_ds, batch_size=args.batch_size, shuffle=False,
+                        num_workers=args.num_workers, pin_memory=True)
+    sq_err, n_px = 0.0, 0
+    with torch.no_grad():
+        for x, y in loader:
+            x, y = x.to(device), y.to(device)
+            with torch.autocast(device_type=device, dtype=torch.float16,
+                                enabled=(device == "cuda")):
+                pred = model(x)
+            sq_err += torch.sum((pred.float() - y.float()) ** 2).item()
+            n_px += y.numel()
+    unet_rmse = (sq_err / n_px) ** 0.5
+
+    # --- log-distance baseline: fit on train, evaluate on test ---
+    train_ds = RadioMapSeerDataset(split="train", img_size=args.img_size)
+    a, b = fit_log_distance(train_ds, n_fit=200)
+    print(f"Baseline fit: gray = {a:.3f} + ({b:.3f}) * log10(dist)")
+
+    bsq, bn = 0.0, 0
+    for i in range(len(test_ds)):
+        x, y = test_ds[i]
+        pred = baseline_predict(x[1].numpy(), a, b)
+        truth = y.squeeze().numpy()
+        bsq += float(np.sum((pred - truth) ** 2))
+        bn += truth.size
+    base_rmse = (bsq / bn) ** 0.5
+
+    print("\n" + "=" * 54)
+    print("  EVALUATION ON UNSEEN TEST CITIES")
+    print("=" * 54)
+    print(f"  U-Net (deep learning)  RMSE: {unet_rmse:.4f} gray | {unet_rmse*DB_DYNAMIC_RANGE:5.2f} dB")
+    print(f"  Log-distance baseline  RMSE: {base_rmse:.4f} gray | {base_rmse*DB_DYNAMIC_RANGE:5.2f} dB")
+    print(f"  Improvement over baseline:   {100*(base_rmse-unet_rmse)/base_rmse:.1f}%")
+    print("=" * 54)
+
+
+if __name__ == "__main__":
+    main()

src/export_onnx.py

@@ -0,0 +1,45 @@
+"""
+export_onnx.py — export the trained U-Net to ONNX and verify it matches PyTorch.
+Run from the project root:  python src\\export_onnx.py
+"""
+import numpy as np
+import torch
+
+from model import UNet
+from dataset import RadioMapSeerDataset
+
+CKPT = "checkpoints/best.pt"
+OUT = "web/radio_unet.onnx"
+
+# --- load the trained model ---
+ckpt = torch.load(CKPT, map_location="cpu", weights_only=False)
+base = ckpt.get("args", {}).get("base", 64)
+model = UNet(2, 1, base=base)
+model.load_state_dict(ckpt["model_state"])
+model.eval()
+
+# --- export to ONNX (fixed 1x2x256x256 input) ---
+dummy = torch.randn(1, 2, 256, 256)
+torch.onnx.export(
+    model, dummy, OUT,
+    input_names=["input"], output_names=["coverage"],
+    opset_version=17,
+)
+print(f"Exported {OUT}")
+
+# --- verify ONNX output matches PyTorch on a real test sample ---
+import onnxruntime as ort
+
+ds = RadioMapSeerDataset(split="test", img_size=256)
+x, _ = ds[0]
+xb = x.unsqueeze(0).numpy().astype(np.float32)
+
+with torch.no_grad():
+    torch_out = model(x.unsqueeze(0)).numpy()
+
+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)))
+print(f"Max |PyTorch - ONNX| on a test sample: {max_diff:.2e}")
+print("OK — outputs match" if max_diff < 1e-3 else "WARNING: large mismatch, stop and check")

src/inspect_features.py

@@ -0,0 +1,72 @@
+"""
+inspect_features.py — visualize U-Net feature maps at each encoder level for one
+demo case, to check whether the transmitter ("tower") registers through the depth.
+
+Run:  python src\\inspect_features.py
+(or paste these blocks into Jupyter Lab cells to poke at individual channels)
+"""
+import numpy as np
+import torch
+import matplotlib.pyplot as plt
+from pathlib import Path
+
+from dataset import RadioMapSeerDataset
+from model import UNet
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+ckpt = torch.load("checkpoints/best.pt", map_location=device, weights_only=False)
+base = ckpt.get("args", {}).get("base", 64)
+model = UNet(2, 1, base=base).to(device)
+model.load_state_dict(ckpt["model_state"])
+model.eval()
+
+# capture each level's output with forward hooks
+acts = {}
+def hook(name):
+    def fn(m, i, o): acts[name] = o.detach().float().cpu()
+    return fn
+
+model.inc.register_forward_hook(hook("inc 256"))
+model.down1.register_forward_hook(hook("down1 128"))
+model.down2.register_forward_hook(hook("down2 64"))
+model.down3.register_forward_hook(hook("down3 32"))
+model.down4.register_forward_hook(hook("down4 16 (bottleneck)"))
+
+# one demo case from the test set
+ds = RadioMapSeerDataset(split="test", img_size=256)
+x, y = ds[0]
+tx = np.unravel_index(int(np.argmax(x[1].numpy())), (256, 256))  # (row, col)
+with torch.no_grad():
+    pred = model(x.unsqueeze(0).to(device))
+
+print("Transmitter at (row, col):", tx)
+for k, v in acts.items():
+    print(f"{k}: shape {tuple(v.shape)}")
+
+levels = list(acts.keys())
+fig, axes = plt.subplots(2, 5, figsize=(20, 8))
+
+axes[0, 0].imshow(x[0].numpy(), cmap="gray")
+axes[0, 0].scatter([tx[1]], [tx[0]], c="red", marker="x", s=50)
+axes[0, 0].set_title("input: buildings + Tx"); axes[0, 0].axis("off")
+axes[0, 1].imshow(y.squeeze().numpy(), cmap="viridis")
+axes[0, 1].set_title("ground truth"); axes[0, 1].axis("off")
+axes[0, 2].imshow(pred.squeeze().cpu().numpy(), cmap="viridis")
+axes[0, 2].set_title("prediction"); axes[0, 2].axis("off")
+axes[0, 3].axis("off"); axes[0, 4].axis("off")
+
+# bottom row: max activation over channels at each level (best for spotting the Tx)
+for ax, name in zip(axes[1], levels):
+    fmap = acts[name][0].amax(dim=0).numpy()      # [h, w]
+    h = fmap.shape[0]
+    ax.imshow(fmap, cmap="magma")
+    ax.scatter([tx[1] * h / 256], [tx[0] * h / 256],
+               facecolors="none", edgecolors="cyan", s=70, linewidths=1.5)
+    ax.set_title(name, fontsize=10); ax.axis("off")
+
+fig.suptitle(f"U-Net feature maps (max over channels) — Tx at {tx}", fontsize=13)
+fig.tight_layout()
+Path("outputs").mkdir(exist_ok=True)
+fig.savefig("outputs/feature_maps.png", dpi=110, bbox_inches="tight")
+print("saved outputs/feature_maps.png")

src/model.py

@@ -0,0 +1,98 @@
+"""
+model.py — U-Net for radio coverage map prediction.
+
+Input : [B, 2, H, W]  (buildings + transmitter location)
+Output: [B, 1, H, W]  (predicted coverage map, values in [0, 1])
+
+Classic encoder-decoder with skip connections (Ronneberger et al., 2015),
+adapted for image-to-image regression instead of segmentation.
+"""
+
+import torch
+import torch.nn as nn
+
+
+class DoubleConv(nn.Module):
+    """(conv 3x3 -> BatchNorm -> ReLU) applied twice."""
+    def __init__(self, in_ch, out_ch):
+        super().__init__()
+        self.block = nn.Sequential(
+            nn.Conv2d(in_ch, out_ch, kernel_size=3, padding=1, bias=False),
+            nn.BatchNorm2d(out_ch),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(out_ch, out_ch, kernel_size=3, padding=1, bias=False),
+            nn.BatchNorm2d(out_ch),
+            nn.ReLU(inplace=True),
+        )
+
+    def forward(self, x):
+        return self.block(x)
+
+
+class Down(nn.Module):
+    """Halve resolution (max-pool) then DoubleConv."""
+    def __init__(self, in_ch, out_ch):
+        super().__init__()
+        self.pool = nn.MaxPool2d(2)
+        self.conv = DoubleConv(in_ch, out_ch)
+
+    def forward(self, x):
+        return self.conv(self.pool(x))
+
+
+class Up(nn.Module):
+    """Upsample, concatenate the skip connection, then DoubleConv."""
+    def __init__(self, in_ch, skip_ch, out_ch):
+        super().__init__()
+        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
+        return self.conv(x)
+
+
+class UNet(nn.Module):
+    def __init__(self, in_channels=2, out_channels=1, base=64):
+        super().__init__()
+        self.inc   = DoubleConv(in_channels, base)   # 256
+        self.down1 = Down(base,     base * 2)        # 128
+        self.down2 = Down(base * 2, base * 4)        # 64
+        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.up2 = Up(base * 4, base * 4, base * 2)
+        self.up3 = Up(base * 2, base * 2, base)
+        self.up4 = Up(base,     base,     base)
+
+        self.outc = nn.Conv2d(base, out_channels, kernel_size=1)
+
+    def forward(self, x):
+        x1 = self.inc(x)
+        x2 = self.down1(x1)
+        x3 = self.down2(x2)
+        x4 = self.down3(x3)
+        x5 = self.down4(x4)
+
+        x = self.up1(x5, x4)
+        x = self.up2(x, x3)
+        x = self.up3(x, x2)
+        x = self.up4(x, x1)
+        return torch.sigmoid(self.outc(x))
+
+
+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)
+    y = model(x)
+    print("Input :", tuple(x.shape))
+    print("Output:", tuple(y.shape), "| range:", round(float(y.min()), 3), "-", round(float(y.max()), 3))
+    if device == "cuda":
+        print(f"Peak VRAM (forward, batch 4): {torch.cuda.max_memory_allocated()/1e9:.2f} GB")

src/plot_curves.py

@@ -0,0 +1,18 @@
+import csv
+import matplotlib.pyplot as plt
+
+epochs, tr, va = [], [], []
+with open("checkpoints_diag/history.csv") as f:
+    for row in csv.DictReader(f):
+        epochs.append(int(row["epoch"]))
+        tr.append(float(row["train_mse"]))
+        va.append(float(row["val_mse"]))
+
+plt.figure(figsize=(7, 4))
+plt.plot(epochs, tr, label="train MSE", marker="o")
+plt.plot(epochs, va, label="val MSE", marker="s")
+plt.xlabel("epoch"); plt.ylabel("MSE"); plt.legend(); plt.grid(alpha=0.3)
+plt.title("Training vs. validation loss")
+plt.tight_layout()
+plt.savefig("outputs/loss_curves.png", dpi=120)
+print("saved outputs/loss_curves.png")

src/test.py

@@ -0,0 +1,10 @@
+import numpy as np
+from PIL import Image
+
+base = r"D:\Nextcloud-youfu\dev\code\radio-coverage-dl\data\png"
+ant = np.asarray(Image.open(base + r"\antennas\600_0.png").convert("L"))
+bld = np.asarray(Image.open(base + r"\buildings_complete\600.png").convert("L"))
+
+print("天线图里出现的数值:", np.unique(ant))
+print("天线亮点的像素个数:", int((ant > 127).sum()))
+print("建筑图里出现的数值:", np.unique(bld))

src/train.py

@@ -0,0 +1,120 @@
+"""
+train.py — train the U-Net on RadioMapSeer coverage maps.
+
+Examples:
+  # smoke test: tiny subset, 1 epoch — verify the loop and measure VRAM/speed
+  python src\\train.py --max-samples 200 --epochs 1 --batch-size 16
+
+  # full run (size it after the smoke test)
+  python src\\train.py --epochs 30 --batch-size 16
+"""
+
+import argparse
+import time
+from pathlib import Path
+
+import torch
+import torch.nn as nn
+from torch.utils.data import DataLoader
+
+from dataset import RadioMapSeerDataset
+from model import UNet
+
+
+def evaluate(model, loader, criterion, device, amp):
+    model.eval()
+    total, n = 0.0, 0
+    with torch.no_grad():
+        for x, y in loader:
+            x, y = x.to(device), y.to(device)
+            with torch.autocast(device_type=device, dtype=torch.float16, enabled=amp):
+                loss = criterion(model(x), y)
+            total += loss.item() * x.size(0)
+            n += x.size(0)
+    return total / n
+
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument("--epochs", type=int, default=30)
+    p.add_argument("--batch-size", type=int, default=16)
+    p.add_argument("--lr", type=float, default=1e-3)
+    p.add_argument("--img-size", type=int, default=256)
+    p.add_argument("--base", type=int, default=64)
+    p.add_argument("--max-samples", type=int, default=None,
+                   help="Limit samples for quick smoke tests")
+    p.add_argument("--num-workers", type=int, default=4)
+    p.add_argument("--ckpt-dir", default="checkpoints")
+    args = p.parse_args()
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    amp = (device == "cuda")
+    print(f"Device: {device} | AMP: {amp}")
+
+    train_ds = RadioMapSeerDataset(split="train", img_size=args.img_size,
+                                   max_samples=args.max_samples)
+    val_max = None if args.max_samples is None else max(20, args.max_samples // 5)
+    val_ds = RadioMapSeerDataset(split="val", img_size=args.img_size,
+                                 max_samples=val_max)
+    print(f"Train: {len(train_ds)} | Val: {len(val_ds)}")
+
+    dl_kwargs = dict(batch_size=args.batch_size, num_workers=args.num_workers,
+                     pin_memory=True, persistent_workers=(args.num_workers > 0))
+    train_loader = DataLoader(train_ds, shuffle=True, **dl_kwargs)
+    val_loader = DataLoader(val_ds, shuffle=False, **dl_kwargs)
+
+    model = UNet(in_channels=2, out_channels=1, base=args.base).to(device)
+    criterion = nn.MSELoss()
+    optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
+    scaler = torch.amp.GradScaler(enabled=amp)
+
+    ckpt_dir = Path(args.ckpt_dir)
+    ckpt_dir.mkdir(exist_ok=True)
+    best_val = float("inf")
+
+    import csv
+    history_path = ckpt_dir / "history.csv"
+    with open(history_path, "w", newline="") as f:
+        csv.writer(f).writerow(["epoch", "train_mse", "val_mse"])
+
+    for epoch in range(1, args.epochs + 1):
+        model.train()
+        t0 = time.time()
+        running, n = 0.0, 0
+        for i, (x, y) in enumerate(train_loader, 1):
+            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(model(x), y)
+            scaler.scale(loss).backward()
+            scaler.step(optimizer)
+            scaler.update()
+            running += loss.item() * x.size(0)
+            n += x.size(0)
+            if i % 50 == 0:
+                print(f"  epoch {epoch} step {i}/{len(train_loader)} | loss {running/n:.5f}")
+
+        train_loss = running / n
+        val_loss = evaluate(model, val_loader, criterion, device, amp)
+        dt = time.time() - t0
+        print(f"[epoch {epoch}/{args.epochs}] train MSE {train_loss:.5f} | "
+              f"val MSE {val_loss:.5f} | 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}"])
+
+        if val_loss < best_val:
+            best_val = val_loss
+            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})")
+
+        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})")
+
+
+if __name__ == "__main__":
+    main()

src/visualize.py

@@ -0,0 +1,82 @@
+"""
+visualize.py — render model predictions vs ground truth on unseen test maps.
+
+Loads checkpoints/best.pt and saves comparison panels to outputs/.
+
+Usage:
+  python src\\visualize.py
+  python src\\visualize.py --n 6
+"""
+
+import argparse
+from pathlib import Path
+
+import numpy as np
+import torch
+import matplotlib.pyplot as plt
+
+from dataset import RadioMapSeerDataset
+from model import UNet
+
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument("--ckpt", default="checkpoints/best.pt")
+    p.add_argument("--n", type=int, default=6, help="number of samples to render")
+    p.add_argument("--img-size", type=int, default=256)
+    p.add_argument("--out-dir", default="outputs")
+    p.add_argument("--seed", type=int, default=0)
+    args = p.parse_args()
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    ckpt = torch.load(args.ckpt, map_location=device, weights_only=False)
+    base = ckpt.get("args", {}).get("base", 64)
+    model = UNet(in_channels=2, out_channels=1, base=base).to(device)
+    model.load_state_dict(ckpt["model_state"])
+    model.eval()
+    print(f"Loaded {args.ckpt} (epoch {ckpt.get('epoch', '?')}, "
+          f"val MSE {ckpt.get('val_loss', float('nan')):.5f})")
+
+    ds = RadioMapSeerDataset(split="test", img_size=args.img_size)  # unseen cities
+    rng = np.random.default_rng(args.seed)
+    idxs = rng.choice(len(ds), size=args.n, replace=False)
+
+    out_dir = Path(args.out_dir)
+    out_dir.mkdir(exist_ok=True)
+
+    rmses = []
+    for k, idx in enumerate(idxs):
+        x, y = ds[int(idx)]
+        with torch.no_grad():
+            pred = model(x.unsqueeze(0).to(device)).cpu().squeeze().numpy()
+        buildings, antenna = x[0].numpy(), x[1].numpy()
+        truth = y.squeeze().numpy()
+        err = np.abs(pred - truth)
+        rmse = float(np.sqrt(np.mean((pred - truth) ** 2)))
+        rmses.append(rmse)
+
+        fig, axes = plt.subplots(1, 5, figsize=(18, 4))
+        panels = [
+            (buildings, "City map (buildings)", "gray",    1.0),
+            (antenna,   "Transmitter",          "gray",    1.0),
+            (truth,     "Ground-truth coverage","viridis", 1.0),
+            (pred,      "Predicted coverage",   "viridis", 1.0),
+            (err,       "Absolute error",       "magma",   float(err.max())),
+        ]
+        for ax, (img, title, cmap, vmax) in zip(axes, panels):
+            ax.imshow(img, cmap=cmap, vmin=0, vmax=vmax)
+            ax.set_title(title, fontsize=11)
+            ax.axis("off")
+        fig.suptitle(f"Test sample {idx} — RMSE {rmse:.4f}", fontsize=13)
+        fig.tight_layout()
+        out_path = out_dir / f"prediction_{k:02d}.png"
+        fig.savefig(out_path, dpi=110, bbox_inches="tight")
+        plt.close(fig)
+        print(f"  saved {out_path}  (RMSE {rmse:.4f})")
+
+    print(f"\nDone. {args.n} panels in {out_dir}/  | mean RMSE {np.mean(rmses):.4f}")
+
+
+if __name__ == "__main__":
+    main()

web/index.html

@@ -0,0 +1,295 @@
+<!DOCTYPE html>
+<html lang="en">
+<head>
+<meta charset="utf-8">
+<meta name="viewport" content="width=device-width, initial-scale=1">
+<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; }
+  * { 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; }
+  .lang { cursor:pointer; border:1px solid var(--line); background:transparent; color:var(--muted);
+          padding:5px 12px; border-radius:20px; font-size:.85rem; }
+  .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; }
+  h2 { font-size:1.15rem; margin:36px 0 10px; }
+  p { margin:0 0 12px; }
+  .metrics { display:flex; gap:14px; flex-wrap:wrap; margin:18px 0 8px; }
+  .metric { background:var(--panel); border-radius:10px; padding:12px 16px; flex:1 1 150px; }
+  .metric .v { font-size:1.5rem; font-weight:700; color:var(--accent); }
+  .metric .l { color:var(--muted); font-size:.85rem; }
+  .demo { background:var(--panel); border-radius:14px; padding:20px; margin-top:14px; }
+  .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;
+           border-radius:8px; background:#000; cursor:crosshair; display:block; }
+  .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;
+            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; }
+  .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>
+<div class="wrap">
+  <div class="topbar">
+    <button class="lang" data-lang="en">EN</button>
+    <button class="lang" data-lang="zh">中文</button>
+  </div>
+
+  <h1 data-i18n="title"></h1>
+  <p class="sub" data-i18n="subtitle"></p>
+
+  <div class="metrics">
+    <div class="metric"><div class="v">4.1 dB</div><div class="l" data-i18n="m_rmse"></div></div>
+    <div class="metric"><div class="v">71%</div><div class="l" data-i18n="m_improve"></div></div>
+    <div class="metric"><div class="v">~13M</div><div class="l" data-i18n="m_params"></div></div>
+    <div class="metric"><div class="v">&lt;2 s</div><div class="l" data-i18n="m_speed"></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>
+      <select id="mapSel"></select>
+      <span class="hint" data-i18n="click_hint"></span>
+    </div>
+    <div class="row">
+      <div class="col">
+        <h3 data-i18n="input_title"></h3>
+        <canvas id="inputCanvas" width="256" height="256"></canvas>
+      </div>
+      <div class="col">
+        <h3 data-i18n="output_title"></h3>
+        <canvas id="outputCanvas" width="256" height="256"></canvas>
+        <div class="legend"></div>
+        <div class="legend-labels"><span data-i18n="legend_weak"></span><span data-i18n="legend_strong"></span></div>
+      </div>
+    </div>
+    <div class="status" id="status"></div>
+  </div>
+
+  <h2 data-i18n="results_h"></h2>
+  <p data-i18n="results_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>
+
+<script>
+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.",
+    m_rmse: "RMSE on unseen cities", m_improve: "lower error vs. classical model",
+    m_params: "parameters", m_speed: "in-browser inference",
+    problem_h: "The problem",
+    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?",
+    how_h: "How it works",
+    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.",
+    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",
+    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.",
+    hood_h: "Under the hood",
+    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.",
+    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_error: (m) => "Error: " + m
+  },
+  zh: {
+    title: "基于深度学习的无线信号覆盖预测",
+    subtitle: "一个 U-Net 神经网络直接预测基站在城市中的下行信号覆盖——并在你的浏览器中实时运行。",
+    m_rmse: "未见城市上的 RMSE", m_improve: "较经典模型误差降低",
+    m_params: "参数量", m_speed: "浏览器端推理",
+    problem_h: "问题背景",
+    problem_p: "在建设基站之前，移动运营商需要预估信号在城市中的传播情况——哪里覆盖良好，哪里因建筑遮挡形成信号「阴影」。传统方法要么使用快速但粗糙的经验公式（如 Hata 模型），完全忽略建筑几何结构；要么使用物理精确但缓慢的射线追踪仿真，每张地图可能耗时数分钟。本项目要回答的问题是：神经网络能否在毫秒级时间内，达到接近慢速仿真的预测精度？",
+    how_h: "实现原理",
+    how_p: "模型是一个 U-Net——带跳跃连接的卷积编码器-解码器。它以两个图像通道作为输入（城市建筑布局与发射机位置），输出一张图像：预测的信号覆盖热力图，越亮表示信号越强，越暗表示越弱。网络从未被告知任何传播公式，它完全从样本数据中自行学会了信号随距离衰减、以及在建筑后方形成阴影的规律。",
+    demo_h: "在线体验",
+    demo_p: "选择一个城市布局，点击任意位置放置发射机。模型完全在你的浏览器中运行，不到一秒即可渲染出预测覆盖。把发射机放在建筑附近，观察信号阴影是如何形成的。",
+    select_label: "城市布局：", click_hint: "点击地图以放置发射机。",
+    input_title: "城市地图——点击放置发射机", output_title: "预测覆盖",
+    legend_weak: "弱", legend_strong: "强",
+    results_h: "实验结果",
+    results_p: "在训练中从未见过的测试城市上评估，模型达到 4.1 dB 的均方根误差（RMSE）——相比只能建模距离、对建筑「视而不见」的经典对数距离路径损耗模型（14.3 dB），误差降低了 71%。这一差距正是神经网络的价值所在：它学会了经典模型无法表达的建筑遮蔽效应。",
+    hood_h: "技术细节",
+    s_data_k: "数据集", s_data_v: "RadioMapSeer——基于真实城市（柏林、伦敦、特拉维夫等）、5.9 GHz、56,000+ 张射线追踪仿真覆盖图。",
+    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_error: (m) => "错误：" + m
+  }
+};
+
+let currentLang = (navigator.language || "en").toLowerCase().startsWith("zh") ? "zh" : "en";
+
+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;
+
+const statusEl = document.getElementById("status");
+const inCanvas = document.getElementById("inputCanvas");
+const outCanvas = document.getElementById("outputCanvas");
+const inCtx = inCanvas.getContext("2d", { willReadFrequently: true });
+const outCtx = outCanvas.getContext("2d");
+const mapSel = document.getElementById("mapSel");
+
+const VIRIDIS = [[68,1,84],[72,40,120],[62,74,137],[49,104,142],[38,130,142],
+                 [31,158,137],[53,183,121],[109,205,89],[180,222,44],[253,231,37]];
+function viridis(t){
+  t = Math.max(0, Math.min(1, t));
+  const x = t*(VIRIDIS.length-1), i = Math.floor(x), f = x-i;
+  const a = VIRIDIS[i], b = VIRIDIS[Math.min(i+1, VIRIDIS.length-1)];
+  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(){
+  applyLang(currentLang);
+  MAP_IDS.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));
+
+  ort.env.wasm.wasmPaths = "https://cdn.jsdelivr.net/npm/onnxruntime-web/dist/";
+  ort.env.wasm.numThreads = 1;
+
+  statusEl.textContent = T[currentLang].st_loading;
+  session = await ort.InferenceSession.create("radio_unet.onnx", { executionProviders: ["wasm"] });
+  statusEl.textContent = T[currentLang].st_ready;
+  await loadMap(MAP_IDS[0]);
+}
+
+function loadMap(id){
+  return new Promise((resolve)=>{
+    const img = new Image();
+    img.onload = ()=>{
+      inCtx.drawImage(img, 0, 0, SIZE, SIZE);
+      const data = inCtx.getImageData(0,0,SIZE,SIZE).data;
+      buildings = new Float32Array(SIZE*SIZE);
+      for(let p=0; p<SIZE*SIZE; p++) buildings[p] = data[p*4] / 255;
+      outCtx.clearRect(0,0,SIZE,SIZE);
+      txRC = null;
+      resolve();
+    };
+    img.src = "maps/map_"+id+".png";
+  });
+}
+
+inCanvas.addEventListener("click", async (e)=>{
+  if(!session || !buildings) return;
+  const rect = inCanvas.getBoundingClientRect();
+  const col = Math.floor((e.clientX-rect.left) / rect.width  * SIZE);
+  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;
+  const tensor = new ort.Tensor("float32", input, [1,2,SIZE,SIZE]);
+
+  const t0 = performance.now();
+  const out = await session.run({ input: tensor });
+  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);
+  statusEl.textContent = T[currentLang].st_done(dt.toFixed(0));
+}
+
+function drawTxMarker(){
+  if(!txRC) return;
+  const [r,c] = txRC;
+  outCtx.beginPath(); outCtx.arc(c, r, 3, 0, 2*Math.PI);
+  outCtx.fillStyle = "#ff3b3b"; outCtx.fill();
+  outCtx.lineWidth = 1; outCtx.strokeStyle = "#fff"; outCtx.stroke();
+}
+
+init().catch(err=>{ statusEl.textContent = T[currentLang].st_error(err.message); console.error(err); });
+</script>
+</body>
+</html>