如何从3D NumPy/Torch数组中提取连通分量及对应bounding boxes

方案1：基于Scipy实现（适配NumPy数组） #

scipy.ndimage.label原生支持任意维度的二值数组连通分量标记，3D场景可以直接使用，无需额外适配。

代码示例 #

import numpy as np
from scipy.ndimage import label, generate_binary_structure

# ---------- 2D示例验证（和你给出的输入匹配）----------
seg_mask_2d = np.array([
    [1, 0, 0, 0, 0],
    [0, 1, 0, 0, 0],
    [1, 1, 1, 0, 0],
    [1, 1, 0, 1, 0],
    [1, 1, 0, 0, 1]
], dtype=np.int32)

# 2D连通性设置：第二个参数取1为4连通，取2为8连通，这里用8连通和示例输出匹配
struct_2d = generate_binary_structure(2, 2)
connected_components_2d, num_components_2d = label(seg_mask_2d, structure=struct_2d)
print(connected_components_2d)
# 输出和你给出的示例逻辑一致，仅标签编号顺序可能不同

# ---------- 3D场景直接使用 ----------
# 假设你的3D掩码形状为 (D, H, W)，dtype为int/bool
seg_mask_3d = np.random.randint(0, 2, size=(32, 128, 128), dtype=np.int32)

# 3D连通性设置：第二个参数取1为6连通，取2为26连通，按需选择
struct_3d = generate_binary_structure(3, 2)
connected_components_3d, num_components_3d = label(seg_mask_3d, structure=struct_3d)

从连通分量提取边界框 #

得到标记好的连通分量数组后，遍历每个标签即可得到对应边界框：

def get_3d_bounding_boxes(label_array):
    bboxes = []
    # 跳过背景标签0
    for label_id in range(1, label_array.max() + 1):
        # 提取当前标签的所有坐标
        z, y, x = np.where(label_array == label_id)
        # 生成边界框格式 [z_min, y_min, x_min, z_max, y_max, x_max]
        bbox = (z.min(), y.min(), x.min(), z.max(), y.max(), x.max())
        bboxes.append(bbox)
    return bboxes

bboxes_3d = get_3d_bounding_boxes(connected_components_3d)

方案2：基于Torchvision实现（适配PyTorch张量） #

如果你的掩码是PyTorch张量，不需要转NumPy，可以直接用torchvision.ops.connected_components实现，支持GPU加速，适合训练流程中直接调用。

代码示例 #

import torch
from torchvision.ops import connected_components

# 3D掩码张量，形状为 (D, H, W)，bool/int类型，设备可以是cuda
seg_mask_3d_torch = torch.randint(0, 2, size=(32, 128, 128), dtype=torch.int32, device='cuda')

# 连通分量标记，3D场景下connectivity参数取1为6连通，取3为26连通
connected_components_3d_torch = connected_components(seg_mask_3d_torch.unsqueeze(0).unsqueeze(0), connectivity=3).squeeze()
num_components_3d_torch = connected_components_3d_torch.max().item()

注意事项 #

• OpenCV的cv2.connectedComponents仅支持2D数组，3D场景需要自行逐层处理，效率较低，不推荐使用。
• 连通性参数按需选择：3D场景下6连通仅判断上下左右前后6个相邻像素，26连通会判断三维空间内所有相邻的像素。

内容的提问来源于stack exchange，提问作者Rylan Schaeffer