Spaces:
Running on Zero
Running on Zero
File size: 5,984 Bytes
af8bc9f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 | #!/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()
|