纯Python编写K-means算法，提取图片中的主体颜色

在文章使用python提取图片中的主体颜色，我使用sklearn快速实现了一个K-means算法，来识别图片的主体颜色。

今天我用Python手写了一个K-means算法，来完成同样的功能。

聚类是数据挖掘中一种非常重要的学习流派，指将未标注的样本数据中相似的分为同一类，正所谓“物以类聚，人以群分”。

K-means是聚类算法中最为简单、高效的，核心思想：由用户指定k个初始中心，以作为聚类的类别，重复迭代更新数据集的中心直至算法收敛。

代码有点长，关键代码我会进行分析，其余代码没必要进行分析。

首先，我们定义一个计算距离的方法，我采用的是欧氏距离。

def get_euclidean_distance(_pixel, _center):
    """
    获得欧氏距离
    :param _pixel: 像素
    :param _center: 中心
    :return: 欧氏距离
    """
    d_pow_2 = 0
    for _channel_index in range(n_channels):
        d_pow_2 += pow(_pixel[_channel_index] - _center[_channel_index], 2)
    return np.sqrt(d_pow_2)

然后定义一个方法来计算中心，方便之后进行迭代。

def get_nearest_center(_pixel):
    """
    获得最近的中心
    :param _pixel: 像素
    :return: 最近的中心的index
    """
    min_center_d = get_euclidean_distance(_pixel, centers[0])
    min_center_index = 0
    for _center_index in range(1, k):
        d = get_euclidean_distance(_pixel, centers[_center_index])
        if d < min_center_d:
            min_center_d = d
            min_center_index = _center_index
    return min_center_index

下面的模块用于计算新的中心，体现了K-means的关键思想。

def cal_new_center():
    """
如果有的中心没有最邻近的，那就说明选到了俩一样颜色的中心（虽然概率很小）那就重新选一个
    :return: 新的中心
    """

    # 归零
    center_counts = np.zeros(k, dtype=int)
    _centers = np.zeros((k, n_channels), dtype=int)

    # 对每个类逐通道求和
    for _pixel_index in range(n_pixels):
        center_counts[labels[_pixel_index]] += 1
        for _channel_index in range(n_channels):
            _centers[labels[_pixel_index]][_channel_index] += img[_pixel_index][_channel_index]

    # 对每个中心算出均值
    for _center_index in range(k):
        if center_counts[_center_index] > 0:
            for _channel_index in range(n_channels):
                _centers[_center_index][_channel_index] /= center_counts[_center_index]
        else:
            # 要是中心选重了就重新再选一个（虽然概率很小）
            _centers[_center_index] = img[random.randint(0, n_pixels - 1)]
            print('WARNING: Center %d has no pixel, re-choose center randomly...' % _center_index)

    return _centers

上面的代码已经完成了K-means的关键部分，其他的代码都是在对图像像素进行遍历处理，以及得到K-means聚类后进行输出。代码已经上传在github上了：https://github.com/shulisiyuan/mainColor

我使用这张图片进行了测试，设置的中心点K=4

下面是代码跑出来的结果，可以看出下面的四种颜色主要就是图片中的颜色。

这种方法跑的比较慢，有很多可以优化的地方，后续有时间再慢慢优化。

打算收集大量图片进行分析，构建一个相似颜色图片的聚类。简单实现Google的图图片颜色聚类功能。