使用 sklearn 库中的 KMeans 实现彩色图像聚类分割

sklearn 库是最为常用且经典的机器学习库，里面封装了许多机器学习算法，此篇文章使用此库中的 KMeans 算法，从而实现图像的聚类分割。

本文不讲理论，只谈应用。

引入库

除了 sklearn 库之外，还需要一些图像处理的库，我引入了如下几个库：

import numpy as np
import pylab as plt
import PIL.Image as image
from sklearn.cluster import KMeans

读入图片

我使用了 pylab 库来读入图片：

img = plt.imread('E:/图片/壁纸/001.jpg')

此时读入的 img 是一个三维 numpy 数组，其形状为 (height, width, 3) ，其中3是指通道数，即 RGB 三个通道。

但是，KMeans 传入的参数必须是二维数组，故，还需要将其打散为二维：

img1 = img.reshape((img.shape[0]*img.shape[1], 3))

构建 `KMeans`

在此构建时，只需要使用最简单的方法即可：

k = 3
kmeans = KMeans(n_clusters=k)

参数有很多，我在构建的时候除了 n_clusters 都使用的默认值：

n_clusters：聚类数，即聚为几类

使用 `KMeans` 聚类

然后，使用 fit() 进行训练：

kmeans.fit(img1)

聚类之后，有很多参数，比较重要的，以及此处需要用到的主要有俩：

kmeans.labels_：聚类的结果，是一个一维 numpy 数组，包含了每一个数据所属的类别。比如，如果聚3类，便可能是：[0, 0, 0, 1, 2, ..., 2, 1]
kmeans.cluster_centers_：迭代完成后的聚类中心，是一个2维的 numpy 数组，形状为 (k, 3) 【k 是聚类数，3通道】

重新填色

聚类完成之后，需要将每个像素点重新填色，将同一类的像素点均填为此类聚类中心的颜色。

在此之前，首先需要得到图片的高度和宽度：

height = img.shape[0] 
width = img.shape[1]

首先用 image.new() 重新创建一个图片，其语法如下：

pic_new = image.new("RGB", (width, height))

第一个参数：一个字符串，表示图片类型。RGB 表示常见的 RGB 彩图；如果是L ，表示为灰度图。
第二个参数，一个长度为2的元组：高宽。

然后需要用 putpixel() 方法来填充像素，但是在此之前，还需要处理几个小细节：

更改数据类型

RGB图中，每个通道都是 0-255 之间的整数，但是，kmeans.cluster_centers_ 中元素类型却是 float64，故在填充之前，还需要小小处理一番，将元素变为 int32 类型的。

center = np.zeros([k, 3])

for i in range(k):
    for j in range(3):
        center[i, j] = kmeans.cluster_centers_[i, j]
center = center.astype(np.int32)

直接转变类型不太合适，因为 kmeans.cluster_centers_ 毕竟是类似于一个属性值的东西，而且这个名字太长，换一个简短的也是好的。故重新复制一份再使用 astype 更改数据类型即可。

恢复被打平的结果

上面便提到，kmeans.labels_ 是一个一维数组，但是图片是二维的，所以将其恢复过来即可：

label = kmeans.labels_.reshape((height, width))

然后便可以填充像素了：

for i in range(height):
    for j in range(width):
        pic_new.putpixel((j, i), tuple((center[label[i][j]])))

这里需要注意 putpixel() 方法，其的两个参数：

第一个参数，是新图片的坐标，如果输入的是 (i, j) 实际上填充的是 (j, i) 处的像素，故需要倒着写^[1]。
第二个参数，在 RGB 图像的时候，需要一个长度为3的元组，故需要tuple() 将其由 numpy 转换为元组。

最后保存图片即可：

pic_new.save("D:/K=3.jpg", "JPEG")

结果展示

我使用了王者荣耀大乔的图片来做测试：

原图

K = 3

K = 4

K = 8

见别人的博客 ↩