HOG+SVM进行图片分类

1、介绍

通过提取图像的HOG特征，对HOG特征训练一个线性分类器，对图片进行分类。

HOG特征对图像的大小有严格的要求，不同大小的图像求出的HOG特征有所不同，所以在对图片求取特征时要统一图片的大小。我这里使用图片的大小为128*128*1大小的灰度图像，所以在进行测试时要使用相同大小的图片进行测试，另外获取HOG特征的参数要保持相同。

首先检查图像的大小是否为128*128*1的形状，如果图像的大小不相同则把图像裁剪为128*128*1的形状。然后计算所有图像的HOG特征，训练一个线性SVM模型，并对数据进行测试，计算正确率。

2、代码

训练过程的代码如下：

import os
import time
import cv2
import joblib
import numpy as np
from PIL import Image
from skimage.feature import hog
from sklearn.svm import LinearSVC

# 第一个是你的类别   第二个是类别对应的名称   输出结果的时候方便查看
label_map = {
    1: 'cat',
    2: 'chick',
    3: 'snack',
    4: 'unknown'
}
train_image_path = 'image'
test_image_path = 'test_img'
# 训练集标签的位置
train_label_path = os.path.join(f'{train_image_path}', 'train.txt')
# 测试集标签的位置
test_label_path = os.path.join(f'{test_image_path}', 'train.txt')
image_height = 128
image_width = 128
train_feat_path = 'train/'
test_feat_path = 'test/'
model_path = 'model/'


def get_name_label(file_path):
    print("read label from ", file_path)
    name_list = []
    label_list = []
    with open(file_path) as f:
        for line in f.readlines():
            # 一般是name label  三部分，所以至少长度为3  所以可以通过这个忽略空白行
            if len(line) >= 3:
                name_list.append(line.split(' ')[0])
                label_list.append(line.split(' ')[1].replace('\n', '').replace('\r', ''))
                if not str(label_list[-1]).isdigit():
                    print("label必须为数字，得到的是：", label_list[-1], "程序终止，请检查文件")
                    exit(1)
    return name_list, label_list


def get_image_list(filePath, nameList):
    print('read image from ', filePath)
    img_list = []
    for name in nameList:
        temp = cv2.imread(os.path.join(filePath, name), 0)
        img_list.append(temp.copy())
    return img_list


def get_hog_feat(image_list):
    i = 0
    hog_feature_list = []
    for image in image_list:
        if (image.shape != (image_width, image_height, 1)
                and image.shape != (image_width, image_height)):
            print(f"num of {i} img size not is"
                  f" {(image_width, image_width, 1)}.{image.shape=}")
            exit(-1)
        fd = hog(image, orientations=12, block_norm='L1', pixels_per_cell=[8, 8],
                 cells_per_block=[4, 4], visualize=False, transform_sqrt=True)
        hog_feature_list.append(fd)
    return hog_feature_list


def train(train_hog_fea_list, train_label_list):
    start_time = time.time()
    features = train_hog_fea_list
    labels = train_label_list
    print("Training a Linear LinearSVM Classifier.")
    clf = LinearSVC()
    clf.fit(features, labels)
    end_time = time.time()
    cost_time = end_time - start_time
    print(f"train fit cost {cost_time:.2f}s.")

    # 下面的代码是保存模型的
    if not os.path.exists(model_path):
        os.makedirs(model_path)
    joblib.dump(clf, model_path + 'model')
    print("训练之后的模型存放在model文件夹中")

    correct_num = 0
    for label, feature in zip(labels, features):
        feature = feature.reshape((1, -1))
        tmp_label = clf.predict(feature)
        if tmp_label == label:
            correct_num += 1
    print(f"correct rate is {correct_num/len(labels)*100=:.2f}%")
    return clf


def predict_one(test_image, model):
    test_hog_feature = get_hog_feat([test_image])[0]
    test_hog_feature = test_hog_feature.reshape((1, -1))
    result = model.predict(test_hog_feature)
    result = label_map[int(result[0])]
    return result


def predict_all(test_hog_feature, model):
    results = []
    for hog_feature in test_hog_feature:
        hog_feature = hog_feature.reshape((1, -1))
        result = model.predict(hog_feature)
        result = label_map[int(result)]
        results.append(result)
    return results


if __name__ == '__main__':
    train_name, train_label = get_name_label(train_label_path)
    print(train_name)
    print(train_label)
    test_name, test_label = get_name_label(test_label_path)
    print(test_name)
    print(test_label)
    train_image_list = get_image_list(train_image_path, train_name)
    test_image_list = get_image_list(test_image_path, test_name)

    train_hog_feature_list = get_hog_feat(train_image_list)
    test_hog_feature_list = get_hog_feat(test_image_list)

    model = train(train_hog_feature_list, train_label)
    result = predict_one(test_image_list[0], model)
    results = predict_all(test_hog_feature_list, model)
    print(result)
    print(results)

预测过程的代码：

import os

import cv2
import joblib
from train import predict_one, predict_all

user_image_path = "user_img/save_img/1.jpg"
img = cv2.imread(user_image_path, 0)
print(img.shape)
model = joblib.load("model/model")
print(model)
# result = predict_one(img, model)
# print(result)

for root, dirs, files in os.walk(user_image_path[:-5]):
    print(root, dirs, files)
    for file in files:
        # print(root.join())
        image = cv2.imread(root + file, 0)
        result = predict_one(image, model)
        print(file, result)

3、资源下载