| """Collator emitting the AC-3 frozen batch tensor contract. |
| |
| Consumes per-triple samples in the public `csi_path_*` schema and batches |
| them into the AC-3 contract: |
| scene_idx [Q] int64 |
| tx_scene_idx [Q] int64 |
| rx_xyz_norm [Q, 3] float32 |
| tx_xyz_norm [Q, 3] float32 |
| scene_extent [Q, 3] float32 |
| gt_num_paths [Q] int64 |
| gt_delay_ns [Q, M_max] float32 |
| gt_peak_db [Q, M_max] float32 |
| gt_path_mask [Q, M_max] uint8 |
| gt_truncated [Q] int64 |
| |
| Zero-path triples pass through (mask=0; num_paths=0). |
| """ |
|
|
| from __future__ import annotations |
|
|
| from typing import Any, Sequence |
|
|
| import numpy as np |
| import torch |
|
|
|
|
| def _to_float(x: Any) -> torch.Tensor: |
| return torch.as_tensor(x, dtype=torch.float32) |
|
|
|
|
| def triple_collate(samples: Sequence[dict[str, Any]]) -> dict[str, Any]: |
| if not samples: |
| raise ValueError("triple_collate received empty sample list") |
|
|
| Q = len(samples) |
| |
| M_max = max(int(np.asarray(s["csi_path_valid"]).shape[0]) for s in samples) |
|
|
| scene_seen: dict[Any, int] = {} |
| tx_seen: dict[tuple[Any, Any], int] = {} |
| scene_voxel_levels: list[Any] = [] |
|
|
| scene_idx = torch.zeros((Q,), dtype=torch.long) |
| tx_scene_idx = torch.zeros((Q,), dtype=torch.long) |
| tx_scene_voxel_idx = torch.zeros((Q,), dtype=torch.long) |
| tx_xyz_norm = torch.zeros((Q, 3), dtype=torch.float32) |
| rx_xyz_norm = torch.zeros((Q, 3), dtype=torch.float32) |
| scene_extent = torch.zeros((Q, 3), dtype=torch.float32) |
| gt_num_paths = torch.zeros((Q,), dtype=torch.long) |
| gt_delay_ns = torch.zeros((Q, M_max), dtype=torch.float32) |
| gt_peak_db = torch.zeros((Q, M_max), dtype=torch.float32) |
| gt_path_mask = torch.zeros((Q, M_max), dtype=torch.uint8) |
| gt_truncated = torch.zeros((Q,), dtype=torch.long) |
|
|
| for q, s in enumerate(samples): |
| sid = s["scene_id"] |
| if sid not in scene_seen: |
| scene_seen[sid] = len(scene_seen) |
| scene_voxel_levels.append(s.get("voxel_level")) |
| scene_idx[q] = scene_seen[sid] |
| tx_scene_voxel_idx[q] = scene_idx[q] |
|
|
| key = (sid, s["tx_id"]) |
| if key not in tx_seen: |
| tx_seen[key] = len(tx_seen) |
| tx_scene_idx[q] = tx_seen[key] |
|
|
| tx_xyz_norm[q] = _to_float(s["tx_xyz_norm"]) |
| rx_xyz_norm[q] = _to_float(s["rx_xyz_norm"]) |
| scene_extent[q] = _to_float(s["scene_extent"]) |
|
|
| count = int(s["csi_path_count"]) |
| gt_num_paths[q] = count |
| gt_truncated[q] = int(s.get("csi_path_truncation_count", 0)) |
|
|
| m = int(np.asarray(s["csi_path_valid"]).shape[0]) |
| gt_delay_ns[q, :m] = _to_float(s["csi_path_delay_ns"]) |
| gt_peak_db[q, :m] = _to_float(s["csi_path_peak_db"]) |
| gt_path_mask[q, :m] = torch.as_tensor(s["csi_path_valid"], dtype=torch.uint8) |
|
|
| return { |
| "scene_idx": scene_idx, |
| "tx_scene_idx": tx_scene_idx, |
| "tx_scene_voxel_idx": tx_scene_voxel_idx, |
| "tx_xyz_norm": tx_xyz_norm, |
| "rx_xyz_norm": rx_xyz_norm, |
| "scene_extent": scene_extent, |
| "gt_num_paths": gt_num_paths, |
| "gt_delay_ns": gt_delay_ns, |
| "gt_peak_db": gt_peak_db, |
| "gt_path_mask": gt_path_mask, |
| "gt_truncated": gt_truncated, |
| "scene_voxel_levels": scene_voxel_levels, |
| "unique_scene_count": len(scene_seen), |
| "unique_scene_tx_count": len(tx_seen), |
| } |
|
|
|
|
| __all__ = ["triple_collate"] |
|
|
