matplotlib画3dbox

• 画box3d
• xyz轴比例设置为1:1:1

from time import time
from elasticsearch_dsl import Object
import numpy as np
import matplotlib.path as mpltPath
from matplotlib import pyplot as plt, patches
from mpl_toolkits.mplot3d import Axes3D
from mpl_toolkits import mplot3d as p3d

from scipy.spatial.transform import Rotation as Rotation
from icecream import ic


def get_lidar_box_corners(box: np.ndarray):
    """
    Args:
        box: [x, y, z, l, w, h, yaw]
    Returns:
        corners: 8 * 3
           0...3
          .   .
         .   .
        1...2  

           4...7
          .   .
         .   .
        5...6

    """
    l2, w2, h2 = box[3:6] * 0.5
    corner_in_obj = np.array([
        [l2, w2, h2],
        [-l2, w2, h2],
        [-l2, -w2, h2],
        [l2, -w2, h2],

        [l2, w2, -h2],
        [-l2, w2, -h2],
        [-l2, -w2, -h2],
        [l2, -w2, -h2],

    ]).T

    # Rotation.from_euler('zyx', [cam_yaw, cam_pitch, cam_roll])
    R_obj_to_rig = Rotation.from_euler('zyx', [box[6], 0, 0]).as_matrix()

    T_obj_to_rig = box[0:3].reshape(-1, 1)
    ic(T_obj_to_rig)
    corners = R_obj_to_rig @ corner_in_obj + T_obj_to_rig
    return corners.T




def draw_box3d(ax: Axes3D, corners: np.ndarray):
    assert corners.shape == (8, 3)
    indice = [0, 1, 2, 3, 0, 4, 5, 6, 7, 4, 5, 1, 2, 6, 7, 3]
    new_corners = corners[indice]
    ax.plot(new_corners[:, 0], new_corners[:, 1], new_corners[:, 2])
    for i, point in enumerate(corners):
        ax.text(point[0], point[1], point[2], str(i))

    # ax.set_box_aspect((np.ptp(new_corners[:, 0]), np.ptp(new_corners[:, 1]), np.ptp(new_corners[:, 2])))

    ic(new_corners.shape)

def set_aspect_equal(ax):
    """ 
    Fix the 3D graph to have similar scale on all the axes.
    Call this after you do all the plot3D, but before show
    """
    X = ax.get_xlim3d()
    Y = ax.get_ylim3d()
    Z = ax.get_zlim3d()
    a = [X[1]-X[0],Y[1]-Y[0],Z[1]-Z[0]]
    b = np.amax(a)
    ax.set_xlim3d(X[0]-(b-a[0])/2,X[1]+(b-a[0])/2)
    ax.set_ylim3d(Y[0]-(b-a[1])/2,Y[1]+(b-a[1])/2)
    ax.set_zlim3d(Z[0]-(b-a[2])/2,Z[1]+(b-a[2])/2)
    ax.set_box_aspect(aspect = (1,1,1))

def test_get_lidar_box_corners():
    boxes = np.array([
        [5, 10, 0, 5, 2, 1.8, np.pi / 6],
        [-1, -3, 0, 4, 1.5, 1, np.pi / 2],

    ])

    plt.rcParams["figure.figsize"] = (10,10)

    fig = plt.figure()
    # ax = p3d.Axes3D(fig)
    ax = fig.add_subplot(projection='3d')
    # ax.set_aspect('auto')
    # ax.set_box_aspect(aspect = (1,1,1))


    for i, box in enumerate(boxes):
        ic(i, box)
        corners_in_rig = get_lidar_box_corners(box)
        ic(corners_in_rig)
        
        draw_box3d(ax, corners_in_rig)
    # ax.set_box_aspect(aspect = (1,1,1))
    set_aspect_equal(ax)
    plt.show()

        

if __name__ == "__main__":
    test_get_lidar_box_corners()
    exit(-1)