| """ |
| Visualization Utils |
| |
| Author: Xiaoyang Wu (xiaoyang.wu.cs@gmail.com) |
| Please cite our work if the code is helpful to you. |
| """ |
|
|
| import os |
|
|
| try: |
| import open3d as o3d |
| except ImportError: |
| o3d = None |
| import numpy as np |
| import torch |
|
|
|
|
| def to_numpy(x): |
| if isinstance(x, torch.Tensor): |
| x = x.clone().detach().cpu().numpy() |
| assert isinstance(x, np.ndarray) |
| return x |
|
|
|
|
| def get_point_cloud(coord, color=None, verbose=True): |
| if not isinstance(coord, list): |
| coord = [coord] |
| if color is not None: |
| color = [color] |
|
|
| pcd_list = [] |
| for i in range(len(coord)): |
| coord_ = to_numpy(coord[i]) |
| if color is not None: |
| color_ = to_numpy(color[i]) |
| pcd = o3d.geometry.PointCloud() |
| pcd.points = o3d.utility.Vector3dVector(coord_) |
| pcd.colors = o3d.utility.Vector3dVector( |
| np.zeros_like(coord_) if color is None else color_ |
| ) |
| pcd_list.append(pcd) |
| if verbose: |
| o3d.visualization.draw_geometries(pcd_list) |
| return pcd_list |
|
|
|
|
| def get_line_set(coord, line, color=(1.0, 0.0, 0.0), verbose=True): |
| coord = to_numpy(coord) |
| line = to_numpy(line) |
| colors = np.array([color for _ in range(len(line))]) |
| line_set = o3d.geometry.LineSet() |
| line_set.points = o3d.utility.Vector3dVector(coord) |
| line_set.lines = o3d.utility.Vector2iVector(line) |
| line_set.colors = o3d.utility.Vector3dVector(colors) |
| if verbose: |
| o3d.visualization.draw_geometries([line_set]) |
| return line_set |
|
|
|
|
| def save_point_cloud(coord, color=None, file_path="pc.ply", logger=None): |
| os.makedirs(os.path.dirname(file_path), exist_ok=True) |
| coord = to_numpy(coord) |
| if color is not None: |
| color = to_numpy(color) |
| pcd = o3d.geometry.PointCloud() |
| pcd.points = o3d.utility.Vector3dVector(coord) |
| pcd.colors = o3d.utility.Vector3dVector( |
| np.ones_like(coord) if color is None else color |
| ) |
| o3d.io.write_point_cloud(file_path, pcd) |
| if logger is not None: |
| logger.info(f"Save Point Cloud to: {file_path}") |
|
|
|
|
| def save_bounding_boxes( |
| bboxes_corners, color=(1.0, 0.0, 0.0), file_path="bbox.ply", logger=None |
| ): |
| bboxes_corners = to_numpy(bboxes_corners) |
| |
| points = bboxes_corners.reshape(-1, 3) |
| |
| box_lines = np.array( |
| [ |
| [0, 1], |
| [1, 2], |
| [2, 3], |
| [3, 0], |
| [4, 5], |
| [5, 6], |
| [6, 7], |
| [7, 0], |
| [0, 4], |
| [1, 5], |
| [2, 6], |
| [3, 7], |
| ] |
| ) |
| lines = [] |
| for i, _ in enumerate(bboxes_corners): |
| lines.append(box_lines + i * 8) |
| lines = np.concatenate(lines) |
| |
| color = np.array([color for _ in range(len(lines))]) |
| |
| line_set = o3d.geometry.LineSet() |
| line_set.points = o3d.utility.Vector3dVector(points) |
| line_set.lines = o3d.utility.Vector2iVector(lines) |
| line_set.colors = o3d.utility.Vector3dVector(color) |
| o3d.io.write_line_set(file_path, line_set) |
|
|
| if logger is not None: |
| logger.info(f"Save Boxes to: {file_path}") |
|
|
|
|
| def save_lines( |
| points, lines, color=(1.0, 0.0, 0.0), file_path="lines.ply", logger=None |
| ): |
| points = to_numpy(points) |
| lines = to_numpy(lines) |
| colors = np.array([color for _ in range(len(lines))]) |
| line_set = o3d.geometry.LineSet() |
| line_set.points = o3d.utility.Vector3dVector(points) |
| line_set.lines = o3d.utility.Vector2iVector(lines) |
| line_set.colors = o3d.utility.Vector3dVector(colors) |
| o3d.io.write_line_set(file_path, line_set) |
|
|
| if logger is not None: |
| logger.info(f"Save Lines to: {file_path}") |
|
|
