kmeans聚类实现图像分割

#!/usr/bin/env python
# encoding: utf-8
'''
@author: lele Ye
@contact: [email protected]
@software: pycharm 2018.2
@file: kmeans.py
@time: 2019/1/7 19:36
@desc:基于聚类的图像分割
'''
import numpy as np
from PIL import Image

from sklearn.cluster import KMeans


def lodaData(filePath):
    '''
    加载图片数据
    :param filePath: 图片路径
    :return:
    '''
    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 = map(lambda x:x/256.0, img.getpixel((i, j)))
            data.append([x, y, z])
    f.close()

    return np.mat(data), m, n


imgData, row, col = lodaData("bull.png")
# 图片当中的颜色（包括背景共有）3类
km = KMeans(n_clusters=3)
# 聚类获得每个像素所属的类别
label = km.fit_predict(imgData)
label = label.reshape([row, col])
# 创建一张新的灰度图以保存聚类后的结果
pic_new = Image.new("RGB", (row, col))
# 根据类别向图片中添加像素值
for i in range(row):
    for j in range(col):
        # 如果属于背景的类别，填充红色
        if label[i][j] == 0:
            pic_new.putpixel((i, j), (202, 12, 22))
        # 输入白色的牛角和文本类别，填充绿色
        elif label[i][j] == 1:
            pic_new.putpixel((i, j), (29, 209, 107))
        # 属于第三类的牛头类别填充蓝色
        else:
            pic_new.putpixel((i, j), (53, 106, 195))
# 以JPEG格式保存图像
pic_new.save("result-bull-4.jpg", "JPEG")

原始图像：

分割以后的结果：