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.
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()