Python训练Kmeans算法预测图像的主色

Question

I am trying to create a model that will predict the dominant color in the image using K-means clustering. I have the data all set up, but I am unsure how I can proceed after fitting the model. Thanks

 from sklearn.cluster import KMeans
 import h5py


 train_data = h5py.File('x_train.h5','r')
 test_data = h5py.File('x_test.h5','r')

 x_train = train_data['train'][:]
 x_test = test_data['test'][:]

 print(x_train.shape) # (429-number of images, 416-height,416-width, 3-channels)

 x_train = x_train/255.0
 x_test = x_test/255.0

 X_train = x_train.reshape(len(x_train),-1)
 X_test = x_test.reshape(len(x_test),-1)

 kmeans = KMeans(n_clusters = 5)
 # Fitting the model to training set
 kmeans.fit(X_train)

 #------edit------

 pred = kmeans.predict(X_test[0])

 labels=pred.labels_
 labels=list(labels)

 centroid=pred.cluster_centers_

 percent=[]
 for i in range(len(centroid)):
     x=labels.count(i)
     x=x/(len(labels))
     percent.append(x)

 get_label_index = percent.index(max(percent))

 get_rgb_of_dominant_color = centroid[get_label_index][:]

 print(get_rgb_of_dominant_color)

Answer 1

This is one approach that I can think of. Suppose, you are fixing the clusters as "5" as done in your code.

Identify the 5 cluster centroids using: kmeans.cluster_centers_

Rank the cluster centroids in the order of 1 to 5 based on the number of datapoints associated with each of them. The cluster centroid with the highest number of datapoints associated to it would be the dominant one. Use the cluster centroid's RBG value and visualize to see the color.

---

EDITED - Added code for detailed explanation

---

Below is the code where I have loaded an image and then trying to find the most dominant color.

import numpy as np
import pandas as pd
from PIL import Image
import matplotlib.pyplot as plt
from matplotlib.pyplot import imshow
from sklearn.cluster import KMeans
%matplotlib inline

#Load the image
arr_img = np.array(Image.open("beach.bmp"), dtype='int32')
plt.imshow(arr_img)

#reshape array from 3D to 2D
r, c, l = arr_img.shape
reshape_img = np.reshape(arr_img, (r*c, l), order="C")
#fit the model with 5 clusters
kmeans = KMeans(n_clusters = 5 ,max_iter=1000, init='random')
kmeans.fit(reshape_img)
# Looking at the labels and their associated data points
unique, counts = np.unique(kmeans.labels_, return_counts=True)
print("The labels are: ",unique)
print("Count of  items: ",counts)
# Find the most dense cluster label
idx = np.where(counts == counts.max())[0]
# Pick the mose dense cluster centroid
s = tuple(map(int,kmeans.cluster_centers_[idx][0]))
# Visualize the color
plt.imshow([[s]])

You can see that Kmeans has identified Blue as the most dominant color correctly.

Answer 2

I was taught that you'd first check the number of kmeans you'd need for looking when the inertia improvement drops below 20%, which could be done by:

test = []
K = range(1,10)

for k in K:
    model = KMeans(n_clusters=k)
    model.fit(X)
    test.append(model.inertia_)

for index, x in enumerate(test):
    if index == 0:
        continue
    else:
        print(index, (((test[index - 1] - x) / test[index - 1]) * 100))

But I guess for you it would be the number of different colors. I would advice you to use that number then, which I assume is 5 looking at your code.

After that, to predict the color, you would do

preds = kmeans.predict(whateverYouWantToPredict)

Those would be your final predictions then, but note that this is an unsupervised method . You could use the predictions of this by appending the predictions to the dataframe and then use that data to train a supervised model for another prediction.