mayavi 显示bbox 和 朝向

1. 代码

import numpy as np
import mayavi.mlab
import pickle
import time

def plot3Dbox(corners):
    for i in range(corners.shape[0]):
        corner = corners[i]
        idx = np.array([0, 1, 2, 3, 0, 4, 5, 6, 7, 4, 5, 1, 2, 6, 7, 3])
        x = corner[0, idx]
        y = corner[1, idx]
        z = corner[2, idx]
        mayavi.mlab.plot3d(x, y, z, color=(1, 0, 0), colormap='Spectral', representation='wireframe',line_width=2)

def boxes3d_to_corners3d(boxes3d, rotate=True):
    """
    :param boxes3d: (N, 7) [x, y, z, h, w, l, ry]
    :param rotate:
    :return: corners3d: (N, 8, 3)
    """
    boxes_num = boxes3d.shape[0]
    h, w, l = boxes3d[:, 3], boxes3d[:, 4], boxes3d[:, 5]
    x_corners = np.array([l / 2., l / 2., -l / 2., -l / 2., l / 2., l / 2., -l / 2., -l / 2.], dtype=np.float32).T  # (N, 8)
    z_corners = np.array([w / 2., -w / 2., -w / 2., w / 2., w / 2., -w / 2., -w / 2., w / 2.], dtype=np.float32).T  # (N, 8)

    y_corners = np.zeros((boxes_num, 8), dtype=np.float32)
    y_corners[:, 4:8] = -h.reshape(boxes_num, 1).repeat(4, axis=1)  # (N, 8)

    if rotate:
        ry = boxes3d[:, 6]
        zeros, ones = np.zeros(ry.size, dtype=np.float32), np.ones(ry.size, dtype=np.float32)
        rot_list = np.array([[np.cos(ry), zeros, -np.sin(ry)],
                             [zeros,       ones,       zeros],
                             [np.sin(ry), zeros,  np.cos(ry)]])  # (3, 3, N)
        R_list = np.transpose(rot_list, (2, 0, 1))  # (N, 3, 3)

        temp_corners = np.concatenate((x_corners.reshape(-1, 8, 1), y_corners.reshape(-1, 8, 1),
                                       z_corners.reshape(-1, 8, 1)), axis=2)  # (N, 8, 3)
        rotated_corners = np.matmul(temp_corners, R_list)  # (N, 8, 3)
        x_corners, y_corners, z_corners = rotated_corners[:, :, 0], rotated_corners[:, :, 1], rotated_corners[:, :, 2]

    x_loc, y_loc, z_loc = boxes3d[:, 0], boxes3d[:, 1], boxes3d[:, 2]

    x = x_loc.reshape(-1, 1) + x_corners.reshape(-1, 8)
    y = y_loc.reshape(-1, 1) + y_corners.reshape(-1, 8)
    z = z_loc.reshape(-1, 1) + z_corners.reshape(-1, 8)

    corners = np.concatenate((x.reshape(-1, 8, 1), y.reshape(-1, 8, 1), z.reshape(-1, 8, 1)), axis=2)

    return corners.astype(np.float32)


def show_points(gt_data):

    for num in range(0,len(gt_data)):

        print(num)
        pointcloud = gt_data[num]['pts_raw']
        print(pointcloud.shape)
        print(type(pointcloud))

        x = pointcloud[:, 0]  # x position of point
        xmin = np.amin(x, axis=0)
        xmax = np.amax(x, axis=0)

        y = pointcloud[:, 1]  # y position of point
        ymin = np.amin(y, axis=0)
        ymax = np.amax(y, axis=0)

        z = pointcloud[:, 2]  # z position of point
        zmin = np.amin(z, axis=0)
        zmax = np.amax(z, axis=0)

        print(xmin,xmax,ymin,ymax,zmin,zmax)

        d = np.sqrt(x ** 2 + y ** 2)  # Map Distance from sensor

        vals = 'height'

        if vals == "height":
            col = z
        else:
            col = d

        fig = mayavi.mlab.figure(bgcolor=(0, 0, 0), size=(640, 500))
        
        mayavi.mlab.points3d(x, y, z,
                             mode="point",
                             # 灰度图的伪彩映射
                             colormap='spectral',  # 'bone', 'copper', 'gnuplot'
                             color=(1, 1, 1),   # Used a fixed (r,g,b) instead
                             figure=fig,
                             )
        bbox = gt_data[num]["box"]
        if num == 0:
            bbox[6] += np.pi

        print(bbox)

        bbox = bbox[np.newaxis, :]
        corners_3d = boxes3d_to_corners3d(bbox)
        print(corners_3d)
        print(corners_3d.shape)

        _x = (corners_3d[0][0][0] + corners_3d[0][1][0]) / 2
        _y = (corners_3d[0][0][1] + corners_3d[0][4][1]) / 2
        _z =  corners_3d[0][1][2]

        _x = (corners_3d[0][0][0] + corners_3d[0][1][0]) / 2 - bbox[0][0]
        _y = (corners_3d[0][0][1] + corners_3d[0][4][1]) / 2 - bbox[0][1] + (corners_3d[0][0][1] + corners_3d[0][4][1]) / 2 / 2
        _z =  corners_3d[0][1][2] - bbox[0][2]

        print(_x,_y,_z)
        mayavi.mlab.quiver3d(bbox[0][0], bbox[0][1] - (corners_3d[0][0][1] + corners_3d[0][4][1]) / 2 / 2 , bbox[0][2],
                             _x, _y, _z, line_width=6, colormap="copper", mode='arrow', scale_factor=3) # quiver3d
        mayavi.mlab.points3d(bbox[0][0], bbox[0][1] - (corners_3d[0][0][1] + corners_3d[0][4][1]) / 2 / 2 , bbox[0][2], color=(0, 1, 0), mode="sphere", scale_factor=1.5)
        # mayavi.mlab.points3d((corners_3d[0][0][0] + corners_3d[0][1][0]) / 2, (corners_3d[0][0][1] + corners_3d[0][4][1]) / 2, corners_3d[0][1][2], color=(0, 1, 0), mode="sphere", scale_factor=1)

        # draw 3d bbox
        corners_3d = np.array(corners_3d).transpose(0, 2, 1)
        plot3Dbox(np.array(corners_3d))

        mayavi.mlab.view(270.0, 160.0, 50,focalpoint=[0, 0, 10])
        mayavi.mlab.show(stop=True)
        # mayavi.mlab.savefig(filename= '/home/seivl/picture/arrow/' + str(num) + '.png',figure=fig) # for save
        mayavi.mlab.close(all=True)
        time.sleep(1)

    print("done")

def main():

    f2 = open('/home/seivl/result.pkl', 'rb')
    gt_data = pickle.load(f2)

    show_points(gt_data)

if __name__ == '__main__':
    main()

2. 结果