Spaces:
Running on Zero
Running on Zero
| #!/usr/bin/env python3 | |
| # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. | |
| # SPDX-License-Identifier: Apache-2.0 | |
| """Download and prepare one official NVIDIA PAI example for the 40k tokenizer.""" | |
| import argparse | |
| from io import BytesIO | |
| import json | |
| from pathlib import Path | |
| import tarfile | |
| import numpy as np | |
| import torch | |
| import torch.nn.functional as F | |
| import yaml | |
| REPO_ROOT = Path(__file__).resolve().parents[2] | |
| DATA_REPO = "nvidia/Cosmos-Transfer-LidarGen-Example" | |
| CLIP = "061dbb50-f392-11ee-8e3c-00044baf7cbc_1712352201101085_1712352221101085" | |
| def load_ranges(path, frame_ids): | |
| clip = path.stem | |
| frames = [] | |
| with tarfile.open(path, "r") as archive: | |
| members = {Path(member.name).name: member for member in archive.getmembers() if member.isfile()} | |
| for frame in frame_ids: | |
| arrays = [] | |
| for field in ("row", "col", "range"): | |
| name = f"{clip}.{frame}.lidar_{field}.npz" | |
| with archive.extractfile(members[name]) as handle: | |
| with np.load(BytesIO(handle.read()), allow_pickle=False) as data: | |
| arrays.append(data["arr_0"]) | |
| row, col, values = arrays | |
| dense = np.zeros((128, 3600), dtype=np.float32) | |
| dense[row.astype(np.int64), col.astype(np.int64)] = values | |
| # Training: scatter-min over triplets, then nearest-exact 1200 -> 1024. | |
| groups = dense.reshape(128, 1200, 3) | |
| reduced = np.where(groups == 0, 1e3, groups).min(axis=-1) | |
| reduced[reduced == 1e3] = 0 | |
| valid = np.isfinite(reduced) & (reduced >= 2.5) & (reduced <= 100) | |
| reduced = np.where(valid, reduced, 0).astype(np.float32) | |
| resized = F.interpolate( | |
| torch.from_numpy(reduced)[None, None], size=(128, 1024), mode="nearest-exact" | |
| )[0, 0].numpy() | |
| frames.append(resized) | |
| return np.stack(frames) | |
| def sensor_rays(): | |
| with (REPO_ROOT / "assets/lidar/lidar_model.yaml").open() as handle: | |
| config = yaml.safe_load(handle) | |
| elevations = config["lidar_models"]["spinning"]["HESAI-Pandar128"]["inclinations_deg"] | |
| phi, theta = np.meshgrid( | |
| np.radians(elevations), np.linspace(np.pi, -np.pi, 1024, endpoint=False), indexing="ij" | |
| ) | |
| return np.stack([ | |
| np.cos(phi) * np.cos(theta), np.cos(phi) * np.sin(theta), np.sin(phi) | |
| ], axis=-1).astype(np.float32) | |
| def save_sample(path, payload, provenance): | |
| path.parent.mkdir(parents=True, exist_ok=True) | |
| with path.with_suffix(".npz.tmp").open("wb") as handle: | |
| np.savez_compressed(handle, **payload) | |
| path.with_suffix(".npz.tmp").replace(path) | |
| path.with_suffix(".json").write_text(json.dumps(provenance, indent=2) + "\n") | |
| print(f"Prepared {path} ({len(payload['range_m'])} frames)", flush=True) | |
| def main(): | |
| parser = argparse.ArgumentParser(description=__doc__) | |
| parser.add_argument("--output-root", type=Path, required=True) | |
| parser.add_argument("--source-root", type=Path, help="Existing lidar_dataset_release directory; skips downloads.") | |
| parser.add_argument("--hdmap-root", type=Path, help="Optional projected HDMap TAR directory.") | |
| parser.add_argument("--revision", default="main", help="Dataset revision (default: main).") | |
| args = parser.parse_args() | |
| torch.set_num_threads(4) | |
| if args.source_root is None: | |
| from huggingface_hub import hf_hub_download | |
| download_root = args.output_root / "source" | |
| for filename in ( | |
| "LICENSE", "README.md", | |
| f"lidar_dataset_release/lidar/{CLIP}.tar", | |
| f"lidar_dataset_release/metadata/{CLIP}.npz", | |
| ): | |
| hf_hub_download(DATA_REPO, filename, repo_type="dataset", revision=args.revision, local_dir=download_root) | |
| source = download_root / "lidar_dataset_release" | |
| else: | |
| source = args.source_root | |
| with np.load(source / "metadata" / f"{CLIP}.npz", allow_pickle=False) as metadata: | |
| frame_ids = metadata["frame_indices"].astype(str) | |
| timestamps = metadata["timestamps_list"].copy() | |
| poses = metadata["pose_list"].copy() | |
| if timestamps.shape != (len(frame_ids), 2) or poses.shape != (len(frame_ids), 2, 4, 4): | |
| raise ValueError("PAI frame IDs, timestamps, and start/end poses must align") | |
| if np.any(np.diff(timestamps[:, 0]) <= 0): | |
| raise ValueError("PAI timestamps must be strictly increasing") | |
| ranges = load_ranges(source / "lidar" / f"{CLIP}.tar", frame_ids) | |
| payload = dict( | |
| range_m=ranges, valid_mask=ranges > 0, ray_directions=sensor_rays(), | |
| dataset=np.asarray("pandar"), sensor_name=np.asarray("NVIDIA PAI"), | |
| sensor_dataset_id=np.int64(0), fps=np.float32(10), | |
| timestamp_us=timestamps[:, 0] - timestamps[0, 0], frame_indices=frame_ids, | |
| clip_name=np.asarray(CLIP), has_prepend=np.bool_(False), | |
| min_range_m=np.float32(2.5), max_range_m=np.float32(100), | |
| ) | |
| provenance = dict( | |
| source_repo=DATA_REPO, clip=CLIP, visualization_yaw_deg=0, | |
| processing="scatter_min: 128x3600 -> 128x1200; nearest-exact -> 128x1024; valid range [2.5,100] m", | |
| rays="Calibrated Pandar128 elevations and uniform final-grid azimuths, matching training", | |
| motion_compensation=False, source_revision=args.revision, | |
| ) | |
| save_sample(args.output_root / "examples/nvidia_pai/sample_000.npz", payload, provenance) | |
| if args.hdmap_root is not None: | |
| hdmap_path = args.hdmap_root / f"{CLIP}.tar" | |
| hdmap = load_ranges(hdmap_path, frame_ids) | |
| condition = np.where(hdmap > 0, 2 * (hdmap - 2.5) / 97.5 - 1, -1).astype(np.float32) | |
| save_sample( | |
| args.output_root / "hdmap_examples/nvidia_pai/sample_000.npz", | |
| {**payload, "hdmap": condition, "pose_sensor_start_end": poses}, | |
| {**provenance, "hdmap_source": str(hdmap_path.resolve())}, | |
| ) | |
| if __name__ == "__main__": | |
| main() | |