Robocom qualifiers use opencv to realize color main color recognition (source code python+analysis)

1 Introduction:

Requirements : Program to realize the color recognition of the specific content in the picture. Programming language: C++ or python. Determines whether the dominant color of the image is red or blue. Basic data: The organizing committee provides more than 200 pictures as the recognition training set, and 200 pictures for the competition are provided separately. Competition requirements: Write code to randomly read 20 pictures from 200 pictures, and output whether the main color of the above 20 pictures is red or blue in text form in the terminal.

 

2. Source code

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. Summary

The above is the code I used, thank you for watching, I hope it will help you.

Guess you like

Origin blog.csdn.net/qq_58130152/article/details/127475487