Spaces:
Running on Zero
Running on Zero
| # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. | |
| # SPDX-License-Identifier: Apache-2.0 | |
| """Shared processed-PAI batch contract for training and inference.""" | |
| import numpy as np | |
| import torch | |
| def make_pai_batch(sample, length, *, input_frames=None, hdmap=True, boxes=True, device="cpu"): | |
| ranges = sample["range_m"][:length] | |
| valid = sample["valid_mask"][:length] & np.isfinite(ranges) & (ranges >= 2.5) & (ranges <= 100) | |
| video = np.where(valid, 2 * (ranges - 2.5) / 97.5 - 1, -1).astype(np.float32) | |
| if input_frames is not None: | |
| video[input_frames:] = -1 # Never expose future GT during inference. | |
| poses = sample["pose_sensor_start_end"][:length] | |
| tensor = lambda value: torch.as_tensor(value, device=device) | |
| batch = dict( | |
| video=tensor(video)[None, None], is_preprocessed=True, | |
| lidar_hdmap=tensor(sample["hdmap"][:length].astype(np.float32))[None, None] | |
| if hdmap else torch.full((1, 1, length, 128, 1024), -1., device=device), | |
| sensor_dataset_id=tensor([0]).long(), sensor_ray_directions=tensor(sample["ray_directions"])[None].float(), | |
| sensor_depth_bounds=tensor([[2.5, 100.0]]).float(), | |
| lidar_to_world=tensor(poses[:, 0])[None].float(), | |
| pose_sensor_start_end=tensor(np.linalg.inv(poses[:, 1]) @ poses[:, 0])[None].float(), | |
| fps=tensor([10.0]), num_frames=tensor([length]), image_size=tensor([[128, 1024]]), | |
| padding_mask=torch.zeros(1, 1, 128, 1024, device=device), | |
| t5_text_embeddings=torch.zeros(1, 512, 1024, device=device), | |
| t5_text_mask=torch.zeros(1, 512, device=device), | |
| ) | |
| batch["t5_text_mask"][:, 0] = 1 | |
| if boxes: | |
| valid_boxes = sample["box_valid"][:length].astype(bool) | |
| tracks = sample["box_track_id"][:length] | |
| dynamic_ids = np.unique(tracks[valid_boxes & sample["box_is_moving"][:length]]) | |
| dynamic = valid_boxes & np.isin(tracks, dynamic_ids) | |
| for key in ("object_to_world", "lwh", "track_id"): | |
| array = sample[f"box_{key}"][:length] | |
| packed = np.zeros((length, 256) + array.shape[2:], dtype=array.dtype) | |
| if key == "track_id": | |
| packed.fill(-1) | |
| for t in range(length): | |
| selected = array[t, dynamic[t]] | |
| if len(selected) > 256: | |
| raise ValueError("More than 256 dynamic boxes in a sweep") | |
| packed[t, :len(selected)] = selected | |
| batch[f"moving_box_{key}"] = tensor(packed)[None] | |
| batch["moving_box_valid"] = tensor(np.arange(256)[None] < dynamic.sum(axis=1)[:, None])[None] | |
| return batch | |