OpenCV K-means 实现主色提取

机器视觉CV

发布于 2019-07-15 19:20:43

2.9K0

发布于 2019-07-15 19:20:43

文章被收录于专栏：机器视觉CV机器视觉CV

建议阅读时长 10 分钟

说在前面

最近大量使用到 K-means 算法，一直使用的是 Sklearn 库来实现，没想到 OpenCV 库里面也有这个算法，故找了个例子实现一下。

本节内容：

OpenCV 中的 cv.kmeans 函数及参数介绍
实现一个图片的主色提取，并按照比例大小生成彩色卡片

OpenCV kmeans 算法

1retval, bestLabels, centers = cv.kmeans(data, K, bestLabels, criteria, attempts, flags[, centers])

主要参数说明如下： input

data: 聚类的数据
K: 聚类数量
bestLabels：输出的标签，即每个数据聚类的标签是什么
criteria：聚类算法迭代终止的条件，( type, max_iter, epsilon) https://docs.opencv.org/master/d1/d5c/tutorial_py_kmeans_opencv.html
- type：终止条件类型
- max_iter:迭代次数
- epsilon：终止的精度
attempts：https://docs.opencv.org/master/d0/de1/group__core.html#ga276000efe55ee2756e0c471c7b270949
flags：聚类初始中心设定的标志
[, centers]

output

retval
bestLabels:聚类的结果
centers：聚类的中心

代码如下：

 1import numpy as np
 2import cv2 as cv
 3
 4image = cv.imread(r'f:/jupyter/OpenCV4.0/sample_img/1.jpg')
 5cv.imshow("input", image)
 6h, w, ch = image.shape
 7print((image.shape))
 8
 9# 构建图像数据
10# K-means 只能处理向量形状的数据，不能处理矩阵型数据，
11# 这里 reshape(-1, 3) 代表图片的所有像素点，而每个像素点有三个特征（即三个通道）
12data = image.reshape((-1, 3))
13data = np.float32(data)
14
15# K-means 算法停止条件
16# 一个元组，传入 cv.kmeans()，( type, max_iter, epsilon ) type 见下面链接，max_iter 是最大迭代次数，epsilon 要达到的精度
17# https://docs.opencv.org/master/d1/d5c/tutorial_py_kmeans_opencv.html
18criteria = (cv.TERM_CRITERIA_EPS + cv.TERM_CRITERIA_MAX_ITER, 10, 1.0)
19num_clusters = 4  # 聚类的数量
20ret, label, center = cv.kmeans(data, num_clusters, None, criteria,
21                               num_clusters, cv.KMEANS_RANDOM_CENTERS)
22
23# 生成主色彩条形卡片，大小是：高：50 宽：原图的宽
24card = np.zeros((50, w, 3), dtype=np.uint8)
25clusters = np.zeros([num_clusters], dtype=np.int32)
26for i in range(len(label)):
27    clusters[label[i][0]] += 1  # 计算每个类别共有多少个
28clusters = np.float32(clusters) / float(h * w)  # 计算概率
29center = np.int32(center) # 因为像素值是 0-255 故对其聚类中心进行强制类型转换
30
31x_offset = 0
32for c in np.argsort(clusters)[::-1]: # 这里对主色按比例从大到小排序 [::-1] 代表首尾反转 如[1,2,3] -> [3, 2, 1]
33    dx = np.int(clusters[c] * w) # 这一类转换成色彩卡片有多宽
34    b = center[c][0] # 这一类对应的中心，即 RGB 三个通道的值
35    g = center[c][1]
36    r = center[c][2]
37    cv.rectangle(card, (x_offset, 0), (x_offset + dx, 50),
38                 (int(b), int(g), int(r)), -1) # 每个主色画出一个矩形
39    x_offset += dx # 偏置就是每个主色的宽度
40
41cv.imshow("color table", card)
42cv.waitKey(0)
43cv.destroyAllWindows()