基于聚类算法 K-means的图像分割

Posted 2021-09-09 404detective

import numpy as np
from PIL import Image
from sklearn.cluster import KMeans

def lodaData(filePath):
    f = open(filePath, 'rb')
    data = []
    img = Image.open(f)
    m, n = img.size
    for i in range(m):
        for j in range(n):
            #注意PNG图片有4个通道来接收，JPG为3个，
            #避免报错随便用个变量接受即可
            x, y, z ,a= 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("图片路径")
# 图片当中的颜色（包括背景共有）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), (255, 0, 0))
        # 输入白色的牛角和文本类别，填充绿色
        elif label[i][j] == 1:
            pic_new.putpixel((i, j), (0, 255, 0))
        # 属于第三类的牛头类别填充蓝色
        else:
            pic_new.putpixel((i, j), (0, 0, 255))
# 以JPEG/PNG格式保存图像
pic_new.save("生成后保存新的图片路径", "PNG")

效果对比如下：