Color segmentation using Kmeans, Opencv Python

Question

I have an image with a letter and a shape of different colors. I need to perform kmeans clustering on them and then provide two different images, one with just the shape regenerated and the other with just the Letter color regenerated. Here is a sample original image and what I need to achieve. Original Image

Shape color regenerated And similarly the other one with just the white R.

I have successfully performed kmeans clustering algorithm, How do I access the labels and cluster idx to regenerate the desired results? Can someone please illustrate with a sample code. Here is the code. Thanks in advance.

import numpy as np
import cv2

img = cv2.imread("/home/manohar/Outputs/Targets/m-0.PNG",1)
cv2.imshow("original",img)

Z = img.reshape((-1,3))

# convert to np.float32
Z = np.float32(Z)

# Here we are applying k-means clustering so that the pixels around a colour are consistent and gave same BGR/HSV values 


# define criteria, number of clusters(K) and apply kmeans()
criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 10, 1.0)
# We are going to cluster with k = 2, because the image will have just two colours ,a white background and the colour of the patch
K = 3
attempts=10
ret,label,center=cv2.kmeans(Z,K,None,criteria,attempts,cv2.KMEANS_PP_CENTERS)

# Now convert back into uint8
#now we have to access the labels to regenerate the clustered image
center = np.uint8(center)
res = center[label.flatten()]
res2 = res.reshape((img.shape))
#res2 is the result of the frame which has undergone k-means clustering


cv2.imshow("res2",res2)
cv2.waitKey()
cv2.destroyAllWindows()

Answer 1

Ok, so if you want to change to black all pixels that were classified as "white" by K-Means you need to first see which label corresponds to the "white" category. You can achieve this by looking to what index in center (which contains the resulting k centers) does the white color (255,255,255) belongs to, like this:

white_index = 0
for b,g,r in center:
    #check if it is white
    if b == g == r == 255:
        #then this is the white category
        break
    white_index = white_index + 1

Do note that as k-means is an unsupervized method your category may not necessarily be exactly white (it could be something like 250,249,254). So you should consider this when searching for the index, that is you should look for the color closer to white. You could achieve this by comparing your colors by applying the color distance formula.

However in my experience, I believe that center already sorts the resulting centers in some way (greater pixel values tend to appear first), and I have noticed the centers closer to white have lower indexes (therefore white may be index 0). But it is better to check for sure.

Now that you know what index corresponds to the white color you can see which pixels were classified to that in the label variable. For this there are several ways you could do it and surely some more efficient than others, one approach could be:

#Copy image to modify it
img_copy = img[:]
#Reshape label to match the original image
new_label = label.reshape(img.shape)
#iterate over new_label, and if the category is white
#paint pixel corresponding to row,col black
row = 0
for i in new_label:
    column = 0
    for category in i:
        if category == white_index:
            #paint that pixel black
            img_copy[row,column] = [0, 0, 0]
        column = column + 1
    row = row + 1

#now show edited image
cv2.imshow("Object Removed",img_copy)
cv2.waitKey()
cv2.destroyAllWindows()

Edit: The above code obtains an image with the detected color removed (blacked out). In order to obtain its complement, that is an image with only the detected object visible, there are several things you could do. One approach could be to obtain a copy of the img and black out the pixels that are not of that color (without category == white_index), like:

#Obtain another copy of image to modify it
extract = img[:]
#Now do the opposite, iterate over new_label, and if the category is  Not white
#paint pixel corresponding to row,col black
row = 0
for i in new_label:
    column = 0
    for category in i:
        if category != white_index:
            #paint that pixel black
            extract[row,column] = [0, 0, 0]
        column = column + 1
    row = row + 1

#show or do whatever you want...

Another more complex way of obtaining the extracted object only is to obtain the contour of your object (by thresholding and using cv2.findContours()), and then apply cv2.boundingRect() so you can obtain a smaller clip with the object (by clipping img with the given rectangle). You could check that link where it is explained in detail.