这个方法很慢。简而言之，它采用字典phase_color_labels，将任意名称映射到对应于RGB值的3元素列表，并将输入图像的每个像素映射到phase_color_labels字典中最接近的任何像素值。我还没有弄清楚是否有一个矢量化版本可以运行得更快。
image变量只是一个 numpy 数组 [H， W， Channels]。
def map_pixels_to_discrete_values(image, phase_color_labels):
"""
Takes an image with floating point pixel values and maps each pixel RGB value
to a new value based on the closest Euclidean distance to one of the RGB sets
in the phase_label input dictionary.
"""
mapped_image = np.copy(image)
for i in range(mapped_image.shape[0]):
for j in range(mapped_image.shape[1]):
min_distance = np.inf
min_distance_label = None
for phase_name, phase_color in phase_color_labels.items():
r = phase_color[0]
g = phase_color[1]
b = phase_color[2]
rgb_distance = (mapped_image[i, j, 0] - r)**2 + (mapped_image[i, j, 1] - g)**2 + (mapped_image[i, j, 2] - b)**2
if rgb_distance < min_distance:
min_distance = rgb_distance
min_distance_label = phase_name
mapped_image[i, j, :] = phase_color_labels[min_distance_label]
return mapped_image
为了使用 Numpy 快速完成任务，您通常希望避免循环并将尽可能多的工作推送到 Numpy 的矩阵操作中。
我的答案的基本思路：
- 从
phase_color_labels获取颜色作为ndarray，phase_colors。 - 使用 Numpy 的广播来计算"外部距离"，即数组中每个图像与
phase_colors中每种颜色之间的欧氏距离。 - 找到每个像素距离最低的颜色的索引，并将其用作
phase_colors中的索引。
phase_colors = np.array([color for color in phase_color_labels.values()])
distances = np.sqrt(np.sum((image[:,:,np.newaxis,:] - phase_colors) ** 2, axis=3))
min_indices = distances.argmin(2)
mapped_image = phase_colors[min_indices]
第三行需要一些额外的解释。首先，请注意phase_names和phase_colors都有形状(L, C)，其中L是标签的数量，C是通道的数量。
image[:,:,np.newaxis,:]在第二个和第三个轴之间插入一个新轴，因此生成的数组具有形状(H, W, 1, C)。- 当从形状
(H, W, 1, C)数组中减去形状(L, C)数组时，Numpy 会将数组广播到形状(H, W, L, C)。您可以在此处找到有关Numpy广播语义的更多详细信息。 - 然后，沿轴 3 取总和会产生一个形状
(H, W, L)数组。 - (平方和平方根都不会影响数组的形状。
在第四行中，在轴 2 上使用argmin然后将数组缩小为形状(H, W)，每个值都是从缩小轴L开始的索引 - 换句话说，索引为phase_colors。
作为额外的改进，由于平方根是一个单调递增的函数，它不会改变哪个距离最小，因此您可以完全删除它。
请注意，对于大image和phase_color_labels，广播的内存成本可能会很明显，这也可能导致性能问题。