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

Posted 2023-02-23

【中文标题】Python训练Kmeans算法预测图像的主色

【英文标题】：Python training Kmeans algorithm to predict the dominant color of a image

【发布时间】：2021-08-14 11:27:31

【问题描述】：

我正在尝试创建一个模型，该模型将使用 K-means 聚类预测图像中的主色。我已经设置了所有数据，但我不确定在拟合模型后如何继续。谢谢

 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)

【参考方案1】：

这是我能想到的一种方法。假设您按照代码中的做法将集群修复为“5”。

使用以下方法识别 5 个集群质心： kmeans.cluster_centers_

根据与每个聚类质心相关联的数据点的数量，按 1 到 5 的顺序对聚类质心进行排名。与其关联的数据点数量最多的集群质心将是主要的。使用簇质心的 RBG 值并可视化查看颜色。

---

已编辑 - 添加代码以进行详细说明

---

下面是我加载图像然后试图找到最主要的颜色的代码。

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

您可以看到 Kmeans 正确地将蓝色识别为最主要的颜色。

【讨论】：

您好，您可以查看已编辑的部分吗，我尝试添加集群方法

嗨@DDank。我觉得你的方法很好。 .为了清楚起见，我添加了一些代码。

【参考方案2】：

有人告诉我，当惯性改进下降到 20% 以下时，您首先要检查所需的 kmeans 数量，这可以通过以下方式完成：

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

但我猜对你来说这将是不同颜色的数量。我建议你使用那个数字，我假设是 5 查看你的代码。

之后，要预测颜色，您会这样做

preds = kmeans.predict(whateverYouWantToPredict)

那将是您的最终预测，但请注意这是unsupervised method。您可以通过将预测附加到数据帧来使用此预测，然后使用该数据训练 supervised model 以进行另一个预测。

【讨论】：

X_test[0] 的形状是 (519168 , ) 但预测需要二维数组