SUNRGB 数据可视化

1 代码

import numpy as np
import pickle
import open3d

f2 = open('/home/seivl/data/val_results_2.pkl','rb')
data2 = pickle.load(f2)

def rotz(t):
    """Rotation about the z-axis."""
    c = np.cos(t)
    s = np.sin(t)
    return np.array([[c, -s,  0],
                     [s,  c,  0],
                     [0,  0,  1]])

def my_compute_box_3d(center, size, heading_angle):
    R = rotz(heading_angle)
    l,w,h = size
    l = l/2
    w = w/2
    h = h/2
    x_corners = [-l,l,l,-l,-l,l,l,-l]
    y_corners = [w,w,-w,-w,w,w,-w,-w]
    z_corners = [h,h,h,h,-h,-h,-h,-h]
    corners_3d = np.dot(R, np.vstack([x_corners, y_corners, z_corners]))
    corners_3d[0,:] += center[0]
    corners_3d[1,:] += center[1]
    corners_3d[2,:] += center[2]
    return np.transpose(corners_3d)

def main():

    vis = open3d.Visualizer()
    vis.create_window(window_name="sun")
    vis.get_render_option().point_size = 10

    for temp_i in range(len(data2["gt"])):
   	    show_flag = 1
        if (len(data2["gt"][temp_i]) == len(data2["pred"][temp_i])):
        
         	pointcloud = open3d.PointCloud()
            pointcloud.points = open3d.utility.Vector3dVector(data2["points"][temp_i][:, 0:3]) 
            pointcloud.colors = open3d.utility.Vector3dVector(data2["points"][temp_i][:, 3:6]) 

            for i in range(len(data2["gt"][temp_i])):
                bbox = data2["gt"][temp_i][i]
                corners_3d = my_compute_box_3d(bbox[0:3], bbox[3:6], bbox[6])

                bbox_lines = [[0, 1], [1, 2], [2, 3], [3, 0], [4, 5], [5, 6], [6, 7], [7, 4], [0, 4], [1, 5], [2, 6], [3, 7]]
                colors = [[0, 0, 0] for _ in range(len(bbox_lines))] 

                bbox = open3d.LineSet()
                bbox.lines  = open3d.Vector2iVector(bbox_lines)
                bbox.colors = open3d.Vector3dVector(colors)
                bbox.points = open3d.Vector3dVector(corners_3d)
                vis.add_geometry(bbox)

            for i in range(len(data2["pred"][temp_i])):
                bbox = data2["pred"][temp_i][i]
                corners_3d = my_compute_box_3d(bbox[0:3], bbox[3:6], bbox[6])

                bbox_lines = [[0, 1], [1, 2], [2, 3], [3, 0], [4, 5], [5, 6], [6, 7], [7, 4], [0, 4], [1, 5], [2, 6], [3, 7]]
                colors = [[0, 1, 0] for _ in range(len(bbox_lines))]  

                bbox = open3d.LineSet()
                bbox.lines  = open3d.Vector2iVector(bbox_lines)
                bbox.colors = open3d.Vector3dVector(colors)
                bbox.points = open3d.Vector3dVector(corners_3d)
                vis.add_geometry(bbox)

            vis.add_geometry(pointcloud)

            while show_flag:
                vis.update_geometry()
                vis.poll_events()
                vis.update_renderer()
                time.sleep(10)
                show_flag = 0
                
if __name__ == '__main__':
    main()

2 结果