高斯混合模型是一种概率式的聚类方法,它假定所有的数据样本x由k个混合多元高斯分布组合成的混合分布生成。
其中高斯分布的概率密度函数如下:
现在的问题就是如何求 α , μ , σ \alpha,\mu,\sigma α,μ,σ
import numpy as np
import matplotlib.pyplot as plt
from sklearn.mixture import GaussianMixture
import cv2
import matplotlib as mpl
import glob
import os
def draw_polygon(img, pts):
for pt in pts.reshape(-1, 2).astype(np.int32):
# import ipdb; ipdb.set_trace()
cv2.circle(img, tuple(pt),2, (0, 0, 255))
def find_target_contour(mask, ncls):
biggest_contours = []
all_areas = []
area_threshold = mask.shape[0] * mask.shape[1] /5
for label in range(ncls):
mask_a = (mask == label).astype(np.uint8)
contours, hierarchy = cv2.findContours(mask_a, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
areas = [cv2.contourArea(item) for item in contours]
idx = np.argmax(areas)
biggest_contours.append(contours[idx].reshape(-1, 2))
all_areas.append(areas[idx])
# only keep 2 region which take big area
biggest_contours = np.asarray(biggest_contours)
all_areas = np.asarray(all_areas)
# import ipdb; ipdb.set_trace()
biggest_contours = biggest_contours[all_areas > area_threshold]
# if len(biggest_contours)
# find whilch contour is closer to bottom of img
centers = np.array([np.mean(item, axis = 0) for item in biggest_contours] )
idx = np.argmax(centers[:, 1])
selected_contours = biggest_contours[idx]
rect = np.min(selected_contours, axis = 0), np.max(selected_contours, axis = 0)
return np.array(rect), selected_contours
def run():
# dataset_root = 'data/'
dataset_root = 'tmp'
imgs = glob.glob(os.path.join(dataset_root, '*.png'))
ncls = 3
for iimg, cur_path in enumerate(imgs):
print(f'[{iimg}/{len(imgs)}].reading {cur_path}')
img = cv2.imread(cur_path)
npts = img.shape[0] * img.shape[1]
classifier = GaussianMixture(ncls)
input_img = cv2.blur(img, (5, 5))
classifier.fit(input_img.reshape(npts, 3))
colors = [0, 80, 125, 180, 255]
labels = classifier.predict(img.reshape(npts, 3))
mask = labels.reshape(*img.shape[:2]).astype(np.uint8)
kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (5,5))
eroded = cv2.erode(mask, kernel)
mask = cv2.dilate(eroded, kernel)
rect, mask_contour = find_target_contour(mask, ncls)
labels = np.expand_dims(labels, -1)
for px in labels:
px[:] = colors[px[0]]
binary_img = labels.reshape(*img.shape[:2])
plt.figure(figsize = (11, 4))
plt.subplot(121)
draw_polygon(img, mask_contour)
rect = rect.astype(np.int32)
cv2.rectangle(img, tuple(rect[0]), tuple(rect[1]), (0, 255, 255), 4)
rgb_img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
plt.imshow(rgb_img)
plt.axis('off')
plt.subplot(122)
plt.imshow(binary_img, cmap = mpl.colormaps['viridis'], interpolation= 'nearest')
plt.axis('off')
plt.show()
run()
从下图可以看出,肤色相近的区域都被识别成了一个区域,对于一些简单的任务GMM还是能很好处理的。