Spaces:
Running on Zero
Running on Zero
| # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. | |
| # SPDX-License-Identifier: Apache-2.0 | |
| """Inference-only raw-vector HDMap conditioning for prepared PAI samples.""" | |
| import numpy as np | |
| import torch | |
| import torch.nn.functional as F | |
| def validate_raw_hdmap(sample): | |
| required = ("map_vertices", "map_offsets", "box_valid", "box_lwh", "box_object_to_world") | |
| for key in required: | |
| if key not in sample: | |
| raise ValueError(f"Missing {key}; rerun scripts/data/prepare_pai_bidir.py") | |
| vertices, offsets = sample["map_vertices"], sample["map_offsets"] | |
| if vertices.ndim != 2 or vertices.shape[1:] != (3,) or not np.isfinite(vertices).all(): | |
| raise ValueError("Raw map vertices must be finite (N,3) world points") | |
| if (offsets.ndim != 1 or len(offsets) < 2 or offsets.dtype.kind not in "iu" | |
| or offsets[0] != 0 or offsets[-1] != len(vertices) or (np.diff(offsets) < 2).any()): | |
| raise ValueError("Raw map offsets must delimit nonempty polylines of at least two vertices") | |
| shape = sample["box_valid"].shape | |
| if len(shape) != 2 or shape[0] != len(sample["range_m"]): | |
| raise ValueError("Raw box masks must be (T,M), aligned with the real sweeps") | |
| if sample["box_lwh"].shape != shape + (3,) or sample["box_object_to_world"].shape != shape + (4, 4): | |
| raise ValueError("Raw box sizes/poses do not match the box mask") | |
| active = sample["box_valid"].astype(bool) | |
| if (not np.isfinite(sample["box_object_to_world"][active]).all() | |
| or not np.isfinite(sample["box_lwh"][active]).all() | |
| or (sample["box_lwh"][active] <= 0).any()): | |
| raise ValueError("Active boxes must have finite poses and positive finite sizes") | |
| def _sample_polyline(vertices, spacing): | |
| # Match LiDARGen's raw-MADS densification, including endpoint handling. | |
| vertices = np.asarray(vertices, dtype=np.float64) | |
| points = [] | |
| for start, end in zip(vertices[:-1], vertices[1:]): | |
| distance = np.linalg.norm(end - start) | |
| count = int(np.clip(np.floor(distance / spacing), 2, 10000)) | |
| fractions = np.linspace(0, 1, count)[:-1] | |
| points.append(start + fractions[:, None] * (end - start)) | |
| points.append(vertices[-1:]) | |
| return np.concatenate(points).astype(np.float32) | |
| def _box_points(sample, frame, spacing): | |
| signs = np.array([[1, 1, 1], [1, -1, 1], [-1, -1, 1], [-1, 1, 1], | |
| [1, 1, -1], [1, -1, -1], [-1, -1, -1], [-1, 1, -1]], dtype=np.float64) | |
| wireframe = [0, 1, 2, 3, 0, 4, 5, 6, 7, 4, 5, 1, 2, 6, 7, 3] | |
| points = [] | |
| for index in np.flatnonzero(sample["box_valid"][frame]): | |
| corners = signs * sample["box_lwh"][frame, index] / 2 | |
| pose = sample["box_object_to_world"][frame, index].astype(np.float64) | |
| corners = corners @ pose[:3, :3].T + pose[:3, 3] | |
| points.append(_sample_polyline(corners[wireframe], spacing)) | |
| return np.concatenate(points) if points else np.empty((0, 3), np.float32) | |
| def _dilate_nearest(normalized, kernel_size): | |
| """Widen bbox pixels without averaging ranges or filling empty neighborhoods.""" | |
| if kernel_size < 1 or kernel_size % 2 != 1: | |
| raise ValueError("BBox dilation must be a positive odd integer") | |
| if kernel_size == 1: | |
| return normalized | |
| valid = torch.isfinite(normalized) & (normalized > -1 + 1e-5) | |
| work = normalized[:, None].masked_fill(~valid[:, None], torch.inf) | |
| nearest = -F.max_pool2d(-work, kernel_size, stride=1, padding=kernel_size // 2) | |
| return torch.where(torch.isfinite(nearest), nearest, -1.)[:, 0] | |
| def rasterize_raw_hdmap(sample, *, bbox_dilation=1, seg_interval=0.1): | |
| """Return a (T,H,W) normalized condition and provenance, without GT LiDAR.""" | |
| from diffusion.training.datasets.lidar_dataset.forecast_warp import ( | |
| TorchSpatialPointMap, compose_normalized_range_channels, | |
| ) | |
| from diffusion.training.models.hdmaptolidar_model_rf import _sensor_tables_from_ray_directions | |
| validate_raw_hdmap(sample) | |
| if not np.isfinite(seg_interval) or seg_interval <= 0: | |
| raise ValueError("Raw HDMap segment spacing must be positive and finite") | |
| if bbox_dilation < 1 or bbox_dilation % 2 != 1: | |
| raise ValueError("BBox dilation must be a positive odd integer") | |
| rays = torch.as_tensor(sample["ray_directions"], dtype=torch.float32, device="cuda") | |
| elevations, azimuths, offsets = _sensor_tables_from_ray_directions(rays, allow_approximate=True) | |
| renderer = TorchSpatialPointMap( | |
| ray_dirs=rays, row_elevations_rad=elevations, column_azimuths_rad=azimuths, | |
| row_azimuth_offsets_rad=offsets, min_range=2.5, max_range=100, | |
| min_value=-1, inverse_depth=False, radius_m=100, voxel_size_m=0.2, device="cuda", | |
| ) | |
| vertices, offsets = sample["map_vertices"], sample["map_offsets"] | |
| static_points = np.concatenate([_sample_polyline(vertices[a:b], seg_interval) | |
| for a, b in zip(offsets[:-1], offsets[1:])]) | |
| poses = sample["pose_sensor_start_end"].astype(np.float32) | |
| motions = np.linalg.inv(poses[:, 1]) @ poses[:, 0] | |
| static, static_hit = renderer.render_points_world( | |
| static_points, poses[:, 0], target_pose_sensor_start_end=motions, | |
| ) | |
| boxes, box_hits = [], [] | |
| for t, pose in enumerate(poses): | |
| ranges, hits = renderer.render_points_world( | |
| _box_points(sample, t, seg_interval), pose[None, 0], | |
| target_pose_sensor_start_end=motions[t:t + 1], | |
| ) | |
| boxes.append(ranges[0, 0]) | |
| box_hits.append(hits[0]) | |
| thick_boxes = _dilate_nearest(torch.stack(boxes), bbox_dilation) | |
| condition = compose_normalized_range_channels( | |
| torch.stack((static[0], thick_boxes)), channel_dim=0, min_value=-1, inverse_depth=False, | |
| )[0].cpu().numpy() | |
| stats = dict( | |
| source="prepared_raw_vectors", bbox_dilation_kernel=bbox_dilation, | |
| map_dilation_kernel=1, segment_spacing_m=seg_interval, | |
| rolling_shutter=True, min_range_m=2.5, max_range_m=100, | |
| static_hit_fraction=float(static_hit.float().mean()), | |
| bbox_hit_fraction_before_dilation=float(torch.stack(box_hits).float().mean()), | |
| bbox_hit_fraction_after_dilation=float((thick_boxes > -1 + 1e-5).float().mean()), | |
| original_hit_fraction=float((sample["hdmap"] > -1 + 1e-5).mean()), | |
| condition_hit_fraction=float((condition > -1 + 1e-5).mean()), | |
| ) | |
| return condition, stats | |