问题背景
在工业识别的场景中,经常会遇到误判与误识别,所以最后输出的结果,需要再通过N分类算法去过滤一遍,确保识别到的物体,是我们想要的,能实现这个方案有很多,传统机器学习与深度神经网络都可以做到,传统机器学习,比如SVM,决策树,深度神经网络,从最简单的卷积到resnet,小样本学习,迁移学习,都可以实现。本文通过SVM对算法进行封装,实现一个图片的分类。
参考资料
https://blog.csdn.net/linmuquan1989/article/details/126315358 SVM 图片分类python实现
解决方案
SVM的使用非常的简单,在这里都做了封装,只要将类别按照索引,从0开始排列,制作数据集,就能使用下文中的代码,进行小尺寸图片的训练与分类。数据集如下图所示:
使用方法,参照main函数,参照进行传参,即可拥有自己的SVM。
# -*- coding: utf-8 -*-
import os
import cv2
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.svm import SVC
import time
def init_svm_instance(root_path,class_num,kernel="rbf"):
# ----------------------------------------------------------------------------------
# 第一步 切分训练集和测试集
# ----------------------------------------------------------------------------------
X = [] # 定义图像名称,磁盘地址
Y = [] # 定义图像分类类标
for i in range(0, int(class_num)):
# 遍历文件夹,读取图片
for f in os.listdir(root_path+"/%s" % i):
# 获取图像名称
X.append(root_path+"//" + str(i) + "//" + str(f))
# 获取图像类标即为文件夹名称
Y.append(i)
X = np.array(X)
Y = np.array(Y)
# 随机率为100% 选取其中的20%作为测试集
X_train, X_test, y_train, y_test = train_test_split(X, Y,
test_size=0.2, random_state=1)
print(len(X_train), len(X_test), len(y_train), len(y_test))
# ----------------------------------------------------------------------------------
# 第二步 图像读取及转换为像素直方图
# ----------------------------------------------------------------------------------
# 训练集
XX_train = []
for i in X_train:
# 读取图像
# print i
image = cv2.imdecode(np.fromfile(i, dtype=np.uint8), cv2.IMREAD_COLOR)
# 图像像素大小一致
img = cv2.resize(image, (64, 64),
interpolation=cv2.INTER_CUBIC)
# 计算图像直方图并存储至X数组
hist = cv2.calcHist([img], [0, 1], None,
[64, 64], [0.0, 255.0, 0.0, 255.0])
XX_train.append(((hist / 255).flatten()))
# ----------------------------------------------------------------------------------
# 第三步 基于支持向量机的图像分类处理
# ----------------------------------------------------------------------------------
# 常见核函数‘linear’, ‘poly’, ‘rbf’, ‘sigmoid’, ‘precomputed’
clf = SVC().fit(XX_train, y_train)
clf = SVC(kernel=kernel).fit(XX_train, y_train)
return clf
def svm_infer(clf,img_disk_path,label_list):
XX_test = []
image = cv2.imdecode(np.fromfile(img_disk_path, dtype=np.uint8), cv2.IMREAD_COLOR)
# 图像像素大小一致
img = cv2.resize(image, (64, 64),
interpolation=cv2.INTER_CUBIC)
# 计算图像直方图并存储至X数组
hist = cv2.calcHist([img], [0, 1], None,
[64, 64], [0.0, 255.0, 0.0, 255.0])
XX_test.append(((hist / 255).flatten()))
start = time.time()
predictions_labels = clf.predict(XX_test)
end = time.time()
print("推理耗时:", (end - start))
return label_list[predictions_labels[0]]
if __name__ == '__main__':
start = time.time()
clf = init_svm_instance('D:/codeRepo/SVMImageClassification/bbld','2')
end = time.time()
print("初始化耗时:", (end - start))
result = svm_infer(clf,'D:/codeRepo/SVMImageClassification/bbld/0/ld_222283.jpg',['ld','ldk'])
print(result)