K-Means聚类及图像压缩

数据链接:https://pan.baidu.com/s/1cz8SihL2HYh_cFudc7y07Q 密码:tr35

import numpy as np
import cv2
import matplotlib.pyplot as plt 

original_img = cv2.imread('./images/ColorfulBird.jpg')
print('图像维度：', original_img.shape)
cv2.imshow('original_img', original_img)
cv2.waitKey(0)
cv2.destroyAllWindows()
height, width, depth = original_img.shape

# 将图像数据点平铺
# 每个数据点为一个3维的样本
pixel_sample = np.reshape(original_img, (height * width, depth))
print('前10个像素点：')
print(pixel_sample[:10, :])
print(pixel_sample.shape)

from sklearn.cluster import KMeans
# 压缩后图片包含的颜色个数，即为聚类的个数
k = 5
kmeans = KMeans(n_clusters=k, random_state=0)
# 训练模型
kmeans.fit(pixel_sample)
# 找到每个3维像素点对应的聚类中心
cluster_assignments = kmeans.predict(pixel_sample)

# 每个样本聚类的结果是对应的聚类中心的 索引
print(kmeans.cluster_centers_)
print(set(cluster_assignments))
# 聚类中心值
cluster_centers = kmeans.cluster_centers_
print(cluster_centers.shape)
print(cluster_centers)

compressed_img = np.zeros( (height, width, depth), dtype=np.uint8)

# 遍历每个像素点，找到聚类中心对应的像素值
pixel_count = 0
for i in range(height):
    for j in range(width):
        # 获取像素点的聚类中心的索引
        cluster_idx = cluster_assignments[pixel_count]

        # 获取聚类中心索引位置上的像素值
        cluster_value = cluster_centers[cluster_idx]

        # 赋值
        compressed_img[i][j] = cluster_value
        pixel_count += 1

cv2.imwrite('./images/compressed_image.jpg', compressed_img, [cv2.IMWRITE_JPEG_QUALITY, 50])
cv2.imshow('original_img', original_img)
cv2.imshow('compressed_img', compressed_img)
cv2.waitKey(0)
cv2.destroyAllWindows()