Image Segmentation using K-means. Image Segmentation from scratch using... | by Kartik Narula | Towards Data Science

While pre-existing libraries(such as OpenCV) save time and effort, implementing the basic algorithms from scratch is another delight.
In this post, I will show the step by step implementation of image segmentation using k-means in python. We train the pipeline on 1100 images across 8 categories sampled from the SUN database. Image segmentation is the grouping of pixels of similar types together. The pipeline can be further extended to classify an image. An image of a park, for instance, will have a greater number of green-colored pixels, for example than an image of a highway. We will talk about classification in a later post.
We established that an image of a park will have a greater number of green-colored pixels. While color is a good way of differentiating pixels, it can create issues at times(water and sky have the same color, for instance). Hence, we need other ways to characterize pixels. In other words, we need to extract useful features from the image. Deep learning algorithms, such as CNN’s find useful features automatically for us but’s that’s outside the scope of this article.
For extracting features, we will use four kinds of image filters with multiple scales- gaussian, derivative of gaussian in the x-direction, derivative of gaussian in the y-direction and the Laplacian of gaussian. Each filter extracts different sets of features from the image. Smaller scales pick up narrower features and larger scales pick up broader features(Think of forest and the trees). We will call the set of filters with different scales as a ‘filter bank’.
Each row in the above figure represents one particular scale and each column represents a different type of filter. Gaussian filter(first column) blurs the image and omits the higher frequencies. Derivative of Gaussian in X direction(second column) picks up the vertical edges. The derivative of gaussian in Y direction(third column) picks up the horizontal edges. The laplacian of gaussian (fourth column) detects regions of rapid intensity changes.
We run the filter bank over an image(which is an array of size H*W*3, where H and W are the height and width of the image, respectively and 3 is the number of channels). We output an array of size(H*W*3F), where F is the size of the filter bank. If we use 4 filters at 5 different scales, F is 4*5=20
From T training images, we get a filter response vector of size T*H*W*3F. To save computation resources, we select alpha random filter responses from each image. Thus, the final size of the filter response vector is alpha* T * F.
We will use this vector of filter responses as the input to our k-means algorithm. It will output K(a hyperparameter) cluster centers. Each pixel of the test image will then be mapped to the nearest cluster center and painted with a corresponding color to form a ‘visual map’. The resulting ‘visual map’ would have similar pixels grouped together and would be the final output of our ‘image segmentation’ algorithm.
Bravo, let’s see the whole pipeline in action. We will start by defining the values of the hyperparameters, which can be tuned later to enhance performance.
filter_scales=[1,1.5,2]K=10alpha=25
We now define the function extract_fiter_responses which returns a vector of size H*W*3F for an image of size H*W*3 and filter bank size F. We use the inbuilt functions in scipy for generating filter_responses.
def extract_filter_responses(filter_scales, img):                   if len(img.shape)==2:     #convert image into 3 channel        img=np.dstack((img,img,img))            modified_img=skimage.color.rgb2lab(img)    r_channel=modified_img[:,:,0]    g_channel=modified_img[:, :,1]    b_channel=modified_img[:, : ,2]    g=[]    for i in range(len(filter_scales)):        modified_r_channel=scipy.ndimage.gaussian_filter(r_channel,      filter_scales[i])        modified_g_channel=scipy.ndimage.gaussian_filter(g_channel, filter_scales[i])        modified_b_channel=scipy.ndimage.gaussian_filter(b_channel, filter_scales[i])                modified_r_dog_X=scipy.ndimage.gaussian_filter(r_channel, filter_scales[i],(0,1))        modified_g_dog_X=scipy.ndimage.gaussian_filter(g_channel, filter_scales[i], (0,1))        modified_b_dog_X=scipy.ndimage.gaussian_filter(b_channel, filter_scales[i], (0,1))            modified_r_dog_Y=scipy.ndimage.gaussian_filter(r_channel, filter_scales[i],(1,0))        modified_g_dog_Y=scipy.ndimage.gaussian_filter(g_channel, filter_scales[i], (1,0))        modified_b_dog_Y=scipy.ndimage.gaussian_filter(b_channel, filter_scales[i], (1,0))                modified_r_log=scipy.ndimage.gaussian_laplace(r_channel, filter_scales[i])        modified_g_log=scipy.ndimage.gaussian_laplace(g_channel, filter_scales[i])        modified_b_log=scipy.ndimage.gaussian_laplace(b_channel, filter_scales[i])                                        a=np.dstack((modified_r_channel, modified_g_channel, modified_b_channel))        b=np.dstack((modified_r_dog_X, modified_g_dog_X, modified_b_dog_X))        c=np.dstack((modified_r_dog_Y, modified_g_dog_Y, modified_b_dog_Y))        d=np.dstack((modified_r_log, modified_g_log, modified_b_log))        filter_response=np.dstack((a,b,c,d))        g.append(filter_response)                            filter_responses=np.dstack(g)            return filter_responses
The next function will extract alpha random filter responses from the vector returned by extract_filter_responses
def compute_dictionary_one_image(alpha, img):       response=extract_filter_responses(filter_scales, img)     d=response.shape[0]*response.shape[1]    response=response.reshape((d,-1))     alphas=np.random.choice(d, alpha)    alphaed_response=response[alphas]       return alphaed_response
The next function pools the alpha responses from each image into a single vector and applies the k-means function from scipy. The function takes the training files path as an argument.
def compute_dictionary(K, alpha, train_files):    m=[]    for i in range(len(train_files)):               img_path = train_files[i]        img = Image.open(img_path)        img = np.array(img).astype(np.float32)/255        re=compute_dictionary_one_image(alpha, img)        m.append(re)            m=np.array(m)    n=m.shape[0]*m.shape[1]    final_response=m.reshape((n,-1))        kmeans=KMeans(n_clusters=K).fit(final_response)    dictionary=kmeans.cluster_centers_    return dictionary
Cool. So now we have K cluster centers. For a test image, we need to map each pixel of the image with the nearest cluster center and plot the resulting visualization. We use the scipy.spatial.distance.cdist to find the index of the closest cluster center. This is illustrated in the get_visual_maps() function below
def get_visual_words(filter_scales, img, dictionary):      response=extract_filter_responses(filter_scales, img)   response=response.reshape(response.shape[0]*response.shape[1],-1)       dist=scipy.spatial.distance.cdist(response, dictionary)       visual_words=np.argmin(dist, axis=1)   visual_words=visual_words.reshape(img.shape[0],img.shape[1])       return visual_words
We will now use these functions to create a pipeline that runs k-means on a training set of 1177 images and check the performance on a test image.
#compute cluster centerscompute_dictionary(K, alpha, train_files)#test on a image      img_path = 'image path'img = Image.open(img_path)img = np.array(img).astype(np.float32)/255wordmap = visual_words.get_visual_words(filter_scales, img, dictionary)plt.imshow(wordmap)
Here are a few sample visualizations of the algorithm. I used Filter scales=[1,1.5, 2], K=25 and alpha=10. We are able to identify different regions of the image. In the first image(of a kitchen) we can identify different regions, the table, stove, etc, and so on. Feel free to tune these parameters and see if you can get even better results!
In the next post, we will extend the pipeline that classifies an image by its type(The first is an image of a kitchen, second is an image of a highway, and so on..)