python实现简单KMeans算法

参考链接：北京理工大学公开课

code：

import numpy as np
import PIL.Image as image # 加载PIL包，用于加载创建图片
from sklearn.cluster import KMeans # 加载KMeans算法

# 该函数实现以二进制打开二进制图片，并以列表的形式返回像素值
def loadData(filepath):
    # "rb"表示读取二进制文件
    f = open(filepath,"rb")
    data = []
    img = image.open(f)
    m,n = img.size
    for i in range(m):
        for j in range(n):
            x,y,z = img.getpixel((i,j))
            data.append([x/256.0,y/256.0,z/256.0])
    f.close()
    return np.mat(data),m,n

imageData,row,col = loadData(".//data//bull.jpg")

# 聚类获得每个像素所属的类别
km = KMeans(n_clusters=3)
label =  km.fit_predict(imageData)
label = label.reshape([row,col])

# 创建一张新的灰度图来保存聚类后的结果
# "L"一般表示灰度图
pic_new = image.new("L",(row,col))

# 根据所属类别想图片中添加灰度值
for i in range(row):
    for j in range(col):
        pic_new.putpixel((i,j),int(256/(label[i][j]+1)))

# 以JPEG格式保存图像
pic_new.save("After_segmentation_bull.jpg","JPEG")

Original_image:

Output_image:

The end.