K-means 寻找图像主体颜色

编辑栏

几个月前，我写了篇文章“纯python编写K-means算法，提取图片中的主体颜色”，找到图片中的显著颜色。今天突发奇想，能不能把k-means用于图片压缩呢？对图像rgb三个通道进行聚类，得到中心点像素，并填充对应的类别区域，我找来了一张上海的图片进行压缩

我的代码如下，由于在jupyte编辑器上写的代码，因此显示图片使用的matplotlib模块

import numpy as np
import matplotlib.pyplot as plt
from skimage import io
import cv2 as cv
import os
from sklearn.cluster import KMeans

img = io.imread('shanghai.jpg')
img = np.array(img, dtype=np.float64) / 255

plt.axis('off')
plt.imshow(img)

outImgPath = '/data/githubFile/outImgBlog'
if not os.path.exists(outImgPath):
    os.makedirs(outImgPath)


# 定义函数完成kemans功能
def KMeansImage(img, k):
    w, h, c = img.shape
    dd = np.reshape(img, (w * h, c))
    km = KMeans(n_clusters=k)
    km.fit(dd)
    labels = km.predict(dd)
    centers = km.cluster_centers_
    new_img = img.copy()
    for i in range(w):
        for j in range(h):
            ij = i * h + j
            new_img[i][j] = centers[labels[ij]]
    return {'new_image': new_img, 'center_colors': centers}


plt.figure(figsize=(12, 9))
plt.imshow(d)
plt.axis('off')
plt.show()
for i in range(200, 1000, 100):
    print('Number of clusters:', i)
    out = KMeansImage(img=img, k=i)
    centers, new_image = out['center_colors'], out['new_image']
    print('new_image.shape:', new_image.shape)
    plt.figure(figsize=(12, 1))
    plt.imshow([centers]);
    plt.axis('off')
    plt.show()

    # 使用算法跑出的中心点，生成一个矩阵，为数据可视化做准备
    result = []
    result_width = 200
    result_height_per_center = 80

    for center_index in range(i):
        result.append(np.full((result_width * result_height_per_center, centers.shape[1]), centers[center_index] * 255,
                              dtype=int))
    result = np.array(result)
    result = result.reshape((result_height_per_center * i, result_width, centers.shape[1]))

    # # 保存图片
    imgName = 'block-%s.jpg' % i
    cv.imwrite(os.path.join(outImgPath, imgName), result)

    plt.figure(figsize=(12, 9))
    plt.imshow(new_image)
    plt.axis('off')
    plt.show()
    imgName = 'outputImg-%s.jpg' % i
    cv.imwrite(os.path.join(outImgPath, imgName), new_image * 255)