SVM训练代码

import cv2
import os
import numpy as np
from numpy.linalg import norm
from args import args
import time

class StatModel(object):
    def load(self, fn):
        self.model = self.model.load(fn)

    def save(self, fn):
        self.model.save(fn)


class SVM(StatModel):
    def __init__(self, C=1, gamma=0.5):
        self.model = cv2.ml.SVM_create()
        self.model.setGamma(gamma)
        self.model.setC(C)
        self.model.setKernel(cv2.ml.SVM_RBF)
        self.model.setType(cv2.ml.SVM_C_SVC)
        # 不能保证包括所有省份

    # 训练svm
    def train(self, samples, responses):
        self.model.train(samples, cv2.ml.ROW_SAMPLE, responses)

    # 调用svm
    def predict(self, samples):
        r = self.model.predict(samples)
        return r[1].ravel()


# 定义参数
SZ = 20  # 训练图片长宽
L = 136
W = 36


# MAX_WIDTH = args.MAX_WIDTH  # 原始图片最大宽度
# Min_Area = args.Min_Area  # 车牌区域允许最大面积
# PROVINCE_START = args.PROVINCE_START
# provinces = args.provinces

# 来自opencv的sample，用于svm训练
def deskew(img):
    m = cv2.moments(img)
    if abs(m['mu02']) < 1e-2:
        return img.copy()
    skew = m['mu11'] / m['mu02']
    M = np.float32([[1, skew, -0.5 * W * skew], [0, 1, 0]])
    img = cv2.warpAffine(img, M, (L, W), flags=cv2.WARP_INVERSE_MAP | cv2.INTER_LINEAR)
    return img


# 来自opencv的sample，用于svm训练
def preprocess_hog(digits):
    samples = []
    for img in digits:
        gx = cv2.Sobel(img, cv2.CV_32F, 1, 0)
        gy = cv2.Sobel(img, cv2.CV_32F, 0, 1)
        mag, ang = cv2.cartToPolar(gx, gy)
        bin_n = 16
        bin = np.int32(bin_n * ang / (2 * np.pi))
        # print('bin',len(bin))
        # print('mag', len(mag))
        bin_cells = bin[:10, :10], bin[10:, :10], bin[:10, 10:], bin[10:, 10:]
        mag_cells = mag[:10, :10], mag[10:, :10], mag[:10, 10:], mag[10:, 10:]

        # bin_cells = mag[:18, :18], mag[18:, :18], mag[:18, 18:], mag[18:, 18:]
        # mag_cells = mag[:18, :18], mag[18:, :18], mag[:18, 18:], mag[18:, 18:]
        # print('bin', len(bin_cells), type(bin_cells))
        # print('mag', len(mag_cells), type(mag_cells))
        hists = [np.bincount(b.ravel(), m.ravel(), bin_n) for b, m in zip(bin_cells, mag_cells)]
        # print('hists', len(hists), type(hists))
        hist = np.hstack(hists)
        # transform to Hellinger kernel
        eps = 1e-7
        hist /= hist.sum() + eps
        hist = np.sqrt(hist)
        hist /= norm(hist) + eps

        samples.append(hist)
    return np.float32(samples)


def train_svm(path):
    # 识别英文字母和数字
    Model = SVM(C=1, gamma=0.5)

    # 英文字母和数字部分训练
    chars_train = []
    chars_label = []

    for root, dirs, files in os.walk(os.path.join(path, 'plate5')):
        if len(os.path.basename(root)) > 3:
            continue
        root_int = os.path.basename(root)
        if root_int == 'has':
            index = 1
        else:
            index = 0
        print('root_int',index)
        for filename in files:
            # print('input：{}'.format(filename))
            filepath = os.path.join(root, filename)
            digit_img = cv2.imread(filepath)
            digit_img = cv2.cvtColor(digit_img, cv2.COLOR_BGR2GRAY)
            digit_img = cv2.resize(digit_img, (136, 36))
            chars_train.append(digit_img)
            chars_label.append(index)

    chars_train = list(map(deskew, chars_train))
    chars_train = preprocess_hog(chars_train)
    chars_label = np.array(chars_label)
    Model.train(chars_train, chars_label)

    if not os.path.exists("plate_svm5.dat"):
        # 保存模型
        Model.save("plate_svm5.dat")
    else:
        # 更新模型
        os.remove("plate_svm5.dat")
        Model.save("plate_svm5.dat")



def show_files(path, all_files):
    # 首先遍历当前目录所有文件及文件夹
    file_list = os.listdir(path)
    # 准备循环判断每个元素是否是文件夹还是文件，是文件的话，把名称传入list，是文件夹的话，递归
    for file in file_list:
        # 利用os.path.join()方法取得路径全名，并存入cur_path变量，否则每次只能遍历一层目录
        cur_path = os.path.join(path, file)
        # 判断是否是文件夹
        if os.path.isdir(cur_path):
            show_files(cur_path, all_files)
        else:
            if cur_path.find(".jpg") > 0:
                all_files.append(cur_path)

    return all_files


def predict_svm(path):
    # 识别英文字母和数字
    t1 = time.time()
    plate_svm = SVM(C=1, gamma=0.5)
    plate_svm.load("plate_svm4.dat")
    # 英文字母和数字部分训练
    correct_count = 0  # 正确数
    contents = show_files(path, [])
    contents = contents[:((int)(len(contents)*0.2))]

    # print("contents",len(contents))
    test_num = (int)(len(contents))
    print("test_num",test_num)

    for img_name in contents:
        digit_img = cv2.imread(img_name)
        digit_img = cv2.cvtColor(digit_img, cv2.COLOR_BGR2GRAY)
        digit_img = cv2.resize(digit_img, (136, 36))
        # print('digit_img',digit_img)
        # cv2.imshow('digit_img',digit_img)
        # cv2.waitKey(0)
        chars_train = preprocess_hog([digit_img])
        # print('chars_train',chars_train)
        resp = plate_svm.predict(chars_train)
        # print('part2',resp[0])
        if resp[0] == 1:
            correct_count += 1

    print('correct_count',correct_count)
    print('test_num',test_num)
    accuracy = correct_count / test_num
    print("accuracy:", accuracy)  # 准确率
    t2 = time.time() - t1
    t3 = t2 / test_num
    print('t2',t2)
    print('t3', t3)


def predict_svm1(img_name):

    plate_svm = SVM(C=1, gamma=0.5)
    plate_svm.load("plate_svm2.dat")
    digit_img = cv2.imread(img_name)
    digit_img = cv2.cvtColor(digit_img, cv2.COLOR_BGR2GRAY)
    digit_img = cv2.resize(digit_img, (136, 36))
    chars_train = preprocess_hog([digit_img])
    resp = plate_svm.predict(chars_train)
    return resp[0]





if __name__ == '__main__':
    # img = cv2.imread('Train/plate/has/plate_AUU093.jpg')
    # img1 = deskew(cv2.cvtColor(img, cv2.COLOR_BGR2GRAY))
    # cv2.imshow('0',img1)
    # cv2.waitKey(0)


    path = 'Train'
    # train_svm(path)
    # train_RandomForestClassifier(path)
    print('完成')

    predict_svm("Train/plate1/has")
    # a = predict_svm1("Train/plate2/has/100.jpg")
    # print(a)
    # test("svm_detect.mat", "Train/plate1/testdata")