图片压缩-KMeans

下面给大家一起分享使用KMeans自动聚类，压缩图片像素点。每种图片可能他们的维度都不同，比如jpg一共有(w,h,3)三维,但是灰度图只有一维（w,h,1）,也有四维的图片(w,h,4)等等。我们可以将这些不同维度的图片映射到一维的空间中。拿三维图片来做解释吧，三维图有rgb三种颜色，首先把每个点r、g、b三个颜色数值提取转换为(w*h,3)，然后使用KMeans训练自动聚类模型，把转换过后的图片矩阵喂入模型得出不同的类别。再找到类别中心点，然后重新创建图片矩阵。代码如下所示：

# coding:utf-8
import matplotlib.pyplot as plt
import numpy as np
from sklearn.cluster import KMeans
from sklearn.utils import shuffle
from sklearn.datasets import load_sample_image
 
imgData = load_sample_image('china.jpg')
 
# imgData = plt.imread('./img/img.png') # ./img/timg.jpg
plt.figure('image1');plt.xticks([]);plt.yticks([])
plt.imshow(imgData)
 
# 压缩图片
def im_resize(imgData,pixel=8):
    w,h,rgb = imgData.shape
    img = 0
    if rgb==3 or rgb == 1:
        img = np.array(imgData,dtype=float)/255
    else:
        img = imgData
 
    img1 = img.reshape(-1,rgb)
    x_train = shuffle(img1)
 
    # 聚类模型建立
    kmeans = KMeans(n_clusters=pixel)
    kmeans.fit(x_train[:2000,:])
    # 分别对原始图片进行分类
    x_class = kmeans.predict(img1)
    # 类别的中心点
    cluster_center = kmeans.cluster_centers_
    image = np.zeros([w,h,rgb])
    n = 0
    for i in range(w):
        for j in range(h):
            index = x_class[n]
            rgbColor = cluster_center[index,:]
            image[i,j] = rgbColor
            n+=1
    return image
 
image_new = im_resize(imgData,pixel=16)
print(image_new.shape)
plt.figure('image2')
plt.imshow(image_new)
plt.xticks([]);plt.yticks([])
plt.show()