如何在NumPy中利用4D掩码将2D数组赋值给4D数组？

问题背景 #

现有形状为(H,W,H-h+1,W-w+1)的4D掩码，其中mask[:,:,dh,dw]对应从图像左上角偏移(dh,dw)的h×w滑动窗口掩码。需要将一个h×w的矩阵放置到掩码指定的位置，实现通过4D掩码给4D数组赋值。

示例掩码：

>>> mask[:,:,0,0]
array([[1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0],
       [1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0],
       [1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0],
       [1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]])

>>> mask[:,:,1,1]
array([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
       [0, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0],
       [0, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0],
       [0, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0],
       [0, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]])

待赋值的h×w矩阵：

array([[ 1,  2,  3,  4,  5],
       [ 6,  7,  8,  9, 10],
       [11, 12, 13, 14, 15],
       [16, 17, 18, 19, 20]])

期望偏移(0,0)时的结果：

>>> result[:,:,0,0]
array([[1, 2, 3, 4, 5, 0, 0, 0, 0, 0, 0, 0],
       [6, 7, 8, 9, 10, 0, 0, 0, 0, 0, 0, 0],
       [11, 12, 13, 14, 15, 0, 0, 0, 0, 0, 0, 0],
       [16, 17, 18, 19, 20, 0, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]])

解决方案 #

方法一：遍历偏移窗口（直观易懂） #

通过遍历每个滑动窗口的偏移(dh,dw)，利用掩码的布尔索引定位目标区域，直接赋值：

import numpy as np

# 定义基础参数
H, W = 10, 12
h, w = 4, 5

# 初始化掩码并填充示例值
mask = np.zeros((H, W, H-h+1, W-w+1), dtype=int)
mask[:h, :w, 0, 0] = 1
mask[1:1+h, 1:1+w, 1, 1] = 1

# 待赋值的h×w矩阵
patch = np.array([
    [1, 2, 3, 4, 5],
    [6, 7, 8, 9, 10],
    [11, 12, 13, 14, 15],
    [16, 17, 18, 19, 20]
])

# 创建结果数组，形状与掩码一致
result = np.zeros_like(mask, dtype=patch.dtype)

# 遍历所有偏移组合
for dh in range(H - h + 1):
    for dw in range(W - w + 1):
        # 获取当前偏移下的有效掩码区域
        valid_area = mask[:, :, dh, dw] == 1
        # 将patch展平后赋值到对应位置
        result[valid_area, dh, dw] = patch.flatten()

# 验证偏移(0,0)的结果
print("result[:,:,0,0]:")
print(result[:, :, 0, 0])

方法二：向量化操作（高效处理大尺寸数据） #

利用NumPy的向量化特性，避免循环，提升处理效率：

import numpy as np

# 定义参数与初始化数据（同方法一）
H, W = 10, 12
h, w = 4, 5
mask = np.zeros((H, W, H-h+1, W-w+1), dtype=int)
mask[:h, :w, 0, 0] = 1
mask[1:1+h, 1:1+w, 1, 1] = 1
patch = np.array([
    [1, 2, 3, 4, 5],
    [6, 7, 8, 9, 10],
    [11, 12, 13, 14, 15],
    [16, 17, 18, 19, 20]
])

result = np.zeros_like(mask, dtype=patch.dtype)

# 生成所有偏移的网格矩阵
dh_grid, dw_grid = np.meshgrid(np.arange(H-h+1), np.arange(W-w+1), indexing='ij')

# 提取掩码有效区域对应的dh和dw值
dh_vals = dh_grid[mask[:h, :w, :, :] == 1]
dw_vals = dw_grid[mask[:h, :w, :, :] == 1]

# 将patch展平并重复对应次数，匹配有效区域的数量
patch_repeated = np.tile(patch.flatten(), (H-h+1)*(W-w+1))

# 直接赋值到所有有效位置
result[mask == 1] = patch_repeated

原理说明 #

两种方法核心都是利用布尔索引定位掩码中的1值区域，再将h×w矩阵的内容映射到这些区域。遍历法适合小数据量，逻辑直观；向量化法通过广播和网格生成，大幅提升大尺寸数据的处理速度。

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