使用 K 均值进行颜色量化

简介

颜色量化是一个在保留图像整体外观的同时减少图像中使用的不同颜色数量的过程。这是通过将相似颜色聚类在一起并用单个颜色值表示它们来实现的。在本实验中，我们将使用 K 均值聚类算法对中国颐和园的图像进行颜色量化。

虚拟机使用提示

虚拟机启动完成后，点击左上角切换到“笔记本”标签以访问 Jupyter Notebook 进行练习。

有时，你可能需要等待几秒钟让 Jupyter Notebook 完成加载。由于 Jupyter Notebook 的限制，操作验证无法自动化。

如果你在学习过程中遇到问题，请随时向 Labby 提问。课程结束后提供反馈，我们将立即为你解决问题。

Skills Graph

加载并显示原始图像

我们将从加载并显示颐和园的原始图像开始。

import numpy as np
import matplotlib.pyplot as plt
from sklearn.datasets import load_sample_image

## 加载颐和园照片
china = load_sample_image("china.jpg")

## 显示原始图像
plt.figure()
plt.axis("off")
plt.title("原始图像")
plt.imshow(china)
plt.show()

将图像转换为浮点数并重塑形状

我们将把图像转换为浮点数，并将其重塑为二维 numpy 数组，以便能够由 K 均值算法进行处理。

## 转换为浮点数，而不是默认的 8 位整数编码。
china = np.array(china, dtype=np.float64) / 255

## 获取图像的尺寸
w, h, d = original_shape = tuple(china.shape)
assert d == 3

## 将图像重塑为二维 numpy 数组
image_array = np.reshape(china, (w * h, d))

拟合 K 均值模型

我们将在图像数据的一个小子样本上拟合 K 均值模型，并使用它来预测整个图像的颜色索引。

from sklearn.cluster import KMeans
from sklearn.utils import shuffle
from time import time

n_colors = 64

## 在数据的一个小子样本上拟合 K 均值模型
print("在数据的一个小子样本上拟合模型")
t0 = time()
image_array_sample = shuffle(image_array, random_state=0, n_samples=1000)
kmeans = KMeans(n_clusters=n_colors, n_init="auto", random_state=0).fit(
    image_array_sample
)
print(f"完成于 {time() - t0:0.3f} 秒。")

## 获取所有点的标签
print("预测整个图像的颜色索引 (k 均值)")
t0 = time()
labels = kmeans.predict(image_array)
print(f"完成于 {time() - t0:0.3f} 秒。")

使用随机码本预测颜色索引

我们将使用随机码本预测整个图像的颜色索引。

## 获取一个随机码本
codebook_random = shuffle(image_array, random_state=0, n_samples=n_colors)

## 使用随机码本预测整个图像的颜色索引
print("预测整个图像的颜色索引 (随机)")
t0 = time()
labels_random = pairwise_distances_argmin(codebook_random, image_array, axis=0)
print(f"完成于 {time() - t0:0.3f} 秒。")

重建图像

我们将使用从 K 均值模型和随机码本获得的码本和标签来重建压缩图像。

def recreate_image(codebook, labels, w, h):
    """从码本和标签重建（压缩）图像"""
    return codebook[labels].reshape(w, h, -1)

## 将原始图像与量化后的图像并排显示
plt.figure()
plt.clf()
plt.axis("off")
plt.title("原始图像 (96,615 种颜色)")
plt.imshow(china)

plt.figure()
plt.clf()
plt.axis("off")
plt.title(f"量化图像 ({n_colors} 种颜色，K 均值)")
plt.imshow(recreate_image(kmeans.cluster_centers_, labels, w, h))

plt.figure()
plt.clf()
plt.axis("off")
plt.title(f"量化图像 ({n_colors} 种颜色，随机)")
plt.imshow(recreate_image(codebook_random, labels_random, w, h))

plt.show()

总结

在本实验中，我们学习了如何使用 K 均值聚类算法对图像进行颜色量化。我们加载并显示了原始图像，将其转换为浮点数并重塑为二维 numpy 数组。我们在图像数据的一个子样本上拟合了 K 均值模型，并使用它来预测整个图像的颜色索引。我们还使用随机码本预测了整个图像的颜色索引。最后，我们使用从 K 均值模型和随机码本获得的码本和标签重建了压缩图像，并将原始图像与量化后的图像并排显示。