使用 K-Means 聚类进行图像分割

原文：towardsdatascience.com/image-segmentation-with-k-means-clustering-1bc53601f033

你可以在这里查看这个项目的笔记本 here

https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/dd72df59bc4201bdd2fffd5cf9c528df.pnghttps://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/0a8b43d044b33bb9bae13ba3b85d7a02.png

LEFT: 原始照片，RIGHT: 分割图像（5 种颜色/段）

观察上面的图像，我们看到一个图像色调分离过滤器，它给图像带来卡通般的视觉效果，但幕后，这个过滤器实际上正在使用一种称为聚类的机器学习算法。

在探索这个过程的原理并了解我们如何在 Python 中实现它之前，让我们先看看为什么我们可能想要这样做。

图像分割

在一张普通照片中，一个像素可以呈现大约 16.7 百万种不同的颜色。然而，在这张处理过的图像中，却只有 5 种不同的颜色。我们已经将所有像素分成 5 个不同的组，将图像分割成这些不同的颜色区域。

我们还减少了图像中的噪声和变化量。因此，如果这要在某些其他机器学习应用中使用，我们已经大大减少了需要处理的数据量，尤其是如果这应用于整个图像库的话。

尽管我们简化了这张图像，但我们仍然保留了大部分重要的结构数据。我们仍然能够识别形状和形式、阴影和亮部，以及许多不同的质感和图案。所有告诉我们这是一张鸟儿站在墙上的图片的信息仍然存在。

这个过程不仅适用于图像。我们可以在其他不同上下文中的各种数据集上使用相同的工具，简化它们并减少所需的数据处理量。

除了减少数据处理需求之外，这个算法还有更多直接的应用，因为它使我们能够更容易地识别图像中的特征。例如，下面图像左侧的湖几乎完全是单一像素颜色，与湖周围的所有像素都不同，这使得它很容易被隔离。

这种技术可以用来追踪极地冰盖的大小；通过获取一张图像，应用这个过滤器来识别冰和海洋像素，然后你可以轻松地计算出被冰覆盖的区域。

https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/3afcac1a3545fa830603dc844ab7a435.pnghttps://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/06fda1610e8e7854d44f7a6627af0bd8.png

LEFT: 原始照片，RIGHT: 分割图像（5 种颜色）

聚类

支持此过程的机器学习技术被称为聚类。这是一个用于在数据集中识别簇或组的过程。例如，对于图像，我们识别出颜色相似的像素组。我们希望找到的组数是算法的一个参数。

聚类主要有三种不同类型：基于密度的、基于质心的和层次化的，每种类型都有许多不同的算法可以根据情况使用。对于这篇文章，我们将实现一种称为K-Means 聚类的基于质心的算法。

https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/773897c7b3d8b32db2bc8e45cd02216d.png

展示聚类算法结果的插图

我们将使用这个聚类算法对由图像数据组成的集合进行操作。图像中的每个像素都可以用 3 个值来表示，这些值对应于每个颜色通道的强度：红色、绿色和蓝色。每个强度值将在 0 到 255 之间。

如果我们取一个像素值，比如[103, 56, 213]，这是一种明亮的紫色，我们可以在一组轴上将其绘制为坐标点（如下所示）。如果我们对图像中的所有像素重复此操作，我们可能会发现在生成的图中开始形成“点群”。所有蓝色像素将彼此靠近（因为它们都有相似的颜色值），绿色像素也是如此，以及黄色像素。结果可能看起来像这样：

https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/2ec5c70ec1d1f729b4f32ac848df3c6f.png

在三维空间中绘制的像素值

虽然从这个固定的视角来看并不容易看出，但很明显，在数据中存在一些簇，其中点密集地聚集在一起。

K-Means 聚类

让我们深入了解 K-Means 聚类是如何工作的。

我们首先选择一个k的值，这将是我们想要将数据分组成的簇的数量。选择正确的值是一个重要的步骤，因为它可以大大影响算法的成功。如果你碰巧知道你的数据点属于多少个不同的类别，那么你应该选择这个值。如果不是，有各种方法（如肘部方法）可以识别k的最佳值。目标是减少簇内变异（簇内的点应该彼此靠近），并增加簇间变异（簇本身尽可能明显）。

对于这个应用，处理图像时，很容易看出所选值如何影响结果。

https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/d8d9b5050890d0f141dffa5d19be3a26.png

不同 k 值的示例：2，3，5，7（从左到右，从上到下）

在我们选择了k的值之后，下一步是设置我们的初始质心点这些代表每个簇的中心。我们为每个质心随机选择一组坐标，在可能的值范围内。

然后重复以下步骤：

计算每个像素与每个质心之间的距离
将每个像素分配给最近的质心
对于每个质心，找到所有分配像素的平均值，这应该对每个维度分别进行。
通过更新它们的坐标到这些平均值来“移动”质心

这会一直重复，直到质心停止移动。这表明簇没有进一步的变化，每个像素都被分配到最近的簇，质心处于最佳位置。

现在我们将用 Python 实现这个算法的简单版本。正如我们在测试中将会看到的那样，这不是一个计算效率高的方法，但它是一个非常直接的方法。

在以后，我们将介绍一种更有效的方法，使用向量计算。这将使我们能够处理更高分辨率的图像，显著减少执行时间。

Python 算法 #1

我们将首先导入我们的图像并将其转换为像素数组。

img=cv2.imread("bird-small.jpg")# swap colour channels so that it is displayed correctlyimg=cv2.cvtColor(img,cv2.COLOR_BGR2RGB)# display imageplt.axis('off')plt.imshow(img);# split into channelsr,g,b=cv2.split(img2)# reformat pixel colours into 1xn arrayr=np.array(r.flatten())g=np.array(g.flatten())b=np.array(b.flatten())

这张图像的维度是 (267, 400, 3) ，即，高度 = 267，宽度 = 400，通道 = 3。我们将通道分成 3 个数组，然后将这些数组重新格式化为 1xn 数组（一长串像素值而不是像素网格）。

例如，红色通道现在看起来是这样的：

[130, 130, 130, ..., 145, 156, 159]包含 267*400 = 106,800 个值

现在我们定义一个距离函数，它将用于计算像素和质心之间的距离。在这个例子中，我们将使用欧几里得距离，如下公式所示：

distance = sqrt( (b0-a0)² + (b1-a1)² + (b2-a2)² )，其中a和b是两个像素，而a0, a1, a2, b0, b1, b2是每个像素的颜色分量。

# euclidean distancedefdist(a,b):returnnp.sqrt((a[0]-b[0])**2+(a[1]-b[1])**2+(a[2]-b[2])**2)

接下来，我们将设置初始质心位置，并为此方法定义一些参数。

# number of clusters/centroidsk=4# random initial starting points within range 0, 255centroids=np.array([[random.randint(0,255),random.randint(0,255),random.randint(0,255)]foriinrange(k)])# number of pixels in the imagepixels=len(r)# exit conditions - stop repeating when max iterations have been reached# or the centroids stop movingmax_iter=8moved=True

现在我们可以构建算法的主体部分。

iter=0whilemovedanditer<=max_iter:iter+=1# cluster assignments, placeholder arrayassignment=[0]*pixels# for each pixelforiinrange(pixels):# compute distance between each pixel and each centroiddistances=[0]*kforjinrange(k):distances[j]=dist(centroids[j],[r[i],g[i],b[i]])# find minimum distance, returns index (0, .., k-1) of nearest centroidnearest=np.argmin(distances)# will look something like:# [0, 0, 1, 0, 1, 2, 2, 0, ...]# with a centroid value assigned to each pixelassignment[i]=nearest prev_centroids=centroids.copy()# for each cluster, calculate mean of allocated points for each dimensionforiinrange(k):# list of array indices of pixels that belong to each clusterind=[jforjinrange(pixels)ifassignment[j]==i]# check cluster assignment is not empty# prevents divide by zero error when calculating meaniflen(ind)!=0:centroids[i][0]=np.mean(r[ind])centroids[i][1]=np.mean(g[ind])centroids[i][2]=np.mean(b[ind])else:centroids[i][0]=0centroids[i][1]=0centroids[i][2]=0# check if centroids have movedifnp.array_equal(centroids,prev_centroids):moved=False

算法完成后，我们现在可以重建图像，将每个像素分配给其簇中心的颜色以产生所需的效果。

# make copy of colour channelsr_copy=np.array(r.copy())g_copy=np.array(g.copy())b_copy=np.array(b.copy())# update pixels to be the colour of their clusterforiinrange(k):ind=[jforjinrange(pixels)ifassignment[j]==i]r_copy[ind]=centroids[i][0]g_copy[ind]=centroids[i][1]b_copy[ind]=centroids[i][2]# compile channelsimg2=np.array([r_copy,g_copy,b_copy])# transpose to group values into pixelsimg2=img2.transpose()# reshape list of pixels into height x widgh x channelsimg2=img2.reshape(img.shape)plt.axis('off')plt.imshow(img2)plt.savefig("bird-small(k4).png",format="png",dpi=600)

测试

如我们所见，图像已被减少到只有 4 种不同的颜色。然而，正如预期的那样，这是一个非常低效的实现，运行时间大约为 30 秒，而且这只是一个低分辨率图像。

//k=4Initial Random Centroids:[[8211780][134182172][214107227][15545225]]Final Centroids:[[837962][13012798][178171127][333131]]

https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/415cfddc200a05d4322f0c6668b0faf4.pnghttps://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/e0d9fe16e27df8cdbe38c80a34f89edf.png

左侧：聚类前，右侧：聚类后（k=4）

Python 算法 #2（使用 NumPy 优化）

我们进行与之前类似的设置，除了我们将图像保持为 3D 数组而不是将其拆分为通道。我们还重写了距离函数，使其能够处理值的向量而不是单个点。这允许更高效的执行，因为一些操作可以并行进行。

img=import_image("corfu.jpg")# split image into channels, reformat h x w x c structureimg=np.array(cv2.split(img))img=img.transpose(1,2,0)# exit conditionsmax_iter=10iter=0moved=True# number of clustersk=7# initial cluster centresclusters=[[random.randint(0,255)foriinrange(3)]forjinrange(k)]# define distance functiondefdist(a,b):return(np.sqrt(np.sum((b-a)**2,2)))

算法的主要部分遵循类似的架构，除了我们使用不同的编码方法来跟踪哪些像素对应于哪个簇，并使用 2D 掩码来索引图像，从而允许更快的访问。

iter=0whileiter<=max_iterandmoved==True:iter+=1# calculate distance between pixels and cluster, for every clusterdistances=[dist(img,clusters[i])foriinrange(k)]# index (0, ..., k) of the nearest cluster centre for each pixel# produces an array the same shape as the image, instead of pixels,# it stores in the index of the nearest cluster# this can be used as a mask later onnearest=np.argmin(distances,0)prev_clusters=clusters.copy()foriinrange(k):# create 1-hot encoded mask of which pixels belong to the clusterind=np.array(np.where(nearest==i,1,0),dtype=bool)# apply mask to image to extract subset of pixelssubset=img[ind]# calculate mean of the identified subset - update cluster centresclusters[i]=[np.round(np.mean(subset[:,0])),np.round(np.mean(subset[:,1])),np.round(np.mean(subset[:,2]))]# remove NaN values - replace with 0ifnp.isnan(clusters[i][0]):clusters[i][0]=0ifnp.isnan(clusters[i][1]):clusters[i][1]=0ifnp.isnan(clusters[i][2]):clusters[i][2]=0ifclusters==prev_clusters:moved=False

现在来显示图像。

# After the final iteration, the cluster centres represent the pixel colour# of each cluster. We apply the final version of the array, nearest, as a# mask to sample colours for each pixelclusters=np.array(clusters,dtype=int)img2=clusters[nearest]# display imageplt.axis('off')plt.imshow(img2)plt.savefig("corfu(k7).png",format="png",dpi=600)

测试

使用之前的图像测试这个版本，我们看到有显著的改进，总执行时间不到 1 秒。这现在使我们能够尝试更大的图像，并且具有更高的max_iter截止点（可能带来更准确的结果）。

下面的图像具有（2003，3456）的维度，给它提供了 6,922,368 个像素。这几乎是之前的 65 倍（106,800 个像素）。

使用这种新的实现，图像的处理时间大约为 40 秒（对于相同的k值）。因此，使用之前的实现处理这个相同的、新的图像，将需要超过 15 分钟。

（值得注意的是，由于代码的结构，性能会受到k选择的影响而略有影响）。

https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/e374354f5e7d7c0c3cbde0abbf7fb665.pnghttps://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/8db5ee30a37922d7892cd74cf37b3c71.png

左：原始图像，右：聚类图像（k=7）