Python实现图像压缩K-Means时质心出现NaN的排查与修复

问题描述 #

我正在完成一项作业，需手动实现K-Means聚类算法并用于Python图像压缩。运行一次迭代后，重新计算的质心常出现一行或多行NaN值。我初步猜测是无数据点分配到聚类导致除零错误，但不确定是否如此，也不知如何解决。

我尝试过调整聚类数量，仍出现NaN。代码改编自同作业中可行的K-Means实现，但不确定适配时是否出错。我已通过打印检查各步骤输入输出，未发现明显问题，且无法使用现成K-Means函数。

我的代码如下：

# Import the necessary libraries
import numpy as np
import os
from os.path import abspath, exists
from PIL import Image
from numpy import asarray
from scipy.sparse import csc_matrix, find

# Read in the data
dirpath = os.getcwd()
image_path = dirpath + '//data/image.bmp'
image_1 = Image.open(image_path)
image_1_array = asarray(image_1)
image_1_array = image_1_array.reshape(-1, image_1_array.shape[-1])

# Randomly initialize cluster centers
k_clusters = 20
centroids = image_1_array[np.random.randint(image_1_array.shape[0], size=(1, k_clusters))[0]]

iterations = 1

for i in range(0,iterations):
    centroids_squared = np.sum(np.power(centroids, 2), axis=1, keepdims=True)
    print(centroids)
    print(centroids_squared)

    # Calculate the difference between data points and centrouds, make assignments
    tmpdiff = (2 * np.dot(centroids, image_1_array.T) - centroids_squared)
    labels = np.argmax(tmpdiff, axis=0)

    # Update centroids
    dp_num = image_1_array.shape[0]
    P = csc_matrix((np.ones(dp_num), (labels, np.arange(0, dp_num, 1))), shape=(k_clusters,dp_num))

    # Count the data points in each cluster center.
    count = P.sum(axis=1)

    # Adjust cluster centers
    centroids = np.array((P.dot(image_1_array) / count))

问题根源与修复方案 #

你的猜测完全正确：NaN的出现就是因为部分聚类没有分配到任何数据点，导致count中对应的值为0，执行除法时触发除零错误，得到NaN。

具体修复步骤 #

1. 检测空聚类：在更新质心前，找出count中值为0的聚类索引。
2. 重新初始化空聚类：对于这些空聚类，从原始图像数据中随机选取新的质心，避免后续迭代继续出错。

修改后的核心代码如下：

# Adjust cluster centers
count = P.sum(axis=1)
# 找出空聚类的索引
empty_clusters = np.where(count == 0)[0]

if len(empty_clusters) > 0:
    # 从原始数据中随机选择新的质心替换空聚类
    new_centroids = image_1_array[np.random.randint(image_1_array.shape[0], size=len(empty_clusters))]
    centroids[empty_clusters] = new_centroids
else:
    # 正常计算质心
    centroids = np.array((P.dot(image_1_array) / count))

补充说明 #

• 调整聚类数量仍出现问题的原因：随机初始化的质心可能刚好落在图像中完全不存在的颜色区域，导致没有像素被分配到该聚类。
• 若想进一步降低空聚类概率，可采用K-Means++初始化（从数据集中选第一个质心，后续质心优先选离现有质心最远的点），但作业中如果允许随机初始化，上述修复已足够解决问题。

内容的提问来源于stack exchange，提问作者DiYage