python - 向量化嵌套 for 循环

Question

我需要向量化两个嵌套的 for 循环，但不知道该怎么做。一种用于灰度图像，一种用于彩色图像。我想用 kuwahara 过滤器过滤图像。您在下面看到的代码是我需要矢量化以获得快速函数的最后一步。

• 数组 img_kuwahara 的形状为 mxn 或 mxnx3(彩色图像)
• 数组index_min的形状为mxn
• 数组均值的形状为 4xmxn(灰度)或 3x4xmxn(彩色)

我需要将 mean 数组中的正确值放入 img_kuwahara 数组中。

您可以使用以下数组作为示例数据:

index_min = np.array([[0, 1, 1, 2, 3],[3, 3, 2, 2, 2],[2, 3, 3, 0, 2],[0, 1, 1, 0, 3],[2, 1, 3, 0, 0]])

mean = np.random.randint(0, 256, size=(4,5,5))(灰度图像)

均值 = np.random.randint(0, 256, size=(3,4,5,5))(彩色图像)

行 = 5，列 = 5

谢谢你的帮助

# Edit gray scale image
if len(image.shape) == 2:

    # Set result image
    img_kuwahara = np.zeros((row, columns), dtype=imgtyp)

    for k in range(0, row):
        for i in range(0, columns):
            img_kuwahara[k, i] = mean[index_min[k, i], k, i]



# Edit color image
if len(image.shape) == 3:

    # Set result image
    img_kuwahara = np.zeros((row, columns, 3), dtype=imgtyp)

    for k in range(0, row):
        for i in range(0, columns):
            img_kuwahara[k, i, 0] = mean[0][index_min[k, i], k, i]
            img_kuwahara[k, i, 1] = mean[1][index_min[k, i], k, i]
            img_kuwahara[k, i, 2] = mean[2][index_min[k, i], k, i]

Answer 1

可以使用 np.meshgrid 对第一个循环进行矢量化:

j, i = np.meshgrid(range(columns), range(rows))
img_kuwahara = mean[index_min[i, j], i, j]

可以使用额外的 np.moveaxis 对第二个循环进行矢量化(假设 mean 在这种情况下实际上是一个 4D 数组，而不是 3D 数组列表；否则只需转换它):

j, i = np.meshgrid(range(columns), range(rows))
img_kuwahara = np.moveaxis(mean, 0, -1)[index_min[i, j], i, j]

除了 np.meshgrid，您还可以使用 np.mgrid (支持更自然的语法):

i, j = np.mgrid[:rows, :columns]