# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import io import os import torch import numpy as np from PIL import Image # open3d is only needed by the PLY-save / statistical-filter helpers below, NOT by # the pure-numpy ``visualize_point_cloud`` proj renderer. Importing it eagerly pulls # ``open3d.ml`` -> sklearn -> pandas, which crashes on numpy-ABI-mismatched containers # ("numpy.dtype size changed"). Guard it so the proj render path (infer_diffusion / # infer_forecast) works everywhere; the o3d-using helpers fail only if actually called. try: import open3d as o3d except Exception: # noqa: BLE001 o3d = None # Configure environment for better Kaleido performance os.environ["KALEIDO_SCOPE_TIMEOUT"] = "60" # Increase timeout to 60 seconds # eye: camera position # center: view direction CAMERA_VIEWS = { "front_view_1": { "eye": {"x": -0.3, "y": 0, "z": 0.2}, "center": {"x": 0.1, "y": 0, "z": 0}, }, "paper_view_down": { "eye": {"x": -0.3, "y": 0, "z": 0.26}, "center": {"x": 0.1, "y": 0, "z": -0.03}, }, "paper_view_down_strong": { "eye": {"x": -0.3, "y": 0, "z": 0.3}, "center": {"x": 0.1, "y": 0, "z": -0.08}, }, "paper_side_view": { "eye": {"x": 0.0, "y": -0.12, "z": 0.04}, "center": {"x": 0.0, "y": 0.0, "z": 0.0}, }, "front_view_2": { "eye": {"x": -0.1, "y": 0, "z": 0.05}, "center": {"x": 0.1, "y": 0, "z": 0}, }, "front_view_3": { "eye": {"x": -0.05, "y": 0, "z": 0.05}, "center": {"x": 0.1, "y": 0, "z": 0}, }, "front_view_4": { "eye": {"x": -0.01, "y": 0, "z": 0.01}, "center": {"x": 0.1, "y": 0, "z": 0}, }, "front_view_5": { "eye": {"x": -0.12, "y": 0, "z": 0.05}, "center": {"x": 0.1, "y": 0, "z": 0}, }, "front_view_6": { "eye": {"x": -0.075, "y": 0, "z": 0.05}, "center": {"x": 0.075, "y": 0, "z": 0.0}, }, "top_down_view_1": { "eye": {"x": 0, "y": -0.05, "z": 0.5}, "center": {"x": 0, "y": -0.05, "z": 0}, }, } VIZ_KWARGS = { "point_size": 0.3, "camera_position": CAMERA_VIEWS["front_view_1"], "width": 1280, "height": 720, "bgcolor": [0.0, 0.0, 0.0], "range": 200, } def visualize_point_cloud( point_cloud, colors, point_size, camera_position, width, height, range=80, opacity=1.0, bgcolor=(0, 0, 0), fov_deg=60.0, antialias_point_size=False, principal_y_offset=0.0, depth_sort=False, depth_sort_bin=0.0, point_opacity=None, ): """ Ultra-fast 2D projection visualization (no rendering overhead). """ # Plotly camera system: eye, center, up are in normalized domain coordinates # Scale eye and center to match the point cloud's coordinate system point_cloud_center = np.array([0, 0, 0]) # Lidar origin is always at (0,0,0) scale = range # Scale eye and center from normalized coordinates to match point cloud scale # eye: camera position (scaled to point cloud coordinates) # center: point the camera looks at (scaled to point cloud coordinates) # up: up direction on the page (normalized vector, default [0, 0, 1]) eye = np.array([camera_position["eye"]["x"], camera_position["eye"]["y"], camera_position["eye"]["z"]]) * scale center = np.array([camera_position["center"]["x"], camera_position["center"]["y"], camera_position["center"]["z"]]) * scale # Get up vector from camera_position if provided, otherwise use default [0, 0, 1] # Up is a direction vector (normalized), not scaled by range if "up" in camera_position: up_world = np.array([camera_position["up"]["x"], camera_position["up"]["y"], camera_position["up"]["z"]]) else: up_world = np.array([0, 0, 1]) # Plotly default # Calculate view direction (from eye to center) forward = center - eye forward = forward / np.linalg.norm(forward) # Calculate right vector (perpendicular to forward and up_world) right = np.cross(forward, up_world) right_norm = np.linalg.norm(right) # Handle edge case when forward is parallel to up_world if right_norm < 1e-6: # Use a different reference vector up_world = np.array([0, 1, 0]) right = np.cross(forward, up_world) right_norm = np.linalg.norm(right) right = right / right_norm # Calculate camera up vector (perpendicular to forward and right) # This ensures up is in the plane defined by forward and up_world up = np.cross(right, forward) up = up / np.linalg.norm(up) # Project points to 2D using a perspective projection (Plotly uses perspective) # Camera coordinate system: right (x), up (y), forward (z) points_relative = point_cloud - eye # Points relative to camera position x_cam = np.dot(points_relative, right) y_cam = np.dot(points_relative, up) z_cam = np.dot(points_relative, forward) if point_opacity is None: point_opacity = np.ones(len(point_cloud), dtype=np.float32) else: point_opacity = np.broadcast_to( np.asarray(point_opacity, dtype=np.float32).reshape(-1), (len(point_cloud),), ) # Keep only points in front of the camera eps = 1e-6 depth_mask = z_cam > eps x_cam = x_cam[depth_mask] y_cam = y_cam[depth_mask] z_cam = z_cam[depth_mask] colors = colors[depth_mask] point_opacity = point_opacity[depth_mask] # Perspective projection # Use fov_deg as vertical field of view, zoom scales the FOV fov_rad = np.deg2rad(fov_deg) tan_half_fov = np.tan(fov_rad / 2.0) aspect = width / height x_norm = (x_cam / (z_cam * tan_half_fov * aspect)) * 0.5 + 0.5 y_norm = ( (y_cam / (z_cam * tan_half_fov)) * 0.5 + 0.5 + float(principal_y_offset) ) # Create image img = np.zeros((height, width, 3), dtype=np.uint8) img[:] = np.array(bgcolor) * 255 # Convert to pixel coordinates x_pix = (x_norm * width).astype(int) y_pix = ((1 - y_norm) * height).astype(int) # Flip Y # Filter valid pixels valid = (x_pix >= 0) & (x_pix < width) & (y_pix >= 0) & (y_pix < height) x_pix = x_pix[valid] y_pix = y_pix[valid] # reshape(-1, 3) keeps an empty/1-D color array broadcastable into the (0, 3) # pixel slice (e.g. when every point is filtered out -> nothing to draw). colors_valid = (colors[valid] * 255).astype(np.uint8).reshape(-1, 3) opacity_valid = np.clip(point_opacity[valid], 0.0, 1.0) if depth_sort: # Painter's algorithm: far points first, near points last. Stable sorting # preserves caller ordering for equal-depth temporal composites. sort_depth = z_cam[valid] if depth_sort_bin > 0.0: sort_depth = np.floor(sort_depth / float(depth_sort_bin)) order = np.argsort(-sort_depth, kind="stable") x_pix = x_pix[order] y_pix = y_pix[order] colors_valid = colors_valid[order] opacity_valid = opacity_valid[order] # Draw points. The legacy path quantizes the requested radius to an integer. # The opt-in antialiased path blends the boundary pixels by their fractional # disk coverage, then restores every projected center at full opacity. if antialias_point_size: radius = max(0.0, float(point_size)) extent = int(np.ceil(radius + 0.5)) for dy in np.arange(-extent, extent + 1): for dx in np.arange(-extent, extent + 1): if dx == 0 and dy == 0: continue coverage = float(np.clip(radius + 0.5 - np.hypot(dx, dy), 0.0, 1.0)) if coverage == 0.0: continue shifted_x = x_pix + dx shifted_y = y_pix + dy shifted_valid = ( (shifted_x >= 0) & (shifted_x < width) & (shifted_y >= 0) & (shifted_y < height) ) target_x = shifted_x[shifted_valid] target_y = shifted_y[shifted_valid] foreground = colors_valid[shifted_valid].astype(np.float32) background = img[target_y, target_x].astype(np.float32) alpha = coverage * opacity_valid[shifted_valid, None] img[target_y, target_x] = np.rint( alpha * foreground + (1.0 - alpha) * background ).astype(np.uint8) foreground = colors_valid.astype(np.float32) background = img[y_pix, x_pix].astype(np.float32) alpha = opacity_valid[:, None] img[y_pix, x_pix] = np.rint( alpha * foreground + (1.0 - alpha) * background ).astype(np.uint8) else: radius_px = max(0, int(round(float(point_size)))) if radius_px == 0: foreground = colors_valid.astype(np.float32) background = img[y_pix, x_pix].astype(np.float32) alpha = opacity_valid[:, None] img[y_pix, x_pix] = np.rint( alpha * foreground + (1.0 - alpha) * background ).astype(np.uint8) else: for dy in np.arange(-radius_px, radius_px + 1): for dx in np.arange(-radius_px, radius_px + 1): if dx * dx + dy * dy > radius_px * radius_px: continue shifted_x = x_pix + dx shifted_y = y_pix + dy shifted_valid = ( (shifted_x >= 0) & (shifted_x < width) & (shifted_y >= 0) & (shifted_y < height) ) target = ( shifted_y[shifted_valid], shifted_x[shifted_valid], ) foreground = colors_valid[shifted_valid].astype(np.float32) background = img[target].astype(np.float32) alpha = opacity_valid[shifted_valid, None] img[target] = np.rint( alpha * foreground + (1.0 - alpha) * background ).astype(np.uint8) return img def vis_point_cloud_to_numpy(point_cloud, colors, point_size=None): kwargs = dict(VIZ_KWARGS) if point_size is not None: kwargs["point_size"] = point_size img = visualize_point_cloud(point_cloud, colors, **kwargs) return img def vis_point_cloud(point_cloud, colors, save_path): img = visualize_point_cloud(point_cloud, colors, **VIZ_KWARGS) img = Image.fromarray(img) img.save(save_path) def o3d_statistical_filter(pts, nb_neighbors=20, std_ratio=1.0, return_ind=False): pcd = o3d.geometry.PointCloud() pcd.points = o3d.utility.Vector3dVector(pts) cl, ind = pcd.remove_statistical_outlier(nb_neighbors=nb_neighbors, std_ratio=std_ratio) inlier_cloud = pcd.select_by_index(ind) # get the numpy point cloud inlier_pts = np.asarray(inlier_cloud.points) if return_ind: # ``ind`` maps inlier rows back to the input order so callers can keep # per-point attributes (e.g. intensity) aligned with the kept points. return inlier_pts, np.asarray(ind, dtype=np.int64) return inlier_pts def save_range_map_to_ply( range_map_input, save_path ): """ Save the range map to a point cloud Input: range_map_input: numpy array of shape (H, W), unnormalized range values """ from utils.lidar_rangemap import range_map_to_point_cloud # get the ray directions and ray directions input_points = range_map_to_point_cloud(range_map_input) n_pts = input_points.shape[0] input_colors = np.array([[1, 0.706, 0]] * n_pts, dtype=np.float32) pcd_ref = o3d.geometry.PointCloud() pcd_ref.points = o3d.utility.Vector3dVector(input_points) pcd_ref.colors = o3d.utility.Vector3dVector(input_colors) o3d.io.write_point_cloud(save_path, pcd_ref) def render_range_map_to_point_cloud( range_map_input, range_map_recon, filter_outlier=False, save_path=None, point_size=None, ray_directions=None, filter_nb_neighbors=20, filter_std_ratio=1.0, ): """ Render the range map to a point cloud. Input: range_map_input: numpy array of shape (N, H, W), unnormalized range values range_map_recon: numpy array of shape (N, H, W), unnormalized range values ray_directions: optional (H, W, 3) unit ray-direction grid. When given it overrides the default Pandar128 geometry (e.g. the Waymo sensor's measured elevation/azimuth tables); must match the range-map H, W. Output: point_cloud: numpy array of shape (N, 3, H, W) """ # unnormalize the range map valid_mask_input = range_map_input > 0 valid_mask_recon = range_map_recon > 0 # get the ray directions (default Pandar128 unless an override is supplied) if ray_directions is None: from utils.lidar_rangemap import ( load_pandar128_elevations, range_map_to_ray_directions, ) elevation_angles = load_pandar128_elevations() assert len(elevation_angles) == range_map_input.shape[1] ray_directions = range_map_to_ray_directions(range_map_input.shape[2], elevation_angles) # shape: (H, W, 3) else: ray_directions = np.asarray(ray_directions, dtype=np.float32) assert ray_directions.shape[:2] == range_map_input.shape[1:3], ( f"ray_directions {ray_directions.shape} does not match range map " f"{range_map_input.shape}" ) ray_origins = np.zeros_like(ray_directions) # shape: (H, W, 3) # get the point cloud bs = range_map_input.shape[0] vis_list = [] for batch_idx in range(bs): c_range_map_input = range_map_input[batch_idx] # shape: (H, W) c_range_map_recon = range_map_recon[batch_idx] # shape: (H, W) c_ray_directions = ray_directions # shape: (H, W, 3) c_ray_origins = ray_origins # shape: (H, W, 3) c_valid_mask_input = valid_mask_input[batch_idx] c_valid_mask_recon = valid_mask_recon[batch_idx] input_points = ( c_ray_origins[c_valid_mask_input] + c_ray_directions[c_valid_mask_input] * c_range_map_input[c_valid_mask_input][:, np.newaxis] ) # shape: (N, 3) recon_points = ( c_ray_origins[c_valid_mask_recon] + c_ray_directions[c_valid_mask_recon] * c_range_map_recon[c_valid_mask_recon][:, np.newaxis] ) # shape: (N, 3) if filter_outlier: recon_points = o3d_statistical_filter( recon_points, nb_neighbors=filter_nb_neighbors, std_ratio=filter_std_ratio, ) n_pts = input_points.shape[0] input_colors = np.array([[1, 0.706, 0]] * n_pts, dtype=np.float32) n_pts = recon_points.shape[0] recon_colors = np.array([[0, 0.651, 0.929]] * n_pts, dtype=np.float32) if save_path is not None: os.makedirs(save_path, exist_ok=True) o3d.io.write_point_cloud( f"{save_path}/input_pcd_{batch_idx}.ply", o3d.geometry.PointCloud( points=o3d.utility.Vector3dVector(input_points) ) ) o3d.io.write_point_cloud( f"{save_path}/recon_pcd_{batch_idx}.ply", o3d.geometry.PointCloud( points=o3d.utility.Vector3dVector(recon_points) ) ) input_img = vis_point_cloud_to_numpy(input_points, input_colors, point_size=point_size) # shape: (H, W, 3) recon_img = vis_point_cloud_to_numpy(recon_points, recon_colors, point_size=point_size) # shape: (H, W, 3) recon_images = np.concatenate([input_img, recon_img], axis=1) # shape: (H, 2W, 3) vis_list.append(recon_images) vis_list = np.array(vis_list) # shape: (bs, H, 2W, 3) return vis_list def render_range_map_to_point_cloud_with_intensity( range_map_input, range_map_recon, intensity_input, intensity_recon, filter_outlier=False, ray_directions=None, filter_nb_neighbors=20, filter_std_ratio=1.0, ): """ Render point clouds colored by intensity (viridis colormap). Input: range_map_input/recon: numpy array of shape (N, H, W), unnormalized range intensity_input/recon: numpy array of shape (N, H, W), values in [0, 1] ray_directions: optional (H, W, 3) unit ray-direction grid. When given it overrides the default Pandar128 geometry (e.g. the Waymo sensor's measured elevation/azimuth tables); must match the range-map H, W. Output: vis_list: numpy array of shape (N, H, 2W, 3) """ from matplotlib import cm cmap = cm.get_cmap("viridis") valid_mask_input = range_map_input > 0 valid_mask_recon = range_map_recon > 0 if ray_directions is None: from utils.lidar_rangemap import ( load_pandar128_elevations, range_map_to_ray_directions, ) elevation_angles = load_pandar128_elevations() assert len(elevation_angles) == range_map_input.shape[1] ray_directions = range_map_to_ray_directions(range_map_input.shape[2], elevation_angles) else: ray_directions = np.asarray(ray_directions, dtype=np.float32) assert ray_directions.shape[:2] == range_map_input.shape[1:3], ( f"ray_directions {ray_directions.shape} does not match range map " f"{range_map_input.shape}" ) bs = range_map_input.shape[0] vis_list = [] for b in range(bs): mask_in = valid_mask_input[b] mask_re = valid_mask_recon[b] pts_in = ray_directions[mask_in] * range_map_input[b][mask_in][:, np.newaxis] pts_re = ray_directions[mask_re] * range_map_recon[b][mask_re][:, np.newaxis] if filter_outlier: pts_re = o3d_statistical_filter( pts_re, nb_neighbors=filter_nb_neighbors, std_ratio=filter_std_ratio, ) colors_in = cmap(np.clip(intensity_input[b][mask_in], 0, 1))[:, :3].astype(np.float32) colors_re = cmap(np.clip(intensity_recon[b][mask_re], 0, 1))[:, :3].astype(np.float32) img_in = vis_point_cloud_to_numpy(pts_in, colors_in) img_re = vis_point_cloud_to_numpy(pts_re, colors_re) vis_list.append(np.concatenate([img_in, img_re], axis=1)) return np.array(vis_list) def range_map_to_structured_point_cloud(range_map: torch.Tensor): """ Convert the range map to a structured point cloud. Input: range_map: torch tensor of shape (N, H, W), unnormalized range values Output: point_cloud: torch tensor of shape (N, 3, H, W) """ from utils.lidar_rangemap import ( load_pandar128_elevations, range_map_to_ray_directions, ) elevation_angles = load_pandar128_elevations() assert len(elevation_angles) == range_map.shape[1] valid_mask = range_map > 0 ray_directions = range_map_to_ray_directions(range_map.shape[2], elevation_angles) # shape: (H, W, 3) ray_origins = np.zeros_like(ray_directions) # shape: (H, W, 3) points = range_map[..., np.newaxis] * ray_directions # shape: (H, W, 3) points = points[range_map > 0] return points