CODEWITHSUNDEEP
pythonopencvnumpymachine-learning
flowfree
flowfree

Reputation: 16462

Efficient way to cluster colors using K-Nearest

I am trying to cluster colors on an image to a predefined classes (black, white, blue, green, red). I'm using the following code:

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()

The code works well and result is shown below. The image on the left is the input and the image on the right is the output.

src dst

However the loop above is very inefficient and make the conversion slow. What is the most efficient Numpy operation to replace the loop above?

Upvotes: 2

Views: 1166

Answers (3)

flowfree
flowfree

Reputation: 16462

I managed to remove the loop using the code below. The code runs very fast, almost similar with the C++ version.

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()

The code generated the correct result as before, with faster execution time.

src dst

Upvotes: 0

ivan_a
ivan_a

Reputation: 613

You can build a look up table. So that you know corresponding class for each color. It doesn't have to be 256x256x256 you can reduce number of bins.

Upvotes: 0

mdurant
mdurant

Reputation: 28683

If you want a simple squared difference measure ("which is the euclidian nearest number), this will work.

Calculate differences

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

(assuming src is y,x,3 in shape)

Pick closest colour index:

index = diff.argmin(axis=2)

New image:

out = colors[index]

If your colours are really to have component values of 0 or 0xff, you can use something like

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

Upvotes: 2

Related Questions