commit 0749daa2f987276f5addbeb0ae9e29f086e8a9ad Author: youfu Date: Sun Jun 7 20:54:13 2026 +0800 v1: U-Net radio coverage prediction (2-channel input, baseline) diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..fc79158 --- /dev/null +++ b/.gitignore @@ -0,0 +1,11 @@ +data/ +checkpoints/ +outputs/ +*.zip +__pycache__/ +*.pyc +.ipynb_checkpoints/ +.DS_Store +*.onnx +checkpoints_diag/ +checkpoints_v2/ \ No newline at end of file diff --git a/src/dataset.py b/src/dataset.py new file mode 100644 index 0000000..067ff6a --- /dev/null +++ b/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))) \ No newline at end of file diff --git a/src/evaluate.py b/src/evaluate.py new file mode 100644 index 0000000..be4253b --- /dev/null +++ b/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() \ No newline at end of file diff --git a/src/export_onnx.py b/src/export_onnx.py new file mode 100644 index 0000000..bf836b0 --- /dev/null +++ b/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") \ No newline at end of file diff --git a/src/inspect_features.py b/src/inspect_features.py new file mode 100644 index 0000000..cda439e --- /dev/null +++ b/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") \ No newline at end of file diff --git a/src/model.py b/src/model.py new file mode 100644 index 0000000..9f62312 --- /dev/null +++ b/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") \ No newline at end of file diff --git a/src/plot_curves.py b/src/plot_curves.py new file mode 100644 index 0000000..63118ca --- /dev/null +++ b/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") \ No newline at end of file diff --git a/src/test.py b/src/test.py new file mode 100644 index 0000000..2dab41e --- /dev/null +++ b/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)) \ No newline at end of file diff --git a/src/train.py b/src/train.py new file mode 100644 index 0000000..15f9d94 --- /dev/null +++ b/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() \ No newline at end of file diff --git a/src/visualize.py b/src/visualize.py new file mode 100644 index 0000000..05f42b3 --- /dev/null +++ b/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() \ No newline at end of file diff --git a/web/index.html b/web/index.html new file mode 100644 index 0000000..3e5eb4e --- /dev/null +++ b/web/index.html @@ -0,0 +1,295 @@ + + + + + +Radio Coverage Prediction — Deep Learning Demo + + + + +
+
+ + +
+ +

+

+ +
+
4.1 dB
+
71%
+
~13M
+
<2 s
+
+ +

+

+ +

+

+ +

+

+
+
+ + + +
+
+
+

+ +
+
+

+ +
+
+
+
+
+
+ +

+

+ +

+
+
+
+
+
+
+
+ +

+
+ PythonPyTorchU-Net / CNN + ONNXONNX Runtime WebNumPy +
+ +
+
+ + + + \ No newline at end of file diff --git a/web/maps/map_600.png b/web/maps/map_600.png new file mode 100644 index 0000000..186f948 Binary files /dev/null and b/web/maps/map_600.png differ diff --git a/web/maps/map_615.png b/web/maps/map_615.png new file mode 100644 index 0000000..d0e32d2 Binary files /dev/null and b/web/maps/map_615.png differ diff --git a/web/maps/map_630.png b/web/maps/map_630.png new file mode 100644 index 0000000..cf71ab8 Binary files /dev/null and b/web/maps/map_630.png differ diff --git a/web/maps/map_645.png b/web/maps/map_645.png new file mode 100644 index 0000000..35ab660 Binary files /dev/null and b/web/maps/map_645.png differ diff --git a/web/maps/map_660.png b/web/maps/map_660.png new file mode 100644 index 0000000..4ddbb3d Binary files /dev/null and b/web/maps/map_660.png differ diff --git a/web/maps/map_675.png b/web/maps/map_675.png new file mode 100644 index 0000000..f287d35 Binary files /dev/null and b/web/maps/map_675.png differ