test.py

# http://stackoverflow.com/questions/24685436/efficient-way-to-cluster-colors-using-k-nearest

import numpy as np
import cv2

def nearest(i, j, src, knn, colors):
    sample = np.reshape(src, (-1, 3)).astype(np.float32)
    retval, result, neighbors, dist = knn.find_nearest(sample, 1)
    return colors[result[0, 0]]

@profile
def main():
    src = cv2.imread('object.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)

    dst = [nearest(i, j, src[i, j], knn, colors)
           for i in range(0, src.shape[0])
           for j in range(0, src.shape[1])]
    #cv2.imshow('src', src)
    #dst = np.reshape(dst, src.shape)
    #cv2.imshow('dst', dst)
    #cv2.waitKey()

if __name__ == '__main__':
    main()