#!/usr/bin/env python
# encoding: utf-8
'''
@author: lele Ye
@contact: 1750112338@qq.com
@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")
原始图像:
分割以后的结果: