| |
| """Smoke-test the packaged WiSER example scene. |
| |
| This script intentionally has two levels: |
| 1. Always validate that the example manifests and assets are readable. |
| 2. Optionally load a checkpoint and report its structure. |
| |
| Full neural inference needs the CUDA sparse backend and the paper checkpoint; |
| for release QA, asset loading and checkpoint readability catch most packaging |
| errors without requiring a large GPU run. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import sys |
| from pathlib import Path |
| from typing import Any |
|
|
| ROOT = Path(__file__).resolve().parents[1] |
| sys.path.insert(0, str(ROOT)) |
|
|
| import torch |
|
|
| from wiser.data.csi_path_targets import MergedPathTargetConfig |
| from wiser.data.dataset import MultiSceneTripleDataset |
| from wiser.data.radiomap_dataset import RadiomapDataset |
|
|
|
|
| def _load_json(path: Path) -> dict[str, Any]: |
| return json.loads(path.read_text()) |
|
|
|
|
| def _resolve_record_paths(records: list[dict[str, Any]], manifest_path: Path) -> list[dict[str, Any]]: |
| out: list[dict[str, Any]] = [] |
| base = manifest_path.parent |
| for rec in records: |
| r = dict(rec) |
| tx_path = r.get("tx_path") |
| if tx_path and not Path(tx_path).is_absolute(): |
| r["tx_path"] = str((base / tx_path).resolve()) |
| out.append(r) |
| return out |
|
|
|
|
| def main() -> None: |
| p = argparse.ArgumentParser() |
| p.add_argument("--example-root", default="example") |
| p.add_argument("--checkpoint", default=None) |
| p.add_argument("--out-json", default="outputs/example_summary.json") |
| p.add_argument("--channels", type=int, default=512) |
| p.add_argument("--max-cir-samples", type=int, default=16) |
| args = p.parse_args() |
|
|
| example_root = Path(args.example_root).resolve() |
| manifests = example_root / "manifests" |
| radiomap_manifest_path = manifests / "radiomap_example.json" |
| cir_manifest_path = manifests / "cir_example.json" |
| scene3d_root = example_root / "data" / "scene3d" |
| wireless_root = example_root / "data" / "wireless" |
|
|
| summary: dict[str, Any] = { |
| "example_root": str(example_root), |
| "radiomap_manifest": str(radiomap_manifest_path), |
| "cir_manifest": str(cir_manifest_path), |
| } |
|
|
| rm_manifest = _load_json(radiomap_manifest_path) |
| rm_records = _resolve_record_paths(rm_manifest["val_heldout"], radiomap_manifest_path) |
| rm_ds = RadiomapDataset( |
| rm_records, |
| channels=args.channels, |
| grid_h=int(rm_manifest.get("grid_h", 36)), |
| grid_w=int(rm_manifest.get("grid_w", 36)), |
| db_floor=-300.0, |
| scene3d_root=scene3d_root, |
| dataset_kind=rm_manifest.get("dataset_kind", "sionna_radiomap"), |
| ) |
| rm_sample = rm_ds[0] |
| summary["radiomap"] = { |
| "num_records": len(rm_ds), |
| "scene_id": rm_sample["scene_id"], |
| "grid_shape": list(rm_sample["gt_radiomap_db"].shape), |
| "valid_cells": int(rm_sample["extent_mask"].sum().item()), |
| "voxel_count": int(rm_sample["voxel_feats"].shape[0]), |
| } |
|
|
| cir_manifest = _load_json(cir_manifest_path) |
| cir_records = _resolve_record_paths(cir_manifest["triples"], cir_manifest_path) |
| cir_manifest_for_loader = dict(cir_manifest) |
| cir_manifest_for_loader["triples"] = cir_records[: args.max_cir_samples] |
| cir_ds = MultiSceneTripleDataset( |
| cir_manifest_for_loader, |
| target_config=MergedPathTargetConfig(), |
| voxel_channels=args.channels, |
| wireless_root=wireless_root, |
| scene3d_root=scene3d_root, |
| dataset_tag=cir_manifest.get("dataset_tag", "voxel_original_csi_path_10cm_1e6"), |
| precompute_stats=False, |
| ) |
| cir_sample = cir_ds[0] |
| summary["cir"] = { |
| "num_records_checked": len(cir_ds), |
| "scene_id": cir_sample["scene_id"], |
| "tx_id": int(cir_sample["tx_id"]), |
| "rx_id": int(cir_sample["rx_id"]), |
| "num_paths": int(cir_sample["csi_path_count"]), |
| "voxel_count": int(cir_sample["voxel_level"]["feats"].shape[1]), |
| } |
|
|
| if args.checkpoint: |
| ckpt_path = Path(args.checkpoint).resolve() |
| ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False) |
| state = ckpt.get("model_state_dict", ckpt.get("state_dict", ckpt)) |
| summary["checkpoint"] = { |
| "path": str(ckpt_path), |
| "top_level_keys": sorted(list(ckpt.keys())) if isinstance(ckpt, dict) else [], |
| "num_state_tensors": len(state) if isinstance(state, dict) else None, |
| "phase_name": ckpt.get("phase_name") if isinstance(ckpt, dict) else None, |
| "epoch": ckpt.get("epoch") if isinstance(ckpt, dict) else None, |
| } |
|
|
| out_json = Path(args.out_json) |
| out_json.parent.mkdir(parents=True, exist_ok=True) |
| out_json.write_text(json.dumps(summary, indent=2)) |
| print(json.dumps(summary, indent=2)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
|
|
