Обработка изображений: векторизация замены элементов массива

Question

Код работает хорошо, но очень медленно.Как я могу векторизовать замену цвета, чтобы избежать использования петли Python for?

processed_image = np.empty(initial_image.shape)
for i, j in np.ndindex(initial_image.shape[:2]):
    l_, a, b = initial_image[i, j, :]
    idx = mapping[a + 128, b + 128]
    a, b = new_colors[tuple(idx)]
    processed_image[i, j] = l_, a, b

У меня есть изображение initial_image в пространстве CIELAB в виде массива фигуры (некоторая высота, некоторая ширина, 3) .Мне нужно создать исправленное изображение, изменив a и b цветовые компоненты изображения, используя mapping.mapping - это массив фигур (255, 255, 2) .Он дает мне индексы, которые можно использовать для получения исправленных a и b цветов из new_colors.new_colors имеет форму (высота стола, ширина стола, 2) .

Решения, которые используют scikit-image, также будут полезны.

Answer 1

Вы можете использовать расширенную индексацию:

# chain the two maps
chained = new_colors[(*np.moveaxis(mapping, 2, 0),)]
# split color channels
c1, *c23 = np.moveaxis(initial_image, 2, 0)
# add 128
c23 = *map(np.add, c23, (128, 128)),
# apply chained map
processed_image_2 = np.concatenate([c1[..., None], chained[c23]], axis=2)