华为ONCE数据集可视化[python + open3d + pcd + json]

'''

Description : 显示一帧点云(pcd格式)和对应的标签文件(json格式)

'''

'''

Usage: python visualize_once_pcd.py --point_json test.json --point_file test.pcd

Example:

python visualize_once_pcd.py

--point_json ./sample/000000.json

--point_file ./sample/000000.pcd

'''



import numpy as np

import open3d as o3d

import matplotlib

import argparse

import json


box_colormap = [

[1, 1, 1],

[0, 1, 0],

[0, 1, 1],

[1, 1, 0],

]


def get_coor_colors(obj_labels):

colors = matplotlib.colors.XKCD_COLORS.values()

max_color_num = obj_labels.max()

color_list = list(colors)[:max_color_num+1]

colors_rgba = [matplotlib.colors.to_rgba_array(color) for color in color_list]

label_rgba = np.array(colors_rgba)[obj_labels]

label_rgba = label_rgba.squeeze()[:, :3]

return label_rgba


def draw_scenes(points, gt_boxes=None, ref_boxes=None, ref_labels=None, ref_scores=None, point_colors=None, draw_origin=True):

vis = o3d.visualization.Visualizer()#可视化器创建

vis.create_window()#创建可视化窗口

vis.get_render_option().point_size = 1.0 #设定点的大小

vis.get_render_option().background_color = np.zeros(3) #设定背景颜色


if draw_origin: #绘制原始坐标系

axis_pcd = o3d.geometry.TriangleMesh.create_coordinate_frame(size=1.0, origin=[0, 0, 0])

vis.add_geometry(axis_pcd) #坐标系可视化


pts = o3d.geometry.PointCloud() #点云创建

pts.points = o3d.utility.Vector3dVector(points[:, :3]) #点坐标赋值

vis.add_geometry(pts) #点云可视化

if point_colors is None:

pts.colors = o3d.utility.Vector3dVector(np.ones((points.shape[0], 3))) #点颜色赋值 白色

else:

pts.colors = o3d.utility.Vector3dVector(point_colors)


if gt_boxes is not None:

vis = draw_box(vis, gt_boxes, (0, 0, 1)) #GT框


if ref_boxes is not None:

vis = draw_box(vis, ref_boxes, (0, 1, 0), ref_labels, ref_scores) #预测框


vis.run() #运行可视化

vis.destroy_window() #关闭可视化窗口


def translate_boxes_to_open3d_instance(box):

center = np.array([box["position"]["x"], box["position"]["y"], box["position"]["z"]])

lwh = np.array([box["dimensions"]["x"], box["dimensions"]["y"], box["dimensions"]["z"]])

axis_angles = np.array([box["rotation"]["x"], box["rotation"]["y"], box["rotation"]["z"]])

rot = o3d.geometry.get_rotation_matrix_from_axis_angle(axis_angles)

box3d = o3d.geometry.OrientedBoundingBox(center, rot, lwh)

line_set = o3d.geometry.LineSet.create_from_oriented_bounding_box(box3d) #创建线段

lines = np.asarray(line_set.lines) #线段索引

lines = np.concatenate([lines, np.array([[1, 4], [7, 6]])], axis=0)#添加额外线段

line_set.lines = o3d.utility.Vector2iVector(lines) #线段索引赋值

return line_set, box3d


def draw_box(vis, boxes, color=(0, 1, 0), ref_labels=None, score=None):

for i in range(len(boxes)):

line_set, box3d = translate_boxes_to_open3d_instance(boxes[i])

if ref_labels is None:

line_set.paint_uniform_color(color)

else:

colormap = box_colormap[int(ref_labels[i])]

line_set.paint_uniform_color(colormap)

vis.add_geometry(line_set)

return vis


def parse_config():

parser = argparse.ArgumentParser(description='arg parser')

parser.add_argument('--point_json', type=str, help='specify the json file with labels')

parser.add_argument('--point_file', type=str, help='specify the point cloud data file (PCD format)')

args = parser.parse_args()

return args


def main():

args = parse_config()

pcd = o3d.io.read_point_cloud(args.point_file)

points = np.asarray(pcd.points)


with open(args.point_json, 'r') as f:

data = json.load(f)


boxes = []

for figure in data['figures']:

boxes.append(figure['geometry'])


draw_scenes(points=points, ref_boxes=boxes)


if __name__ == '__main__':

main()