本文为稀土掘金技术社区首发签约文章,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()
这里是修改预测部分的代码,生成结果如下:
图片被正确分割开来。
四、总结
图像分割算法有很多,聚类是比较简单的一种。在本文的例子中,我们使用颜色作为分割依据,这种方法有几个确定。比如需要提前确定分割的数量、复杂场景的分割效果很差等。
另外这种方式还有一个非常问题,就是在分割时没有考虑位置信息。在前面的例子中,每个像素都是一个样本点,而实际在分割是周围像素对分割的决定也起了非常重要作用,而这是上述算法欠缺的。一直想法是混入位置信息,这里留给读者自己思考。另外一种方法则是使用卷积神经网络,这样就可以很好解决这一问题。后续我们会讨论更复杂的图像分割实现。