Векторизация вложенная для цикла

Question

Мне нужно векторизовать два вложенных цикла for и не знаю, как это сделать. Один для полутоновых изображений и один для цветных изображений. Я хочу отфильтровать изображение с помощью фильтра Кувахара. Код, который вы видите ниже, является последним шагом, который мне нужно векторизовать, чтобы получить быструю функцию.

• массив img_kuwahara имеет форму mxn или mxnx3 (цветное изображение)
• массив index_min имеет форму mxn
• среднее значение массива в форме 4xmxn (серая шкала) или 3x4xmxn (цвет)

Мне нужно получить правильное значение из среднего массива в массив 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]])

среднее = np.random.randint (0, 256, размер = (4,5,5)) (изображения в оттенках серого)

означает = np.random.randint (0, 256, размер = (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, а не списком трехмерных массивов; в противном случае просто преобразуйте его):

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]