阿里巴巴-鹿班实验室算法实现系列——颜色识别

颜色识别部分主要是采用sklearn模块中的kmeans算法实现图像颜色空间的聚类，kmeans算法已有很多博客介绍其算法原理，在此不再详述，kmeans原理可参考此链接

具体代码实现：

from sklearn.cluster import KMeans
import utils
import cv2
import os


def convert_color(value):
    digit = list(map(str, range(10))) + list("ABCDEF")
    if isinstance(value, tuple):                       # 颜色转换，(255,255,255)<->#FFFFF
        string = '#'
        for i in value:
            a1 = i // 16
            a2 = i % 16
            string += digit[a1] + digit[a2]
        return string
    elif isinstance(value, str):                      # 颜色转换，#FFFFF<->(255,255,255)
        a1 = digit.index(value[1]) * 16 + digit.index(value[2])
        a2 = digit.index(value[3]) * 16 + digit.index(value[4])
        a3 = digit.index(value[5]) * 16 + digit.index(value[6])
        return a1, a2, a3


file_path = '.\\imgs\\'
file_list = os.listdir(file_path)
for file in file_list:
    img_name = file_path + file
    ori_image = cv2.imread(img_name)
    image = cv2.cvtColor(ori_image, cv2.COLOR_BGR2RGB)

    image = image.reshape((ori_image.shape[0] * ori_image.shape[1], 3))

    clt = KMeans(n_clusters=8)
    clt.fit(image)

    hist = utils.centroid_histogram(clt)
    bar = utils.plot_colors(hist, clt.cluster_centers_)

    for color in clt.cluster_centers_:
        color = tuple(color.astype('uint8'))
        str_color = convert_color(color)          # rgb值 -> 16进制颜色值
        print(str_color)
    print()
    bar = cv2.cvtColor(bar, cv2.COLOR_RGB2BGR)
    cv2.imshow("img", ori_image)
    cv2.imshow("bar", bar)
    cv2.waitKey(0)