好奇的小老鼠——使用FCM对图片像素进行聚类

import numpy as np
class FCM:
    def __init__(self, K, m=2, eps=0.05):
        # 聚类个数
        self.K = K
        # 加权参数
        self.m = m
        # 终止容限
        self.eps = eps
        # 最优相似度矩阵
        self.U = None
        # 最终的聚类中心
        self.centers = None
        # 目标函数值
        self.obj_val = None
    def init_by_centers(self, data, centers):
        # 初始化隶属度矩阵
        self.U = np.random.random((len(data), self.K))
        # 样本到聚类中心的距离
        dist = np.zeros((len(data), self.K))
        for i, x in enumerate(data):
            for j, c in enumerate(centers):
                dist[i][j] = np.linalg.norm(x - c, 2)
        # 计算新的隶属度矩阵
        for i, x in enumerate(data):
            for j, c in enumerate(centers):
                self.U[i][j] = 1./np.sum((dist[i][j]/dist[i]) ** (2/(self.m-1)))
        # 最初的目标函数值
        self.obj_val = np.sum((self.U ** self.m) * (dist ** 2))
    def train(self, data):
        if len(data.shape) == 1:
            data = data[:, np.newaxis]
        if self.U is None:
            # 初始化隶属度矩阵
            self.U = np.random.random((len(data), self.K))
            # 保证每个样本属于所有类的概率是1
            self.U = np.divide(self.U, np.sum(self.U, axis=1)[:, np.newaxis])
        while True:
            temp_U = self.U ** self.m
            # 计算聚类中心
            self.centers = np.divide(np.dot(temp_U.T, data), np.sum(temp_U.T, axis=1)[:, np.newaxis])
            # 样本到聚类中心的距离
            dist = np.zeros((len(data), self.K))
            for i, x in enumerate(data):
                for j, c in enumerate(self.centers):
                    dist[i][j] = np.linalg.norm(x - c, 2)
            # 计算新的隶属度矩阵
            for i, x in enumerate(data):
                for j, c in enumerate(self.centers):
                    temp_U[i][j] = 1./np.sum((dist[i][j]/dist[i]) ** (2/(self.m-1)))
            # 判断是否收敛
            if np.sum(abs(temp_U-self.U)) < self.eps:
                break
            # 更新隶属度矩阵
            self.U = temp_U
            # 更新目标函数值
            self.obj_val = np.sum((self.U ** self.m) * (dist ** 2))
        # 返回样本最大隶属度对应的类别
        return np.argmax(self.U, axis=1)
    def __str__(self) -> str:
        return f"K：{self.K}, m：{self.m}, eps：{self.eps}, U：{self.U}, centers：{self.centers}, dist: {self.dist}"

from matplotlib import pyplot as plt
import time
import cv2
from fcm import *
# 导入图片
def import_image(img_path):
    # 导入图片并转换成灰度图像,其中的像素值是0-255,uint8类型
    img = cv2.imread(img_path, 0)
    # 图片尺寸
    h, w = img.shape
    img = img.astype('float')
    return img, h, w
def run():
    # 开始计时
    begin = time.time()
    img_path = "./mouse.jpeg"
    # 聚类的数量
    c = 2
    # 导入图像
    img, h, w = import_image(img_path)
    # 去重后的图片像素值
    pixel_data = np.unique(img.flatten())
    # FCM模型
    my_fcm = FCM(K=c, m=2, eps=0.01)
    # 优化
    y = my_fcm.train(pixel_data)
    end = time.time()
    print("总共花费了时间{}".format(end - begin))
    print("像素值聚类中心：{}".format(my_fcm.centers.flatten()))
    print("目标函数值：{}".format(my_fcm.obj_val))
    # 像素值及其对应的类别标签
    pixel_label = dict()
    for d, l in zip(pixel_data, y):
        pixel_label[d] = l
    # 按像素类别生成新的图片
    img_new = np.zeros((h, w))
    # 遍历每一个像素
    for i in range(h):
        for j in range(w):
            img_new[i][j] = int(pixel_label[img[i][j]])
    # 绘制图片
    plt.imshow(img_new, cmap='gray')
    plt.gca().xaxis.set_major_locator(plt.NullLocator())
    plt.gca().yaxis.set_major_locator(plt.NullLocator())
    plt.show()
    # plt.savefig('fcm_mouse.png', dpi=800)
if __name__ == "__main__":
    run()