''' Visualization code for point clouds and 3D bounding boxes with mayavi. Modified by Charles R. Qi Date: September 2017 Ref: https://github.com/hengck23/didi-udacity-2017/blob/master/baseline-04/kitti_data/draw.py ''' import numpy as np import mayavi.mlab as mlab try: raw_input # Python 2 except NameError: raw_input = input # Python 3 def draw_lidar_simple(pc, color=None): ''' Draw lidar points. simplest set up. ''' fig = mlab.figure(figure=None, bgcolor=(0,0,0), fgcolor=None, engine=None, size=(1600, 1000)) if color is None: color = pc[:,2] #draw points mlab.points3d(pc[:,0], pc[:,1], pc[:,2], color, color=None, mode='point', colormap = 'gnuplot', scale_factor=1, figure=fig) #draw origin mlab.points3d(0, 0, 0, color=(1,1,1), mode='sphere', scale_factor=0.2) #draw axis axes=np.array([ [2.,0.,0.,0.], [0.,2.,0.,0.], [0.,0.,2.,0.], ],dtype=np.float64) mlab.plot3d([0, axes[0,0]], [0, axes[0,1]], [0, axes[0,2]], color=(1,0,0), tube_radius=None, figure=fig) mlab.plot3d([0, axes[1,0]], [0, axes[1,1]], [0, axes[1,2]], color=(0,1,0), tube_radius=None, figure=fig) mlab.plot3d([0, axes[2,0]], [0, axes[2,1]], [0, axes[2,2]], color=(0,0,1), tube_radius=None, figure=fig) mlab.view(azimuth=180, elevation=70, focalpoint=[ 12.0909996 , -1.04700089, -2.03249991], distance=62.0, figure=fig) return fig def draw_lidar(pc, color=None, fig=None, bgcolor=(0,0,0), pts_scale=1, pts_mode='point', pts_color=None): ''' Draw lidar points Args: pc: numpy array (n,3) of XYZ color: numpy array (n) of intensity or whatever fig: mayavi figure handler, if None create new one otherwise will use it Returns: fig: created or used fig ''' if fig is None: fig = mlab.figure(figure=None, bgcolor=bgcolor, fgcolor=None, engine=None, size=(1600, 1000)) if color is None: color = pc[:,2] mlab.points3d(pc[:,0], pc[:,1], pc[:,2], color, color=pts_color, mode=pts_mode, colormap = 'gnuplot', scale_factor=pts_scale, figure=fig) #draw origin mlab.points3d(0, 0, 0, color=(1,1,1), mode='sphere', scale_factor=0.2) #draw axis axes=np.array([ [2.,0.,0.,0.], [0.,2.,0.,0.], [0.,0.,2.,0.], ],dtype=np.float64) mlab.plot3d([0, axes[0,0]], [0, axes[0,1]], [0, axes[0,2]], color=(1,0,0), tube_radius=None, figure=fig) mlab.plot3d([0, axes[1,0]], [0, axes[1,1]], [0, axes[1,2]], color=(0,1,0), tube_radius=None, figure=fig) mlab.plot3d([0, axes[2,0]], [0, axes[2,1]], [0, axes[2,2]], color=(0,0,1), tube_radius=None, figure=fig) # draw fov (todo: update to real sensor spec.) fov=np.array([ # 45 degree [20., 20., 0.,0.], [20.,-20., 0.,0.], ],dtype=np.float64) mlab.plot3d([0, fov[0,0]], [0, fov[0,1]], [0, fov[0,2]], color=(1,1,1), tube_radius=None, line_width=1, figure=fig) mlab.plot3d([0, fov[1,0]], [0, fov[1,1]], [0, fov[1,2]], color=(1,1,1), tube_radius=None, line_width=1, figure=fig) # draw square region TOP_Y_MIN=-20 TOP_Y_MAX=20 TOP_X_MIN=0 TOP_X_MAX=40 TOP_Z_MIN=-2.0 TOP_Z_MAX=0.4 x1 = TOP_X_MIN x2 = TOP_X_MAX y1 = TOP_Y_MIN y2 = TOP_Y_MAX mlab.plot3d([x1, x1], [y1, y2], [0,0], color=(0.5,0.5,0.5), tube_radius=0.1, line_width=1, figure=fig) mlab.plot3d([x2, x2], [y1, y2], [0,0], color=(0.5,0.5,0.5), tube_radius=0.1, line_width=1, figure=fig) mlab.plot3d([x1, x2], [y1, y1], [0,0], color=(0.5,0.5,0.5), tube_radius=0.1, line_width=1, figure=fig) mlab.plot3d([x1, x2], [y2, y2], [0,0], color=(0.5,0.5,0.5), tube_radius=0.1, line_width=1, figure=fig) #mlab.orientation_axes() mlab.view(azimuth=180, elevation=70, focalpoint=[ 12.0909996 , -1.04700089, -2.03249991], distance=62.0, figure=fig) return fig def draw_gt_boxes3d(gt_boxes3d, fig, color=(1,1,1), line_width=1, draw_text=True, text_scale=(1,1,1), color_list=None): ''' Draw 3D bounding boxes Args: gt_boxes3d: numpy array (n,8,3) for XYZs of the box corners fig: mayavi figure handler color: RGB value tuple in range (0,1), box line color line_width: box line width draw_text: boolean, if true, write box indices beside boxes text_scale: three number tuple color_list: a list of RGB tuple, if not None, overwrite color. Returns: fig: updated fig ''' num = len(gt_boxes3d) for n in range(num): b = gt_boxes3d[n] if color_list is not None: color = color_list[n] if draw_text: mlab.text3d(b[4,0], b[4,1], b[4,2], '%d'%n, scale=text_scale, color=color, figure=fig) for k in range(0,4): #http://docs.enthought.com/mayavi/mayavi/auto/mlab_helper_functions.html i,j=k,(k+1)%4 mlab.plot3d([b[i,0], b[j,0]], [b[i,1], b[j,1]], [b[i,2], b[j,2]], color=color, tube_radius=None, line_width=line_width, figure=fig) i,j=k+4,(k+1)%4 + 4 mlab.plot3d([b[i,0], b[j,0]], [b[i,1], b[j,1]], [b[i,2], b[j,2]], color=color, tube_radius=None, line_width=line_width, figure=fig) i,j=k,k+4 mlab.plot3d([b[i,0], b[j,0]], [b[i,1], b[j,1]], [b[i,2], b[j,2]], color=color, tube_radius=None, line_width=line_width, figure=fig) #mlab.show(1) #mlab.view(azimuth=180, elevation=70, focalpoint=[ 12.0909996 , -1.04700089, -2.03249991], distance=62.0, figure=fig) return fig if __name__=='__main__': pc = np.loadtxt('mayavi/kitti_sample_scan.txt') fig = draw_lidar(pc) mlab.savefig('pc_view.jpg', figure=fig) raw_input()