【圖像處理5】均值聚類

# -*- coding: utf-8 -*-
import matplotlib.pyplot as plt
import numpy as np
import cv2

"""
k-mean聚類圖像分割算法

1. 這裏FCM 模糊聚類也是相似的，只是將硬分類變成隸屬函數；
2. 但無論是K-mean還是FCM 隸屬函數都很麻煩，這裏用的是灰度值，
基本也就等價於一種自適應的基於閾值的圖像分割了

第一步. 隨機設定一個分類, 初始化C個灰度值
第二步. 計算圖像中所有像素與這個分類的關係，進行分類；
第三步. 將分類後的像素灰度均值作爲新的分類
第四步. 迭代上述三步，直到新的分類點與老的分類點滿足收斂要求

"""
path = '/XXXXX/'
img  = cv2.imread(path + 'lena.jpg')

img_gray       = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
img_gray_array = np.array(img_gray)
img_gray_array = cv2.resize(img_gray_array, (128, 128), interpolation=cv2.INTER_LINEAR)

img_classifier = np.zeros(img_gray_array.shape)

class_num = 3

c_center = np.random.randint(255, size=class_num)

w  = img_gray_array.shape[0]
h  = img_gray_array.shape[1]


def get_min(c1, c2, c3):
    min = c1
    if c2 < min:
        min = c2

    if c3 < min:
        min = c3


    if c1 == min:
        return 0
    elif c2 == min:
        return 1
    else:
        return 2


final_dis     = 999
threshold_dis = 5

while final_dis > threshold_dis:

    #1. 更新分類
    for i in range(w):
        for j in range(h):

            cur_gray = img_gray_array[i][j]
            c1_dis   = abs(cur_gray - c_center[0])
            c2_dis   = abs(cur_gray - c_center[1])
            c3_dis   = abs(cur_gray - c_center[2])

            img_classifier[i][j] = get_min(c1_dis, c2_dis, c3_dis)

    #2. 計算新的聚類點
    c_center_new = np.zeros(c_center.shape);
    final_dis    = 0 

    for t in range(class_num):
        c_center_new[t]  = img_gray_array[np.where(img_classifier == t)].mean()
        final_dis += abs(c_center_new[t] - c_center[t])

    c_center = c_center_new


#繪製結果
for c in range(class_num):
    img_gray_array[np.where(img_classifier == c)] = c_center[c]


plt.imshow(img_gray_array, cmap='gray')
plt.show()