Robocom比赛使用opencv进行颜色识别（除白色外出现的其它几种主要颜色） （附带源码+解析）

1.前言

要求：编写代码，从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, 43, 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, 255, 0), textSize=20):

    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)


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(h, s, 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 = 10  # 聚类的数量
    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)  # 计算概率
    center = np.int32(center)  # 因为像素值是 0-255 故对其聚类中心进行强制类型转换

    tp = []
    # 这里对主色按比例从大到小排序 [::-1] 代表首尾反转 如[1,2,3] -> [3, 2, 1]
    for c in range(10):
        # dx = np.int(clusters[c] * w)  # 这一类转换成色彩卡片有多宽
        b = center[c][0]  # 这一类对应的中心，即 RGB 三个通道的值
        g = center[c][1]
        r = center[c][2]
        main_color = simple_rgb2hsv(r, g, b)
        if not main_color in tp:
            tp.append(main_color)

    words = ""
    for i in tp:
        if i != '白色' and i != '灰色':
            words = words+' '+i
    print(words)
    return words


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 = '107'
        image = cv.imdecode(np.fromfile(f'mydata/平安城市测试集/颜色识别任务2/{num_data}.jpg', dtype=np.uint8), -1)
        image = cv.resize(image,None,fx=0.6,fy=0.6,interpolation=cv.INTER_LINEAR)
        h, w, ch = image.shape # 读取图像的高、宽、通道数
        print(f'图像的高、宽、通道数：{image.shape}')
        image = cv2ImgAddText(image, k_means(image), 0, 0, (0, 0, 0), 50)
        cv.imshow(f'{i}', image)
        cv.waitKey(0)
        cv.destroyAllWindows()

 3.总结

关键代码后面都有解析，大家将就着看看。有什么问题可以在评论区进行积极讨论。