使用K-最近的方法來聚集顏色

Question

我想將圖像上的顏色聚類到預定義的類（黑色，白色，藍色，綠色，紅色）。我正在使用以下代碼：

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) 

# This loop is very inefficient! 
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操作來取代上面的循環？

Answer 1

我設法刪除使用下面的代碼循環。代碼運行速度非常快，幾乎與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() 

該代碼與以前一樣生成正確的結果，執行時間更快。

Answer 2

如果你想要一個簡單的平方差的措施（「這是歐幾里得最接近的數字），這會工作。

計算差異

diff = ((src[:,:,:,None] - colors.T)**2).sum(axis=2) 

（假設src是Y，X，3形狀）

接最接近的顏色指數：

index = diff.argmin(axis=2) 

新形象：

out = colors[index] 

如果你的顏色真的有0或0xFF的元件值，可以使用類似

out = np.where(src>0x88, 0xff, 0) 

Answer 3

你可以建立一個查找表。這樣你就知道每種顏色的相應類。它不一定是256x256x256您可以減少垃圾箱數量。