如何将二维Numpy掩码数组扩展至与多维数据数组尺寸完全匹配？

要实现任意维度数据数组对应的二维掩码扩展，核心是先对齐维度结构，再通过复制填充非掩码维度，完全保留掩码的空间结构。以下是通用实现方案：

方案思路 #

1. 匹配掩码与数据数组的维度结构：先将二维掩码调整为与数据数组维度数量一致的形状，仅保留掩码自身的两个维度大小，其余维度设为1；
2. 按数据数组的各维度大小复制掩码，最终得到与数据形状完全一致的掩码数组。

通用代码实现 #

import numpy as np

def expand_mask_to_data_shape(data, mask):
    data_shape = data.shape
    mask_shape = mask.shape
    assert len(mask_shape) == 2, "掩码必须为二维数组"
    
    # 找到掩码维度在数据形状中的对应索引（假设掩码的两个维度与数据中的某两个维度大小匹配）
    mask_dim_indices = []
    for dim in mask_shape:
        # 找到数据形状中第一个匹配的维度索引
        idx = data_shape.index(dim)
        mask_dim_indices.append(idx)
    
    # 构造掩码的扩展基础形状（其余维度设为1）
    expanded_base_shape = [1] * len(data_shape)
    for idx, mask_dim in zip(mask_dim_indices, mask_shape):
        expanded_base_shape[idx] = mask_dim
    
    # 调整掩码形状并复制至目标尺寸
    mask_expanded = mask.reshape(expanded_base_shape)
    return np.tile(mask_expanded, data_shape)

# 测试示例
data = np.random.rand(11, 1, 1, 224, 464, 1, 224)
mask = np.random.rand(224, 464)
mask_final = expand_mask_to_data_shape(data, mask)

print(mask_final.shape)  # 输出：(11, 1, 1, 224, 464, 1, 224)

方案说明 #

• 无需提前确定数据数组的维度数量，代码会自动适配任意大于等于2的维度；
• 通过形状匹配自动定位掩码在数据数组中的维度位置，确保空间结构完全保留；
• 使用np.tile实现维度复制，相比np.resize不会打乱掩码的空间排布，相比手动np.repeat更通用。

如果你的掩码维度在数据数组中存在多个匹配（比如数据形状中有多个224的维度），可以根据实际场景修改mask_dim_indices的获取逻辑，手动指定掩码对应的维度索引即可。

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