本文为稀土掘金技术社区首发签约文章，30天内禁止转载，30天后未获授权禁止转载，侵权必究！

一、前言

图像分割是计算机视觉中非常重要且基本的任务，在需要应用中都需要使用到图像分割算法。包括自动驾驶、修图、电影特效等。现在有许多成熟的图像分割算法，对于一些简单图像可以使用传统图像处理方法完成分割，而一些复杂场景则需要使用基于深度学习的方法。本文要介绍的是一种机器学习算法的分割方案，即使用KMeans算法完成图像分割。

二、KMeans聚类

2.1 聚类

聚类是一种无监督学习算法，常见的有KMeans、DBSCAN、层次聚类等。聚类的思想就是根据样本的相似度，把某些相似的样本归为一个簇，最终样本会被分成n个簇。在某些算法中，也存在一些无法归为任何一个簇的样本点，这种点被称为离群点 ，或者异常点。

在聚类算法中，距离（相似度）的度量是一个重要问题，通常我们使用几何距离作为度量依据，其计算如下：

<math xmlns="http://www.w3.org/1998/Math/MathML"> d i s t a n c e = x 1 2 − x 2 2 distance = \sqrt{x_1^2 - x_2^2} </math>distance = x12−x22

2.2 KMeans算法

KMeans是一个比较简单的聚类算法，其步骤如下：

确定簇的个数k，随机初始化k个簇中心
计算所有样本与k个簇中心的距离，把样本分配给离它最近的簇中心
根据当前已有的簇，选择簇的中心作为新的簇中心
重复2-4步骤，直到3中簇中心不再更新

KMeans算法可以用下图表示：

2.3 KMeans的实现

在sklearn模块中，提供了KMeans的实现，我们可以直接调用。具体代码如下：

ini 复制代码

from sklearn.cluster import KMeans
from sklearn.datasets import make_blobs
import matplotlib.pyplot as plt
plt.style.use("ggplot")
# 生成数据
X, y = make_blobs(300, 2, centers=3)
# 构建模型
kmeans = KMeans(n_clusters=3)
# 训练
kmeans.fit(X)
# 预测
preds = kmeans.predict(X)
# 绘制结果
plt.subplot(121)
plt.title("origin")
plt.scatter(X[:, 0], X[:, 1], c=y)
plt.subplot(122)
plt.title("cluster")
plt.scatter(X[:, 0], X[:, 1], c=preds)
plt.show()

绘制结果如下：

其中左边是原本的类别情况，右边是聚类结果。在聚类完成后，可以获取各个簇中心，代码如下：

复制代码

kmeans.cluster_centers_

在后面我们会用到簇中心。

三、图像分割

图像分割就是将图像按照一些要求分割成不同的部分。比如人像分割就是把背景和人像分割开，又或者语义分割是将不同含义的内容分割开。图像分割可以被认为是对图片的每个像素进行分类。这一点与前面的聚类有一些相似的地方。

在聚类中，我们是把每个样本归为一个簇，如果把前景、背景各看作一个簇，那么聚类就可以看作是把前面和背景分割的操作。现在的问题就是如何把图片作为输入。

我们可以把图片的每个像素看作一个样本，把图片从h×w×c（高、宽、通道数）转换成size×c，其中c为通道数，即颜色相关的维度。接下来只需要把图片输入聚类算法即可，分割的代码如下：

ini 复制代码

import cv2
from sklearn.cluster import KMeans
import matplotlib.pyplot as plt

# 读取图片
origin = cv2.imread('img.png')
h, w, c = origin.shape
origin = cv2.cvtColor(origin, cv2.COLOR_BGR2RGB)
img = origin.copy().reshape(-1, c)
# 构建模型
kmeans = KMeans(n_clusters=2)
# 训练
kmeans.fit(img)
# 预测
preds = kmeans.predict(img)
plt.figure(figsize=(10, 5))
plt.subplot(121)
plt.title("origin")
plt.imshow(origin)
plt.subplot(122)
plt.title("cluster")
plt.imshow(preds.reshape((h, w)))
plt.show()

上面代码我们把图片的像素作为样本，进行聚类，这样我们就可以根据像素的相似度进行图片的分割操作了。下面是分割结果：

其中左边是原图，由花和叶组成，可以看作前景和背景。右边是分割结果，把两个部分分割开来了。

因为我们使用的是KMeans算法，因此可以得到两个簇的簇中心，并且在这个例子中我们使用颜色作为聚类的样本，因此可以想到簇中心其实就是各个簇的平均颜色。于是我们使用簇中心作为区域的颜色对上面的结果重新绘制，代码如下：

scss 复制代码

# 预测
preds = kmeans.predict(img)
plt.figure(figsize=(10, 5))
plt.subplot(121)
plt.title("origin")
plt.imshow(origin)
plt.subplot(122)
plt.title("cluster")
result = np.zeros_like(origin)
result[(preds == 0).reshape(h, w), :] = kmeans.cluster_centers_[0]
result[(preds != 0).reshape(h, w), :] = kmeans.cluster_centers_[1]
plt.imshow(result)
plt.show()

这里是修改预测部分的代码，生成结果如下：

图片被正确分割开来。

四、总结

图像分割算法有很多，聚类是比较简单的一种。在本文的例子中，我们使用颜色作为分割依据，这种方法有几个确定。比如需要提前确定分割的数量、复杂场景的分割效果很差等。

另外这种方式还有一个非常问题，就是在分割时没有考虑位置信息。在前面的例子中，每个像素都是一个样本点，而实际在分割是周围像素对分割的决定也起了非常重要作用，而这是上述算法欠缺的。一直想法是混入位置信息，这里留给读者自己思考。另外一种方法则是使用卷积神经网络，这样就可以很好解决这一问题。后续我们会讨论更复杂的图像分割实现。

简单场景的图像分割算法