我正在尝试将图像上的颜色聚类到预定义的类别中（黑色、白色、蓝色、绿色、红色）。我使用了以下代码：

import numpy as np
import cv2
src = cv2.imread('objects.png')
colors = np.array([[0x00, 0x00, 0x00],
                   [0xff, 0xff, 0xff],
                   [0xff, 0x00, 0x00],
                   [0x00, 0xff, 0x00],
                   [0x00, 0x00, 0xff]], dtype=np.float32)
classes = np.array([[0], [1], [2], [3], [4]], np.float32)
dst = np.zeros(src.shape, np.float32)
knn = cv2.KNearest()
knn.train(colors, classes)
# 这个循环非常低效！
for i in range(0, src.shape[0]):
    for j in range(0, src.shape[1]):
        sample = np.reshape(src[i,j], (-1,3)).astype(np.float32)
        retval, result, neighbors, dist = knn.find_nearest(sample, 1)
        dst[i,j] = colors[result[0,0]]
cv2.imshow('src', src)
cv2.imshow('dst', dst)
cv2.waitKey()

代码运行良好，结果如下图所示。左图是输入图像，右图是输出图像。

然而，上面的循环非常低效，导致转换速度缓慢。有什么最有效的Numpy操作可以替代上面的循环吗？

回答：

我设法使用下面的代码去除了循环。代码运行非常快，几乎与C++版本相似。

import numpy as np
import cv2
src = cv2.imread('objects.png')
src_flatten = np.reshape(np.ravel(src, 'C'), (-1, 3))
dst = np.zeros(src.shape, np.float32)
colors = np.array([[0x00, 0x00, 0x00],
                   [0xff, 0xff, 0xff],
                   [0xff, 0x00, 0x00],
                   [0x00, 0xff, 0x00],
                   [0x00, 0x00, 0xff]], dtype=np.float32)
classes = np.array([[0], [1], [2], [3], [4]], np.float32)
knn = cv2.KNearest()
knn.train(colors, classes)
retval, result, neighbors, dist = knn.find_nearest(src_flatten.astype(np.float32), 1)
dst = colors[np.ravel(result, 'C').astype(np.uint8)]
dst = dst.reshape(src.shape).astype(np.uint8)
cv2.imshow('src', src)
cv2.imshow('dst', dst)
cv2.waitKey()

代码生成的结果与之前相同，但执行速度更快。