Python kmeans 图像压缩 demo

## pip install numpy ## pip install matplotlib ## pip install sklearn import os import matplotlib.image as mpimg import matplotlib from sklearn.cluster import KMeans import numpy as np # 给定图像压缩后颜色的数量 N = [2, 8, 16] # 给定输入文件 input_file = os.path.join(os.getcwd(), "madagascar.jpg") # 得到输入文件所在文件夹 input_dir = os.path.dirname(input_file) # 得到输入文件的文件名（供输出使用） input_fn = os.path.splitext(os.path.basename(input_file))[0] # 读入图像，得到RGB通道图像矩阵 img = mpimg.imread(input_file) # 将图像矩阵转换成 像素xRGB通道的 二维矩阵(pixel x color) img2D = np.reshape(img, newshape=[-1, 3]) # 对每个颜色通道做一个循环 for n in N: # 初始化KMeans参数 km = KMeans(n_clusters=n) # 对像素进行聚类，特征为RGB通道 km.fit(img2D) # 得到n个聚类中心（RGB颜色） img_cmap = np.asarray(km.cluster_centers_, np.float32) # 得到每个像素所在的聚类中心标签 labels_reshape = np.asarray(np.reshape(km.labels_, np.shape(img)[0:2]), np.uint8) # 定义输出文件的文件名 output_file = os.path.join(input_dir, input_fn + "_%03dcolor.png" % n) # 将聚类中心转换成图像的颜色索引表 my_cmap = matplotlib.colors.ListedColormap(img_cmap / 255.) # 保存压缩后的图像 mpimg.imsave(output_file, labels_reshape, cmap=my_cmap)