opencv-112-利用KMeans图像分割进行背景替换

发表于 | 分类于

知识点

KMeans可以实现简单的证件照片的背景分割提取与替换，大致可以分为如下几步实现

  1. 读入图像建立KMenas样本
  2. 使用KMeans图像分割，指定指定分类数目
  3. 取左上角的label得到背景cluster index
  4. 生成mask区域，然后高斯模糊进行背景替换

代码（c++,python）

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
#include <opencv2/opencv.hpp>
#include <iostream>

using namespace cv;
using namespace std;

int main(int argc, char** argv) {
	Mat src = imread("D:/projects/opencv_tutorial/data/images/toux.jpg");
	if (src.empty()) {
		printf("could not load image...\n");
		return -1;
	}
	namedWindow("input image", WINDOW_AUTOSIZE);
	imshow("input image", src);

	int width = src.cols;
	int height = src.rows;
	int dims = src.channels();

	// 初始化定义
	int sampleCount = width*height;
	int clusterCount = 3;
	Mat labels;
	Mat centers;

	// RGB 数据转换到样本数据
	Mat sample_data = src.reshape(3, sampleCount);
	Mat data;
	sample_data.convertTo(data, CV_32F);

	// 运行K-Means
	TermCriteria criteria = TermCriteria(TermCriteria::EPS + TermCriteria::COUNT, 10, 0.1);
	kmeans(data, clusterCount, labels, criteria, clusterCount, KMEANS_PP_CENTERS, centers);

	Mat mask = Mat::zeros(src.size(), CV_8UC1);
	int index = labels.at<int>(0, 0);
	labels = labels.reshape(1, height);
	for (int row = 0; row < height; row++) {
		for (int col = 0; col < width; col++) {
			int c = labels.at<int>(row, col);
			if (c == index) {
				mask.at<uchar>(row, col) = 255;
			}
		}
	}

	Mat se = getStructuringElement(MORPH_RECT, Size(3, 3), Point(-1, -1));
	dilate(mask, mask, se);
	GaussianBlur(mask, mask, Size(5, 5), 0);
	Mat result = Mat::zeros(src.size(), CV_8UC3);
	for (int row = 0; row < height; row++) {
		for (int col = 0; col < width; col++) {
			float w1 = mask.at<uchar>(row, col) / 255.0;
			Vec3b bgr = src.at<Vec3b>(row, col);
			bgr[0] = w1 * 255.0 + bgr[0] * (1.0 - w1);
			bgr[1] = w1 * 0 + bgr[1] * (1.0 - w1);
			bgr[2] = w1 * 255.0 + bgr[2] * (1.0 - w1);
			result.at<Vec3b>(row, col) = bgr;
		}
	}
	imshow("KMeans-image-Demo", result);
	waitKey(0);
	return 0;
}
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
"""
利用KMeans图像分割进行背景替换
"""

import cv2 as cv
import numpy as np

image = cv.imread('images/toux.jpg')
cv.imshow("input", image)
h, w, ch = image.shape

# 构建图像数据
data = image.reshape((-1, 3))
data = np.float32(data)

# 图像分割
criteria = (cv.TERM_CRITERIA_EPS + cv.TERM_CRITERIA_MAX_ITER, 10, 1.0)
num_clusters = 4
ret, label, center = cv.kmeans(data, num_clusters, None, criteria, num_clusters, cv.KMEANS_RANDOM_CENTERS)

# 生成mask区域
index = label[0][0]
center = np.uint8(center)
color = center[0]
mask = np.zeros((h, w), dtype=np.uint8)
label = np.reshape(label, (h, w))
mask[label == index] = 255

# 高斯模糊
se = cv.getStructuringElement(cv.MORPH_RECT, (3, 3))
cv.dilate(mask, se, mask)
mask = cv.GaussianBlur(mask, (5, 5), 0)
cv.imshow("background-mask", mask)

# 背景替换
result = np.zeros((h, w, ch), dtype=np.uint8)
for row in range(h):
    for col in range(w):
        w1 = mask[row, col] / 255.0
        b, g, r = image[row, col]
        b = w1 * 255 + b * (1.0 - w1)
        g = w1 * 0 + g * (1.0 - w1)
        r = w1 * 255 + r * (1.0 - w1)
        result[row, col] = (b, g, r)
cv.imshow("background-substitution", result)

cv.waitKey(0)
cv.destroyAllWindows()

结果

代码地址

github