如何为scipy.ndimage.label生成任意高维连通结构

需求说明 #

你有任意维度的高维布尔数组（比如形状为(3, 2, 66, 241)的4维数组），想要用scipy.ndimage.label标记True值的连通区域，但需要根据指定的维度分组规则生成连通结构——比如：

• dim_connections=[[0],[1],[2,3]]：维度0内部连通，维度1内部连通，维度2和3之间连通，但这三个分组互相不连通
• dim_connections=[[0,1],[2,3]]：维度0、1构成一个连通子空间，维度2、3构成另一个独立的连通子空间

下面直接给出通用的结构生成函数和验证示例。

核心逻辑 #

scipy.ndimage.label的连通结构是一个(3,3,...,3)（共ndim个3）的数组，每个位置代表对应偏移量是否连通：

• 中心位置（所有维度偏移为0）必须设为1（自身肯定连通）
• 对于任意偏移量，只有当这个偏移的所有非零维度都属于同一个分组时，该位置才设为1，否则为0

函数实现 #

import numpy as np

def construct_arbitrary_structure(ndim, dim_connections):
    # 给每个维度分配对应的分组ID
    dim_to_group = np.zeros(ndim, dtype=int)
    for group_idx, dims in enumerate(dim_connections):
        for dim in dims:
            dim_to_group[dim] = group_idx
    
    # 生成所有可能的偏移组合（每个维度范围-1到1）
    offset_grid = np.meshgrid(*[np.arange(-1, 2)]*ndim, indexing='ij')
    all_offsets = np.stack(offset_grid, axis=-1).reshape(-1, ndim)
    
    # 初始化结构数组
    structure = np.zeros((3,)*ndim, dtype=int)
    
    for offset in all_offsets:
        # 找出偏移中非零的维度
        non_zero_dims = np.where(offset != 0)[0]
        if len(non_zero_dims) == 0:
            # 中心位置，自身连通
            structure[tuple(offset + 1)] = 1
            continue
        
        # 检查所有非零维度是否属于同一个分组
        group_ids = dim_to_group[non_zero_dims]
        if np.all(group_ids == group_ids[0]):
            structure[tuple(offset + 1)] = 1
    
    return structure

示例验证 #

场景1：分组[[0],[1],[2,3]]（4维） #

# 生成4维结构
struct1 = construct_arbitrary_structure(4, [[0],[1],[2,3]])
# 验证维度0的单偏移：(1,0,0,0)位置为1
print(struct1[2,1,1,1])  # 输出1
# 验证维度2+3的联合偏移：(0,0,1,1)位置为1
print(struct1[1,1,2,2])  # 输出1
# 验证跨分组偏移：(1,0,1,0)位置为0
print(struct1[2,1,2,1])  # 输出0

场景2：分组[[0,1],[2,3]]（4维） #

struct2 = construct_arbitrary_structure(4, [[0,1],[2,3]])
# 验证维度0+1的联合偏移：(1,1,0,0)位置为1
print(struct2[2,2,1,1])  # 输出1
# 验证跨分组偏移：(1,0,1,0)位置为0
print(struct2[2,1,2,1])  # 输出0

简单解释 #

1. 维度分组映射：先给每个维度贴分组标签，后续判断偏移是否合法时直接查标签就行
2. 生成偏移组合：用meshgrid生成所有可能的相邻偏移（每个维度左右各1步，中心是0）
3. 连通性判断：
   • 中心位置直接设为1，因为每个点肯定和自己连通
   • 其他偏移只要涉及的维度都在同一个分组里，就设为1，否则不连通

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