Robocom预选赛使用opencv实现颜色主色调识别（源码python+解析）
1.前言：
要求：编程实现对图片中具体内容的颜色识别。编程语言：C++或python。判断图片的主色调是红色还是蓝色。 基础数据：组委会提供200张以上的图片作为识别训练集，比赛的200 张图片另外提供。 比赛要求：编写代码，从200张图片中随机读取20张图片，在终端中以 文字形式输出以上20张图片的主色调是红色还是蓝色。
2.源代码
import numpy as np
import cv2 as cv
import math
import random
from PIL import Image, ImageDraw, ImageFont
colors = {'黑色': [0, 180, 0, 255, 0, 46],
'灰色': [0, 180, 0, 43, 46, 220],
'白色': [0, 180, 0, 30, 221, 255],
'红色': [0, 10, 35, 255, 46, 255],
'橙色': [11, 25, 43, 255, 46, 255],
'黄色': [26, 34, 43, 255, 46, 255],
'绿色': [35, 77, 43, 255, 46, 255],
'青色': [78, 99, 43, 255, 46, 255],
'蓝色': [100, 124, 43, 255, 46, 255],
'紫色': [125, 155, 43, 255, 46, 255]
}
color = ['黑色', '灰色', '白色', '红色', '橙色', '黄色', '绿色', '青色', '蓝色', '紫色']
def cv2ImgAddText(img, text, left, top, textColor=(0, 0, 0), textSize=0):
if (isinstance(img, np.ndarray)): # 判断是否OpenCV图片类型
img = Image.fromarray(cv.cvtColor(img, cv.COLOR_BGR2RGB))
draw = ImageDraw.Draw(img)
return cv.cvtColor(np.asarray(img), cv.COLOR_RGB2BGR)
# RGB颜色空间转HSV颜色空间
def simple_rgb2hsv(r, g, b):
r, g, b = r / 255.0, g / 255.0, b / 255.0
mx = max(r, g, b)
mn = min(r, g, b)
df = mx - mn
if mx == mn:
h = 0
elif mx == r:
h = (60/360) * (((g - b) / df) % 6)
elif mx == g:
h = (60/360) * ((b - r) / df + 2)
elif mx == b:
h = (60/360) * ((r - g) / df+4)
if mx == 0:
s = 0
else:
s = df / mx
v = mx
h, s, v = math.ceil(h*180), math.ceil(s*255), math.ceil(v*255)
print(f'h:{h}, s:{s}, v:{v}')
i = 0
for value in colors.values():
if ((h >= 156) & (h <= 180)) & ((s >= 43) & (s <= 255)) & ((v >= 46) & (v <= 255)):
return "红色"
elif ((h >= value[0]) & (h <= value[1])) & ((s >= value[2]) & (s <= value[3])) & ((v >= value[4]) & (v <= value[5])):
return color[i]
i += 1
# 构建图像数据
# K-means 只能处理向量形状的数据，不能处理矩阵型数据，
# 这里 reshape(-1, 3) 代表图片的所有像素点，而每个像素点有三个特征（即三个通道）
def k_means(image):
data = image.reshape((-1, 3))
data = np.float32(data)
# K-means 算法停止条件
# 一个元组，传入 cv.kmeans()，( type, max_iter, epsilon ) type 见下面链接，max_iter 是最大迭代次数，epsilon 要达到的精度
# https://docs.opencv.org/master/d1/d5c/tutorial_py_kmeans_opencv.html
criteria = (cv.TERM_CRITERIA_EPS + cv.TERM_CRITERIA_MAX_ITER, 10, 1.0)
num_clusters = 2 # 聚类的数量
ret, label, center = cv.kmeans(data, num_clusters, None, criteria, 10, cv.KMEANS_RANDOM_CENTERS)
# 求取所聚类的最大概率的值
clusters = np.zeros([num_clusters], dtype=np.int32)
for i in range(len(label)):
clusters[label[i][0]] += 1 # 计算每个类别共有多少个
clusters = np.float32(clusters) / float(h * w) # 计算概率
print(clusters)
max_data = max(clusters)
print(max_data)
max_num = clusters.tolist().index(max_data)
print(max_num)
# center = np.int32(center) # 因为像素值是 0-255 故对其聚类中心进行强制类型转换
center = np.uint8(center) # 因为像素值是 0-255 故对其聚类中心进行强制类型转换
print('this is center', center)
# sorted(clusters) # 这里对主色按比例从大到小排序 [::-1] 代表首尾反转 如[1,2,3] -> [3, 2, 1]
b = center[max_num][0] # 这一类对应的中心，即 RGB 三个通道的值
g = center[max_num][1]
r = center[max_num][2]
print(f'r:{r}, g:{g}, b:{b}')
main_color = simple_rgb2hsv(r, g, b)
print(main_color + '\n')
return main_color
if __name__ == '__main__':
for i in range(20):
num = random.randint(1, 200)
if num < 10:
num_data = '00' + str(num)
elif num < 100 and num >= 10:
num_data = '0' + str(num)
else:
num_data = str(num)
print(f'第{num_data}张图')
# num_data = '190'
image = cv.imdecode(np.fromfile(f'mydata/平安城市测试集/颜色识别任务1/{num_data}.jpg', dtype=np.uint8), -1)
# image = cv.resize(image,None,fx=1,fy=1,interpolation=cv.INTER_LINEAR)
image = cv.resize(image,(600,600))
h, w, ch = image.shape # 读取图像的高、宽、通道数
print(f'图像的高、宽、通道数：{image.shape}')
image = cv2ImgAddText(image, k_means(image), 0, 0, (0, 255, 0), 50)
cv.imshow(f'{i}', image)
cv.waitKey(0)
cv.destroyAllWindows()
3.总结
以上就是我所运用的代码，感谢大家观看，希望有帮助到大家。