python+opencv实现高斯混合（GMM）图像分割

前言

最近想学习一点图像分割方面的知识。网上看了一圈，发现GMM实现图像分割的算法在C++/opencv上实现的案例很多，暂时没有发现python实现的。结合C++版本的语法和对算法的理解，花了半个下午时间，使用python实现了一下，在此记录一下。我们的目的是给下图这位漂亮小姐姐的证件照换一个绿色的背景。

原理

原理部分这里就不细讲了。小破站up主shuhuai008在这方面有非常详细的讲解，有兴趣的同学可以去看看。主要思想就是对当前图片各像素RGB分布按照最大似然拟合出一个最可能的混合高斯分布，然后在反推这些像素分别来自于哪个高斯分布，进行分类。

https://www.bilibili.com/video/BV13b411w7Xj/?spm_id_from=333.337.search-card.all.click&vd_source=d8ac0ff430c94b044512c4edabc63f35

第一步代码实现

在得到像素分类标签best后，找到像素数量最大（即面积最大）的类别，视为背景

import cv2
import numpy as np
import matplotlib.pyplot as plt
def GMM(img):
    print(img.shape)
    # 将一个像素点的rgb值作为一个单元处理
    data = img.reshape((-1, 3))
    print(data.shape)
    # 转换数据类型
    data = np.float32(data)
    # 生成模型
    em = cv2.ml.EM_create()
    # 设置参数，将像素分成num个类别
    num = 3
    em.setClustersNumber(num)
    em.setCovarianceMatrixType(cv2.ml.EM_COV_MAT_GENERIC)  # 默认
    # 训练返回的第三个元素包含了预测分类标签
    best = em.trainEM(data)[2]
    # 筛选出面积最大的一个分类（即认为背景是面积最大的）
    index = 0
    length = len(data[best.ravel() == 0])
    for i in range(0, num):
        if len(data[best.ravel() == i]) > length:
            length = len(data[best.ravel() == i])
            index = i
    # 设置为绿色
    data[best.ravel() == index] = (0, 255, 0)
    # 将结果转换为图片需要的格式
    data = np.uint8(data)
    oi = data.reshape(img.shape)
    # show('img',img)
    # show('res',oi)
    cv2.imshow('img', img)
    cv2.imshow('res', oi)
    cv2.waitKey()
    
if __name__ == '__main__':
    img = cv2.imread("../img/certificate_photo2.jpg")
    GMM(img)

已经有点意思了！但是周边还是毛糙了一点。

第二步代码实现，周边去毛糙

仔细想一下，其实对于这个问题完全可以直接像素二分类，也就是背景和人物两类像素，这样分类更简单，边界也应该更清晰。于是，修改上述代码参数num为2

num = 2

再试试看，结果会是啥样子呢？
我去，小姐姐直接绿了。不过聪明的各位应该也知道为什么了吧？
原因就是，刚才的算法认为面积最大的是背景。这里小姐姐所占据的像素面积大于背景了，所以算法误认为人像是背景了，那就在改下吧！

第三步代码实现，重新筛选背景

实现起来比较笨，不过也很简单，就是吧面积小的作为背景呗~

    for i in range(0, num):
        if len(data[best.ravel() == i]) < length:
            length = len(data[best.ravel() == i])
            index = i

这次应该没毛病了
嗯！效果还可以。

总结

对于背景和前景的选择，我目前只是实现了非常简单的判断方法，更多精妙的idea欢迎各位大佬在评论区交流。
另外个人觉得GMM分割可能并不是非常擅长与背景分割的工作（至少存在很多bug）。比如背景色不够重，和肤色接近，那么分割必然会出现很多问题。实际上我找了网上其他一些证件照尝试，确实效果不尽如人意。这可能也是网上大量教程案例都选择这张图片的原因吧！将来发现更好的方法我会持续更新的！