如何使用numpy将图像中的颜色快速映射到调色板

Question

我有两个数组。一个是图像阵列，另一个是调色板阵列。两者都具有包含8位RGB通道的元素。我需要用调色板中最接近的颜色替换图像中的每一种颜色。

目前我正在RGB空间中测量距离，这并不理想，但很容易实现。

这是我的实现：

image_array = np.array(image) # converts PIL image, values are uint8
# palette values are also 8-bit but I use int so I don't have to cast types
palette_array = np.array(palette, dtype=[('red', np.int), ('green', np.int), ('blue', np.int)])
mapped_image = np.empty((image_height, image_width, 3), dtype=np.uint8)
for x in range(image_width):
for y in range(image_height):
r, g, b = image_array[y, x]
distances_squared = (r-palette['red'])**2 + (g-palette['green'])**2 + (b-palette['blue'])**2
closest_index = np.argmin(distances_squared)
closest_color = palette.flat[closest_index]
mapped_image[y, x] = closest_color

调色板有4096种随机颜色(无法进行简单转换(。当映射一个600x448大小的图像时，即使在我的核心i5机器上，这也需要大约一分钟的时间。我计划在像树莓派这样的低端设备上使用它，在那里绘制一个小图像大约需要3分钟。

这太慢了。我相信，当使用numpy语法实现整个循环时，这可以大大加快速度，但我不知道如何做到这一点。

如何从原始图像到映射图像，所有这些都是用numpy语法实现的？

Answer 1

您可以尝试使用scipy中的cKDTree函数。

import numpy as np
from scipy.spatial import cKDTree
palette=np.random.randint(0, 255, size=(4096,3), dtype=np.uint8) # random palette
image_in=np.random.randint(0, 255, size=(800, 600, 3), dtype=np.uint8) # random image
size=image_in.shape
vor=cKDTree(palette)
test_points=np.reshape(image_in, (-1,3))
_, test_point_regions = vor.query(test_points, k=1)
image_out=palette[test_point_regions]
np.reshape(image_out, size)

该程序运行约0.8秒。