在 open3d 中创建箭头

Create arrows in open3d

我正在使用 Open3D 可视化一些点云。我想添加在特定点开始和结束的箭头。箭头将可视化我正在处理的一些事情。但是,我还没有找到添加这些箭头的简单方法。

我注意到有一个使用箭头创建笛卡尔坐标系的函数。因此,可以向 3D 可视化添加箭头。

import open3d as o3d
# Create cartesian coordinate 
FOR = o3d.geometry.TriangleMesh.create_coordinate_frame(
    size=10, origin=[0,0,0])
# Visualize FOR
o3d.visualization.draw_geometries([FOR])

我因为找不到在 Open3D 中创建箭头的简单方法而感到沮丧,经过一段时间的努力,我想出了一个解决方案。

import open3d as o3d
import numpy as np

def draw_geometries(pcds):
    """
    Draw Geometries
    Args:
        - pcds (): [pcd1,pcd2,...]
    """
    o3d.visualization.draw_geometries(pcds)

def get_o3d_FOR(origin=[0, 0, 0],size=10):
    """ 
    Create a FOR that can be added to the open3d point cloud
    """
    mesh_frame = o3d.geometry.TriangleMesh.create_coordinate_frame(
    size=size)
    mesh_frame.translate(origin)
    return(mesh_frame)

def vector_magnitude(vec):
    """
    Calculates a vector's magnitude.
    Args:
        - vec (): 
    """
    magnitude = np.sqrt(np.sum(vec**2))
    return(magnitude)


def calculate_zy_rotation_for_arrow(vec):
    """
    Calculates the rotations required to go from the vector vec to the 
    z axis vector of the original FOR. The first rotation that is 
    calculated is over the z axis. This will leave the vector vec on the
    XZ plane. Then, the rotation over the y axis. 

    Returns the angles of rotation over axis z and y required to
    get the vector vec into the same orientation as axis z
    of the original FOR

    Args:
        - vec (): 
    """
    # Rotation over z axis of the FOR
    gamma = np.arctan(vec[1]/vec[0])
    Rz = np.array([[np.cos(gamma),-np.sin(gamma),0],
                   [np.sin(gamma),np.cos(gamma),0],
                   [0,0,1]])
    # Rotate vec to calculate next rotation
    vec = Rz.T@vec.reshape(-1,1)
    vec = vec.reshape(-1)
    # Rotation over y axis of the FOR
    beta = np.arctan(vec[0]/vec[2])
    Ry = np.array([[np.cos(beta),0,np.sin(beta)],
                   [0,1,0],
                   [-np.sin(beta),0,np.cos(beta)]])
    return(Rz, Ry)

def create_arrow(scale=10):
    """
    Create an arrow in for Open3D
    """
    cone_height = scale*0.2
    cylinder_height = scale*0.8
    cone_radius = scale/10
    cylinder_radius = scale/20
    mesh_frame = o3d.geometry.TriangleMesh.create_arrow(cone_radius=1,
        cone_height=cone_height,
        cylinder_radius=0.5,
        cylinder_height=cylinder_height)
    return(mesh_frame)

def get_arrow(origin=[0, 0, 0], end=None, vec=None):
    """
    Creates an arrow from an origin point to an end point,
    or create an arrow from a vector vec starting from origin.
    Args:
        - end (): End point. [x,y,z]
        - vec (): Vector. [i,j,k]
    """
    scale = 10
    Ry = Rz = np.eye(3)
    T = np.array([[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]])
    T[:3, -1] = origin
    if end is not None:
        vec = np.array(end) - np.array(origin)
    elif vec is not None:
        vec = np.array(vec)
    if end is not None or vec is not None:
        scale = vector_magnitude(vec)
        Rz, Ry = calculate_zy_rotation_for_arrow(vec)
    mesh = create_arrow(scale)
    # Create the arrow
    mesh.rotate(Ry, center=np.array([0, 0, 0]))
    mesh.rotate(Rz, center=np.array([0, 0, 0]))
    mesh.translate(origin)
    return(mesh)


# Create a Cartesian Frame of Reference
FOR = get_o3d_FOR()
# Create an arrow from point (5,5,5) to point (10,10,10)
# arrow = get_arrow([5,5,5],[10,10,10])

# Create an arrow representing vector vec, starting at (5,5,5)
# arrow = get_arrow([5,5,5],vec=[5,5,5])

# Create an arrow in the same place as the z axis
arrow = get_arrow()

# Draw everything
draw_geometries([FOR,arrow])

如论文3D-RCNN所述,我们可以计算出2个(单位)向量的对齐操作之间的矩阵,如下式所示: 其中 'r' 定义为 p 和 q 的叉积。 并且'[r]x'由一个斜对称矩阵定义,可以找到here.

所以我们可以使用这种方法来创建所需的箭头或圆柱体:只需使用先前的矩阵作为目标的旋转并对齐中心。

demo可以通过以下方式实现(基于open3d 0.9.0):

import numpy as np 
import open3d as o3d 


def get_cross_prod_mat(pVec_Arr):
    # pVec_Arr shape (3)
    qCross_prod_mat = np.array([
        [0, -pVec_Arr[2], pVec_Arr[1]], 
        [pVec_Arr[2], 0, -pVec_Arr[0]],
        [-pVec_Arr[1], pVec_Arr[0], 0],
    ])
    return qCross_prod_mat


def caculate_align_mat(pVec_Arr):
    scale = np.linalg.norm(pVec_Arr)
    pVec_Arr = pVec_Arr/ scale
    # must ensure pVec_Arr is also a unit vec. 
    z_unit_Arr = np.array([0,0,1])
    z_mat = get_cross_prod_mat(z_unit_Arr)

    z_c_vec = np.matmul(z_mat, pVec_Arr)
    z_c_vec_mat = get_cross_prod_mat(z_c_vec)

    if np.dot(z_unit_Arr, pVec_Arr) == -1:
        qTrans_Mat = -np.eye(3, 3)
    elif np.dot(z_unit_Arr, pVec_Arr) == 1:   
        qTrans_Mat = np.eye(3, 3)
    else:
        qTrans_Mat = np.eye(3, 3) + z_c_vec_mat + np.matmul(z_c_vec_mat,
                                                    z_c_vec_mat)/(1 + np.dot(z_unit_Arr, pVec_Arr))

    qTrans_Mat *= scale
    return qTrans_Mat



if __name__ == "__main__":
    z_unit_Arr = np.array([0,0,1])
    begin = [1, 0, 0]
    end = [1.6, 0.4, 0.8]
    vec_Arr = np.array(end) - np.array(begin)
    vec_len = np.linalg.norm(vec_Arr)


    mesh_frame = o3d.geometry.TriangleMesh.create_coordinate_frame(size=0.6, origin=[0, 0, 0])
    

    mesh_arrow = o3d.geometry.TriangleMesh.create_arrow(
        cone_height= 0.2 * vec_len, 
        cone_radius= 0.06 * vec_len, 
        cylinder_height= 0.8 * vec_len,
        cylinder_radius=  0.04 * vec_len
        )
    mesh_arrow.paint_uniform_color([1,0,1])
    mesh_arrow.compute_vertex_normals()

    mesh_sphere_begin = o3d.geometry.TriangleMesh.create_sphere(radius=0.1, resolution= 20)
    mesh_sphere_begin.translate(begin)
    mesh_sphere_begin.paint_uniform_color([0,1,1])
    mesh_sphere_begin.compute_vertex_normals()

    mesh_sphere_end = o3d.geometry.TriangleMesh.create_sphere(radius=0.1, resolution= 20)
    mesh_sphere_end.translate(end)
    mesh_sphere_end.paint_uniform_color([0,1,1])
    mesh_sphere_end.compute_vertex_normals()

    # mesh_arrow,
    o3d.visualization.draw_geometries(
        geometry_list= [mesh_frame, mesh_arrow, mesh_sphere_begin, mesh_sphere_end],
        window_name= "before", width= 800, height= 600
    )

    rot_mat = caculate_align_mat(vec_Arr)
    mesh_arrow.rotate(rot_mat, center = False)
    o3d.visualization.draw_geometries(
        geometry_list= [mesh_frame, mesh_arrow, mesh_sphere_begin, mesh_sphere_end],
        window_name= "after rotate", width= 800, height= 600
    )
    mesh_arrow.translate(np.array(begin)) # 0.5*(np.array(end) - np.array(begin))
    o3d.visualization.draw_geometries(
        geometry_list= [mesh_frame, mesh_arrow, mesh_sphere_begin, mesh_sphere_end],
        window_name= "after translate", width= 800, height= 600
    )