如何从NumPy三维矩阵数组中筛选含指定类别的子矩阵

你已经生成的horse_mask是形状为(1735, 128, 128)的逐元素布尔数组，仅需两步即可完成目标筛选：

• 沿后两个空间维度做存在性判断，生成第一维的筛选掩码：对每个(128, 128)规格的子矩阵，只要内部存在任意一个值为True的位置，就判定该子矩阵需要保留。对应代码如下：

# 归约后得到形状为(1735,)的布尔数组，标记每个子矩阵是否包含类别9
keep_mask = horse_mask.any(axis=(1, 2))

• 用生成的一维筛选掩码索引原数组的第一维，直接得到过滤后的结果：

# 输出数组形状为(N, 128, 128)，N为含类别9的子矩阵总数
filtered_arr = arr_seg_labs[keep_mask]

该操作不会修改保留下来的子矩阵内部的任何像素值，仅剔除完全不包含类别9的子矩阵，计算效率远高于循环遍历每个子矩阵的写法。

内容的提问来源于stack exchange，提问作者Zoltán Orosz