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

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()