SVM实现小尺寸图片分类

问题背景

在工业识别的场景中，经常会遇到误判与误识别，所以最后输出的结果，需要再通过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)