import numpy as np
from sklearn.cluster import KMeans
from PIL import Image
import matplotlib.pyplot as plt
from copy import deepcopy


# 打开并读取原始图像中像素颜色值，转换为三维数组
imOrigin = Image.open('颜色压缩测试图像.jpg')
dataOrigin = np.array(imOrigin)
# 然后再转换为二维数组，-1表示自动计算该维度的大小
data = dataOrigin.reshape(-1,3)

n_clusters = 64
while n_clusters >= 4:
    plt.cla()
    print(n_clusters)    
    # 使用KMeans算法把所有像素的颜色值划分为4类
    kmeansPredicter = KMeans(n_clusters=n_clusters)
    kmeansPredicter.fit(data)
    # 使用每个像素所属类的中心值替换该像素的颜色
    # temp中存放每个数据所属类的标签
    temp = kmeansPredicter.labels_
    dataNew = kmeansPredicter.cluster_centers_[temp]
    data = deepcopy(dataNew)
    dataNew = np.uint8(dataNew)    
    dataNew.shape = dataOrigin.shape
    plt.imshow(dataNew)
    plt.pause(0.01)
    n_clusters //= 2