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