How to draw disconnected lines in mayavi and customize their colors without loops

Question

I want to plot many boxes in a scene. If I use mlab.plot3d, it will be very slow.

Therefore, I draw many line segments and then connect them to form boxes using plotting-many-lines-example. But, I do not know how to customize their colors.

Answer 1

I have found a solution. We can use mlab.quiver3d to draw many line segments at the same time by assigning mode="2ddash". Below is my code to draw many boxes at the same time with the option to specify the color of each box.

def vis_boxes(
        points: np.ndarray,
        boxes: np.ndarray,
        color: np.ndarray = None,
):
    """

    Args:
        points:  (N, 3)
        boxes:  (M, 6) xyzwlh
        color: (M, 4) RGBA

    Returns:

    """
    num_boxes = len(boxes)
    params = boxes.shape[1]
    if params not in (6, 9):
        raise ValueError(f"invalid number of box parameters {params}, only support 6, 9")

    fig = mlab.figure(bgcolor=(1, 1, 1), size=(1000, 750))

    mlab.points3d(points[:, 0], points[:, 1], points[:, 2], color=(1, 0, 0), figure=fig, scale_factor=0.1)

    center = boxes[:, 0:3]  # (M, 3)
    x_length = boxes[:, 3:4]  # (M, 1)
    y_length = boxes[:, 4:5]  # (M, 1)
    z_length = boxes[:, 5:6]  # (M, 1)
    forward = np.concatenate([np.ones_like(x_length), np.zeros_like(x_length), np.zeros_like(x_length)],
                             axis=-1) * x_length / 2  # (M, 3)
    right = np.concatenate([np.zeros_like(y_length), np.ones_like(y_length), np.zeros_like(y_length)],
                           axis=-1) * y_length / 2  # (M, 3)
    up = np.concatenate([np.zeros_like(z_length), np.zeros_like(z_length), np.ones_like(z_length)],
                        axis=-1) * z_length / 2  # (M, 3)

    def get_box_points(center, forward, right, up):
        p1 = center - forward - right - up  # (M, 3)
        p2 = center - forward + right - up
        p3 = center + forward + right - up
        p4 = center + forward - right - up
        p5 = center - forward - right + up
        p6 = center - forward + right + up
        p7 = center + forward + right + up
        p8 = center + forward - right + up
        start_points = np.stack([p1, p2, p3, p4, p5, p6, p7, p8, p1, p2, p3, p4], axis=-2)  # (M, 12, 3)
        end_points = np.stack([p2, p3, p4, p1, p6, p7, p8, p5, p5, p6, p7, p8], axis=-2)  # (M, 12, 3)
        return start_points, end_points

    src_start_points, src_end_points = get_box_points(center, forward, right, up)
    flow = src_end_points - src_start_points  # (M, 12, 3)
    start_points = np.reshape(src_start_points, (-1, 3))  # (M*12, 3)
    flow = flow.reshape(-1, 3)  # (M*12, 3)
    scalars = np.arange(num_boxes)  # (M, )
    scalars = np.repeat(scalars, 12)  # (M*12, )

    if color is None:
        cmap = matplotlib.colormaps["rainbow"]
        colors = cmap(np.linspace(0, 1, num_boxes)) * 255
        colors = colors.astype(np.uint8)

    q3d = mlab.quiver3d(
        start_points[:, 0],
        start_points[:, 1],
        start_points[:, 2],
        flow[:, 0],
        flow[:, 1],
        flow[:, 2],
        figure=fig,
        line_width=2,
        scale_factor=1,
        mode="2ddash",
        scalars=scalars,
    )
    q3d.glyph.color_mode = "color_by_scalar"  # Color by scalar
    q3d.module_manager.scalar_lut_manager.lut.table = colors
    mlab.draw()
    mlab.show()