在处理一些复杂情况下的颜色识别任务时,通过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()
关于代码部分的讲解和使用过程中的一些注意事项我会抽时间补上,大家也可以私信我询问,共同交流。