102 lines
3.7 KiB
Python
102 lines
3.7 KiB
Python
"""
|
|
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))) |