利用SVM实现复杂环境下颜色识别任务

在处理一些复杂情况下的颜色识别任务时，通过RGB或HSV空间下的阈值法有时很难实现颜色识别任务。我们考虑使用机器学习对颜色进行分类（在这里我们采用支持向量机SVM）。我们事先准备好不同颜色的照片数据集和测试集，利用OpenCV库中的SVM方法对其进行分类，识别颜色的特征向量用的是颜色直方图，当然大家也可以考虑使用颜色矩进行分类（这两种分类方法的函数我都将在下面的代码中给出）。废话不多说，直接上代码。

import cv2
import numpy as np
import os

#统计颜色直方图
def colorhist(src):
    src=cv2.resize(src,(50,50))
    bgr_planes = cv2.split(src)
    histSize = 256
    # 256会被排除
    histRange = (0, 256)
    accumulate = False
    b_hist = cv2.calcHist(bgr_planes, [0], None, [
        histSize], histRange, accumulate=accumulate)
    g_hist = cv2.calcHist(bgr_planes, [1], None, [
        histSize], histRange, accumulate=accumulate)
    r_hist = cv2.calcHist(bgr_planes, [2], None, [
        histSize], histRange, accumulate=accumulate)
    bgr_hist = np.array([b_hist,g_hist,r_hist])
    return bgr_hist.flatten()

# 计算图像的颜色矩
def compute_color_moments(image):
    # 将图像转换为HSV色彩空间
    hsv_image = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)

    # 分割H、S、V通道
    h, s, v = cv2.split(hsv_image)

    # 计算颜色矩
    h_mean = np.mean(h)
    h_std = np.std(h)
    s_mean = np.mean(s)
    s_std = np.std(s)
    v_mean = np.mean(v)
    v_std = np.std(v)

    # 返回颜色矩特征向量
    color_moments = [h_mean, h_std, s_mean, s_std, v_mean, v_std]
    return color_moments


def prepare_dataset():
    X=blueallhist+greenallhist##特征向量
    noblue= len(blueallhist)
    nogreen=len(greenallhist)
    y = []  # 标签列表
    for i in range(noblue):
        y.append([0])     ##0代表蓝色
    for i in range(nogreen):
        y.append([1])    ##1代表绿色

    return np.array(X, dtype=np.float32), np.array(y)


def SVM():
    # 准备数据
    X, y = prepare_dataset()

    # 创建SVM分类器
    svm_classifier = cv2.ml.SVM_create()

    # 设置SVM参数
    svm_classifier.setType(cv2.ml.SVM_C_SVC)
    svm_classifier.setKernel(cv2.ml.SVM_LINEAR)

    # 训练分类器
    svm_classifier.train(X, cv2.ml.ROW_SAMPLE, y)

    # 在测试集上进行预测
    _, y_pred = svm_classifier.predict(testhist)
    print(y_pred)

    # # 计算准确率
    # accuracy = np.mean(y_pred == y)
    # print("准确率：", accuracy)


def getallhist(filepath):
    allhist=[]
    for filename in os.listdir(filepath):
        img=cv2.imread(filepath+"/"+filename)
        allhist.append(colorhist(img))
    return allhist

if __name__ == "__main__":
    blueallhist = getallhist(r"C:\Users\acheng\Desktop\Demo\blue")
    greenallhist = getallhist(r"C:\Users\acheng\Desktop\Demo\green")

    ##准备测试集
    testimg=cv2.imread("green3.png")
    testhist = [colorhist(testimg)]
    testhist = np.array(testhist, dtype=np.float32)

    SVM()

关于代码部分的讲解和使用过程中的一些注意事项我会抽时间补上，大家也可以私信我询问，共同交流。