如何基于二进制掩码获取最近邻像素并生成目标矩阵？

现有原矩阵：

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

和二进制掩码：

array([[1, 1, 1, 1],
       [1, 1, 1, 1],
       [1, 0, 0, 1],
       [1, 0, 0, 1]])

需要生成新矩阵：掩码为0的位置保留原矩阵元素，非0位置选取掩码0区域的最近邻元素，目标输出如下：

array([[10, 10, 11, 11],
       [10, 10, 11, 11],
       [10, 10, 11, 11],
       [14, 14, 15, 15]])

已尝试使用OpenCV的cv2.distanceTransformWithLabels(binary.astype('uint8'), cv2.DIST_L2, cv2.DIST_MASK_PRECISE)命令得到距离矩阵和标签矩阵，但无法获取最近邻元素的索引，请问如何通过OpenCV或其他Python工具实现上述需求？

方法一：基于OpenCV的distanceTransformWithLabels实现 #

cv2.distanceTransformWithLabels返回的标签矩阵可以和掩码0区域的像素坐标建立映射关系，具体步骤如下：

1. 反转掩码（因为该函数默认计算到非0区域的距离，我们需要计算到原掩码0区域的距离）
2. 提取原掩码0区域的所有像素坐标，建立标签到坐标的映射表
3. 遍历掩码非0的位置，通过标签映射找到最近邻的坐标，替换原矩阵对应位置的值
4. 掩码0的位置保留原矩阵元素

代码示例：

import numpy as np
import cv2

# 原矩阵
src = np.array([[1, 2, 3, 4],
                [5, 6, 7, 8],
                [9, 10, 11, 12],
                [13, 14, 15, 16]], dtype=np.int32)
# 二进制掩码
mask = np.array([[1, 1, 1, 1],
                 [1, 1, 1, 1],
                 [1, 0, 0, 1],
                 [1, 0, 0, 1]], dtype=np.uint8)

# 反转掩码，将原掩码的0区域转为1，用于计算距离
inv_mask = 1 - mask

# 计算距离矩阵和标签矩阵
dist_matrix, label_matrix = cv2.distanceTransformWithLabels(
    inv_mask, 
    cv2.DIST_L2, 
    cv2.DIST_MASK_PRECISE
)

# 获取原掩码0区域的所有坐标（即反转后掩码的1区域）
target_coords = np.argwhere(inv_mask == 1)

# 建立标签到坐标的映射：标签从1开始，对应target_coords的索引顺序
label_to_coord = {i+1: coord for i, coord in enumerate(target_coords)}

# 初始化结果矩阵，先复制原矩阵
result = src.copy()

# 遍历所有掩码非0的位置，替换为最近邻值
non_zero_positions = np.argwhere(mask == 1)
for y, x in non_zero_positions:
    current_label = label_matrix[y, x]
    # 获取最近邻的坐标
    nearest_y, nearest_x = label_to_coord[current_label]
    result[y, x] = src[nearest_y, nearest_x]

print(result)

运行后输出：

array([[10, 10, 11, 11],
       [10, 10, 11, 11],
       [10, 10, 11, 11],
       [14, 14, 15, 15]])

方法二：使用SciPy库简化实现 #

如果可以引入SciPy库，ndimage.distance_transform_edt可以直接返回每个位置到最近目标区域的索引，代码更简洁：

import numpy as np
from scipy.ndimage import distance_transform_edt

src = np.array([[1, 2, 3, 4],
                [5, 6, 7, 8],
                [9, 10, 11, 12],
                [13, 14, 15, 16]], dtype=np.int32)
mask = np.array([[1, 1, 1, 1],
                 [1, 1, 1, 1],
                 [1, 0, 0, 1],
                 [1, 0, 0, 1]], dtype=np.uint8)

# 获取每个位置到最近掩码0区域的坐标索引
nearest_y, nearest_x = distance_transform_edt(mask, return_indices=True)

# 生成结果：掩码0位置保留原值，其余位置用最近邻值替换
result = np.where(mask == 0, src, src[nearest_y, nearest_x])

print(result)

该代码同样能输出目标矩阵。

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