from numpy.core.fromnumeric import reshape
from skimage import io
from sklearn import cluster
from sklearn.cluster import KMeans
import numpy as np
#第一步读取图像
image = io.imread("ladybug.png")
# io.imshow(image)
# io.show()
#第二步取出长和宽
rows = image.shape[0]
cols = image.shape[1]
# print("rows:"+int(rows))
#第三步选取reshape，选取kmeans进行训练（fit）
#reshape成为一个二维样本 ,3为三个通道
image = image.reshape(image.shape[0]*image.shape[1],3)
kmeans = KMeans(n_clusters=128,n_init=10,max_iter=200) #分为128个簇，使用cpu数量，最大迭代次数
print("!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!")
kmeans.fit(image) #进行训练
#第四步将labels图像转为灰度图
clusters = np.asarray(kmeans.cluster_centers_,dtype=np.uint8)
labels = np.asarray(kmeans.labels_,dtype=np.uint8)
labels = labels.reshape(rows,cols) #转化为灰度图
print("$$$$$$$$$$$$$$$$$$$$$$$$$$$$$")
print(clusters.shape)
np.save("codebook_test.npy",clusters)
io.imsave("compressed_test.jpg",labels)
image1 = io.imread("compressed_test.jpg")
io.imshow(image1)
io.show()
print("**********************************")
