使用kmean进行图像分割 使用CRFs进行分割后处理

存在需要解决的问题：

1、kmean是随机聚类的，当设置K值后，每次聚类完后分配的标签也是随机的，如果直接根据聚类后的标签给着色，客户每一次聚类后发现颜色都跟上一次不一样，又开始哔哩哔哩了；

2、可通过设置kmean的参数cv2.KMEANS_USE_INITIAL_LABELS，可一定程度上降低每次运行后标签变动太大，但不能彻底解决这个问题；

3、可通过改变kmean聚类的标签来实现，每次在标签图上着色的颜色位置都是固定的。

简单的方案：（1）对聚类标签的数量进行统计，按占比例大小进行从小到大排序；

                        (2）对原聚类的标签使用排序后的索引进行强制打标签；

存在问题：每次运行kmean聚类后，着色颜色基本上能稳定，但是还是偶尔会变。

 

numpy获取数组下标：

index = np.argsort(ratio)

numpy根据两组数据创建字典：

dic = dict(zip(rt,ratio))

import cv2
import numpy as np

import pydensecrf.densecrf as dcrf
try:
    from cv2 import imread, imwrite
except ImportError:
    # 如果没有安装OpenCV，就是用skimage
    from skimage.io import imread, imsave
    imwrite = imsave
from pydensecrf.utils import unary_from_labels, create_pairwise_bilateral, create_pairwise_gaussian


blue = [255, 20, 20]    #蓝色
green = [50, 255, 50]   #绿色
red = [50, 50, 255]     #红色
light_green = [255, 255, 128] #青色
white = [50, 255, 255] #橙色

#颜色字典
color_dit = np.array([blue, green, red, light_green, white])

#标签着色函数
def decode_segmap(predit_map,color_map,num_class):
    img = predit_map[:, :, 0] if len(predit_map.shape) == 3 else predit_map
    img_out = np.zeros(img.shape +(3,))

    if num_class > len(color_map):
        num_class = len(color_map)
    for i in range(num_class):
        img_out[img == i,:] = color_dit[i]

    return img_out

def CRFs(original_image_path, predicted_image_path, CRF_image_path):
    img = original_image_path

    # 将predicted_image的RGB颜色转换为uint32颜色 0xbbggrr
    anno_rgb = predicted_image_path.astype(np.uint32)
    anno_lbl = anno_rgb[:, :, 0] + (anno_rgb[:, :, 1] << 8) + (anno_rgb[:, :, 2] << 16)

    # 将uint32颜色转换为1,2,...
    colors, labels = np.unique(anno_lbl, return_inverse=True)

    # 如果你的predicted_image里的黑色（0值）不是待分类类别，表示不确定区域，即将分为其他类别
    # 那么就取消注释以下代码
    # HAS_UNK = 0 in colors
    # if HAS_UNK:
    # colors = colors[1:]

    # 创建从predicted_image到32位整数颜色的映射。
    colorize = np.empty((len(colors), 3), np.uint8)
    colorize[:, 0] = (colors & 0x0000FF)
    colorize[:, 1] = (colors & 0x00FF00) >> 8
    colorize[:, 2] = (colors & 0xFF0000) >> 16

    # 计算predicted_image中的类数。
    n_labels = len(set(labels.flat))
    # n_labels = len(set(labels.flat)) - int(HAS_UNK) ##如果有不确定区域，用这一行代码替换上一行

    ###########################
    ###     设置CRF模型     ###
    ###########################
    use_2d = False
    # use_2d = True
    ###########################################################
    ##不是很清楚什么情况用2D
    ##作者说“对于图像，使用此库的最简单方法是使用DenseCRF2D类”
    ##作者还说“DenseCRF类可用于通用（非二维）密集CRF”
    ##但是根据我的测试结果一般情况用DenseCRF比较对
    #########################################################33
    if use_2d:
        # 使用densecrf2d类
        d = dcrf.DenseCRF2D(img.shape[1], img.shape[0], n_labels)

        # 得到一元势（负对数概率）
        U = unary_from_labels(labels, n_labels, gt_prob=0.2, zero_unsure=None)
        # U = unary_from_labels(labels, n_labels, gt_prob=0.2, zero_unsure=HAS_UNK)## 如果有不确定区域，用这一行代码替换上一行
        d.setUnaryEnergy(U)

        # 增加了与颜色无关的术语，功能只是位置而已
        d.addPairwiseGaussian(sxy=(3, 3), compat=3, kernel=dcrf.DIAG_KERNEL,
                              normalization=dcrf.NORMALIZE_SYMMETRIC)

        # 增加了颜色相关术语，即特征是(x,y,r,g,b)
        d.addPairwiseBilateral(sxy=(80, 80), srgb=(13, 13, 13), rgbim=img, compat=10,
                               kernel=dcrf.DIAG_KERNEL,
                               normalization=dcrf.NORMALIZE_SYMMETRIC)
    else:
        # 使用densecrf类
        d = dcrf.DenseCRF(img.shape[1] * img.shape[0], n_labels)

        # 得到一元势（负对数概率）
        U = unary_from_labels(labels, n_labels, gt_prob=0.7, zero_unsure=None)
        # U = unary_from_labels(labels, n_labels, gt_prob=0.7, zero_unsure=HAS_UNK)## 如果有不确定区域，用这一行代码替换上一行
        d.setUnaryEnergy(U)

        # 这将创建与颜色无关的功能，然后将它们添加到CRF中
        feats = create_pairwise_gaussian(sdims=(3, 3), shape=img.shape[:2])
        d.addPairwiseEnergy(feats, compat=3, kernel=dcrf.DIAG_KERNEL,
                            normalization=dcrf.NORMALIZE_SYMMETRIC)

        # 这将创建与颜色相关的功能，然后将它们添加到CRF中
        feats = create_pairwise_bilateral(sdims=(80, 80), schan=(13, 13, 13),
                                          img=img, chdim=2)
        d.addPairwiseEnergy(feats, compat=10,
                            kernel=dcrf.DIAG_KERNEL,
                            normalization=dcrf.NORMALIZE_SYMMETRIC)

    ####################################
    ###         做推理和计算         ###
    ####################################

    # 进行5次推理
    Q = d.inference(20)

    # 找出每个像素最可能的类
    MAP = np.argmax(Q, axis=0)

    # 将predicted_image转换回相应的颜色并保存图像
    MAP = colorize[MAP, :]
    #imwrite(CRF_image_path, MAP.reshape(img.shape))

    return MAP.reshape(img.shape)


''''
算法调用函数入库
参数image，表示需要进行聚类的图像
参数k,表示keams的K值，默认大小为5
'''
def segment_kmeans_color_(image,k=5):

    height, weight,_ = image.shape

    image_copy = image.copy()

    # 3个通道展平
    img_flat = image_copy.reshape((image.shape[0] * image.shape[1], 3))
    img_flat = np.float32(img_flat)

    # kmean聚类参数
    criteria = (cv2.TERM_CRITERIA_EPS + cv2.TermCriteria_MAX_ITER, 200, 0.5)
    flags = cv2.KMEANS_USE_INITIAL_LABELS

    # 聚类
    compactness, labels, centers = cv2.kmeans(img_flat, k, None, criteria, 20, flags)

    #把聚类标签还原为原图的shape
    labels = labels.reshape(height, weight)

    #获取聚类族群的个数
    cluster = centers.shape[0]

    #统计聚类后每一类占整个图像像素的比例值
    ratio = []
    for i in range(cluster):
        ratio.append(len(labels[labels==i]) / (height * weight))

    #获取从小到大排列的比例值在ratio中的索引值，即标签值
    index = np.argsort(ratio)

    #由于kmean聚类后分配标签并不是按照类别在原图所占比例大小进行分配的，所以这里进行标签更正
    #以聚类后所占原图比例进行从小到大分配标签
    correct_label = np.zeros(labels.shape)
    for i in range(cluster):
        correct_label[labels==index[i]]=i

    #构造字典
    rt = []
    for k in range(cluster):
        rt.append('ratio' + str(k))
    ratio.sort()
    dic = dict(zip(rt,ratio))

    #对标签进行着色
    color_labels = decode_segmap(correct_label,centers, centers.shape[0])
    #使用条件随机场对分割进行后处理
    color_labels = CRFs(image,color_labels,None)
    #把分割彩色图和原图进行融合
    image_res = cv2.addWeighted(image, 1, color_labels, 0.35, 0)

    cv2.imshow("image", image_res)

    dic['image'] = image_res

    return dic


if __name__ == '__main__':
    for i in range(5):

        image = cv2.imread('G:\\Data\\segment\\20200529165615.png', cv2.IMREAD_COLOR)
        k =5
        res = segment_kmeans_color_(image,k)

        cv2.waitKey(0)