Pyimagesearch CV系列之一：KMeans Color Clustering

一、概述

今天开始，将Pyimagesearch上学习的计算机视觉相关知识整理到博客上。先从一篇简单的开始，使用Python, OpenCV以及sklearn实现颜色聚类，从而找出给定图像中占比最大的颜色。

二、KMeans聚类算法

2.1 算法简介

KMeans算法是一种常用的聚类算法。其目标是将 n 数据点(data point) 划分成k个簇。每一个数据点根据其与各个聚类中心（mean of each cluster）的距离被划分到不同的簇中。同一个簇中的数据点，相比其他簇中的，被认为更加“相似”。

KMeans算法的一个问题在于需要用户事先指定聚类中心的个数。也有一些聚类算法可以自动选择最佳的聚类中心数K（如XMeans），不过不在本post的讨论范围之内。

本任务的思路：给定M×N个像素点，每个像素点有RGB三个特征，将每个像素点视为一个数据点，进行聚类。

2.2 sklearn相关组件

fit: 将KMeans apply到给定数据点的集合：

clt = KMeans(n_clusters=N_CLUSTERS)
clt.fit(data)

clt.labels_：返回每个数据点的标签
即每个数据点经过聚类之后，被划分在哪一个cluster中。标号0,1,2,3等等
clt.cluster_centers_：每个聚类中心的平均度量值
这里的度量值可以理解为，聚类时用于计算的feature，比如这里的RGB三个值。
每个聚类中心的RGB三个值的平均被视为该聚类的中心。

三、实现

实现非常简单，代码如下：

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

def main(img_file, n_clusters):
    img = cv2.imread(img_file)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    plt.subplot(121)
    plt.imshow(img)
    plt.axis('off')
    data = img.reshape(img.shape[0]*img.shape[1], 3)
    clt = KMeans(n_clusters=n_clusters)
    clt.fit(data)
    hist = centroid_histogram(clt)
    bar = plot_colors(hist, clt.cluster_centers_)
    plt.subplot(122)
    plt.imshow(bar)
    plt.axis('off')

其中centroid_histogram和plot_colors是两个辅助函数。定义如下：

import numpy as np

def centroid_histogram(clt):
    '''  生成一个直方图，可视化每个聚类中心的data point数(所占的比例) '''
    numLabels = np.arange(len(np.unique(clt.labels_)) + 1)  #clt.labels_返回所有样本点的聚类标签（即被分在哪一个cluster中）
    hist, _ = np.histogram(clt.labels_, numLabels)
    # 将直方图归一化，从而使其和为一
    hist = hist.astype('float')
    hist /= hist.sum()
    return hist

def plot_colors(hist, centroids):
    bar = np.zeros((50, 300, 3), dtype="uint8")
    startX = 0
    for (percentage, color) in zip(hist, centroids):
        endX = startX + percentage * 300
        cv2.rectangle(bar, (int(startX), 0), (int(endX), 50), color.astype(np.uint8).tolist(), 50)
        startX = endX
    return bar