用K最近邻色彩压缩

# -*- coding: UTF-8 -*-
from sklearn.datasets import load_sample_image
import matplotlib.pyplot as plt
import numpy as np
from sklearn.cluster import MiniBatchKMeans


def plot_pixels(data, title, colors=None, N=10000):  # 前1w个像素
    if colors is None:
        colors = data
    # 随机选择子集
    rng = np.random.RandomState(0)

    i = rng.permutation(data.shape[0])[:N]  # 随机选1万个序号
    colors = colors[i]  # 随机选1w个像素点
    # print(colors)
    R, G, B = data[i].T  # 把R，G，B的值们找出来
    fig, ax = plt.subplots(1, 2, figsize=(16, 6))
    ax[0].scatter(R, G, color=colors, marker='.')
    ax[0].set(xlabel='Red', ylabel='Green', xlim=(0, 1), ylim=(0, 1))
    ax[1].scatter(R, B, color=colors, marker='.')
    ax[1].set(xlabel='Red', ylabel='Blue', xlim=(0, 1), ylim=(0, 1))
    fig.suptitle(title, size=20)


china = load_sample_image("china.jpg")
# ax1 = plt.axes(xticks=[], yticks=[])
# ax1.imshow(china)
print('china.shape: ', china.shape)
# print('china: ', china)
data = china / 255.0  # 化为0-1区间
data = data.reshape(427 * 640, 3)  # 化成n_samplesXn_features
print('data.shape: ', data.shape)
# plot_pixels(data, title='Input color space: 16 million possible colors')  # RGB颜色空间中的像素分布
kmeans = MiniBatchKMeans(16)  # 更快的算法计算数据子集
kmeans.fit(data)
print('kmeans.predict(data).shape:', kmeans.predict(data).shape)  # 降维
# print('kmeans.predict(data):', kmeans.predict(data))  # 降维
# 降维后的kmeans.predict(data)这个列表的各个数字代表kmeans.cluster_centers_簇中心（一共16个色簇）的需要提出来拼接的行索引
new_colors = kmeans.cluster_centers_[kmeans.predict(data)]
print('kmeans.cluster_centers_.shape:', kmeans.cluster_centers_.shape)  # 簇中心，这16个代表色彩

print('new_colors.shape:', new_colors.shape)
plot_pixels(data, colors=new_colors, title="Reduced color space: 16 colors")

china_recolored = new_colors.reshape(china.shape)  # 在原来的图像空间china.shape：（427*640*3）上还原成原始尺寸
fig2, ax2 = plt.subplots(1, 2, figsize=(16, 6), subplot_kw=dict(xticks=[], yticks=[]))
fig2.subplots_adjust(wspace=0.05)
ax2[0].imshow(china)
ax2[0].set_title('Original Image', size=16)
ax2[1].imshow(china_recolored)
ax2[1].set_title('16-color Image', size=16)
plt.show()
"""
# 按照b的顺序用a的行来组成新的矩阵
a = np.mat([[1, 2, 3],
            [4, 5, 6]])
b = np.array([1, 1, 0, 0, 0, 1, 0])
c = a[b]
print('c:\n', c)
"""