utils/visualize/point_cloud.py

# 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