如何在Matplotlib 3D散点图中为数据点绘制边界框

Matplotlib原生支持3D边界框绘制，不需要借助外部图片编辑工具，只需要计算目标点集的坐标极值后绘制立方体线框即可。

实现步骤 #

• 提取需要添加边界框的目标散点集合的x、y、z三轴坐标
• 计算三个维度的最大、最小值，确定边界框的空间范围
• 基于极值生成边界框的8个顶点坐标，再绘制立方体的12条边

可运行示例代码 #

import matplotlib.pyplot as plt
import numpy as np

# 生成示例3D散点数据
np.random.seed(42)
x = np.random.randn(50)
y = np.random.randn(50)
z = np.random.randn(50)

# 筛选要加边界框的点集，本示例筛选x>0的所有点作为目标点
target_points = np.column_stack([x,y,z])[x>0]

def draw_3d_bounding_box(ax, points, color='red', linewidth=2):
    # 计算三个维度的极值，如需空隙可手动加减偏移量，例：x_min = points[:,0].min() - 0.1
    x_min, x_max = points[:,0].min(), points[:,0].max()
    y_min, y_max = points[:,1].min(), points[:,1].max()
    z_min, z_max = points[:,2].min(), points[:,2].max()
    
    # 定义立方体8个顶点坐标
    vertices = [
        [x_min, y_min, z_min], [x_max, y_min, z_min],
        [x_max, y_max, z_min], [x_min, y_max, z_min],
        [x_min, y_min, z_max], [x_max, y_min, z_max],
        [x_max, y_max, z_max], [x_min, y_max, z_max]
    ]
    vertices = np.array(vertices)
    
    # 绘制立方体12条边
    edges = [
        [0,1], [1,2], [2,3], [3,0], # 底面
        [4,5], [5,6], [6,7], [7,4], # 顶面
        [0,4], [1,5], [2,6], [3,7]  # 连接上下底面的竖边
    ]
    for edge in edges:
        ax.plot3D(vertices[edge,0], vertices[edge,1], vertices[edge,2], color=color, linewidth=linewidth)

# 绘制3D散点图
fig = plt.figure()
ax = fig.add_subplot(projection='3d')
ax.scatter(x, y, z, c='blue', label='全体散点')
ax.scatter(target_points[:,0], target_points[:,1], target_points[:,2], c='orange', label='目标散点')

# 绘制边界框
draw_3d_bounding_box(ax, target_points, color='red', linewidth=2)

ax.set_xlabel('X轴')
ax.set_ylabel('Y轴')
ax.set_zlabel('Z轴')
ax.legend()
plt.show()

实现效果 #

可调参数说明 #

• 修改draw_3d_bounding_box函数的color参数调整边界框颜色，支持16进制色值、颜色名等所有Matplotlib兼容的颜色格式
• 调整linewidth参数修改边框粗细，也可以新增linestyle参数设置虚线、点线等边框样式
• 若需要给多个分组分别加框，只需对每个分组的点集单独调用一次绘制函数即可

内容的提问来源于stack exchange，提问作者Colin Warn