# -*- coding: utf-8 -*-
"""
Created on Tue Sep 18 21:22:56 2018

@author: zy
"""

'''
k-means聚类 对图像分割
'''
import numpy as np
import cv2


def  k_means(n,k,img):
    '''
    n为迭代次数
    k为聚类数目
    img为输入图片
    '''
    length = img.shape[0]
    img_new = img.copy()
    #按列合并
    img_new = np.column_stack((img_new,np.ones(length)))
    #随机选择初始点
    center_point = np.random.choice(length, k, replace=False)    
    center = img_new[center_point,:]
    distance = [[] for i in range(k)]   # 距离度量
    for i in range(n):
        for j in range(k):
              distance[j] = np.sqrt(np.sum(np.square(img_new - np.array(center[j])), axis=1))     # 更新距离
        img_new[:,3] = np.argmin(np.array(distance), axis=0)   # 将最小距离的类别标签作为当前数据的类别
        for l in range(k):
            center[l] = np.mean(img_new[img_new[:,3]==l], axis=0)   # 更新聚类中心
    return img_new 
    
    
    
if __name__ == '__main__':
    img = cv2.imread('image/book1.jpg')    
    rows,cols = img.shape[:2]
    sacle = 0.2   
    new_rows = int(rows*sacle)
    new_cols = int(cols*sacle)
    print('img',img.shape)
    img = cv2.resize(img,dsize=(new_cols,new_rows))
    data = img.reshape(-1, 3)
    img_new = k_means(100,5,data)
    print(img_new.shape)
    # data_new = np.delete(data_new, 3, axis=1)
    # print(data_new.shape)
    pic_new = img_new[:,3].reshape(new_rows,new_cols)   # 将多个标签展示出来
    cv2.imshow('',pic_new)
    cv2.waitKey(0)
    cv2.destroyAllWindows()