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.总结

以上就是我所运用的代码，感谢大家观看，希望有帮助到大家。