vgg16网络实现车牌二分类（整理）

参考：https://blog.youkuaiyun.com/qq_37389133/article/details/80707304，感谢

文件结构：

vgg16:

---data      #存放训练数据

     ---Yes

     ---No

---model   #存放训练好的模型

---pred_data   #存放测速数据

    ---Yes

    ---No

---vgg16.npy   #参数文件，直接下载即可,点击此处下载

---vgg16.py     #程序文件，如下

# 代码如下：

import os
import numpy as np
import tensorflow as tf
import skimage.io
import skimage.transform
import matplotlib as mpl

#在ubuntu-server版本上运行，所以不能显示图片，解决办法就是加上下面这句话，ubuntu桌面版应该不需要。
mpl.use('Agg')
#第二步就是在.config/matplotlib下的matplotlibrc文件中添加一行 ckend : Agg
#.config在home文件下，我的是在账户文件夹下 .config一般是隐藏的，用ls -a即可显示。


from matplotlib import pyplot as plt


# 这段代码是将文件夹中的图片读出来，并且统一好规格，224*224的规格

def load_img(path):
    img = skimage.io.imread(path)
    img = img / 255.0
    short_edge = min(img.shape[:2])
    yy = int((img.shape[0] - short_edge) / 2)
    xx = int((img.shape[1] - short_edge) / 2)
    crop_img = img[yy: yy + short_edge, xx: xx + short_edge]
    resized_img = skimage.transform.resize(crop_img, (224, 224))[None, :, :, :]  # shape [1, 224, 224, 3]
    return resized_img


# 这里指从图片中读取数据，并且我们自己规定身体长度为一种结果特征，以此来进行区别和训练，注意这里只是指一种条件，方便训练，事实上也可以多种或者其他的条件！
def load_data():
    #参考博客是分类猫和老虎，所以两个文件夹分别叫tiger和cat
    imgs = {'Yes': [], 'No': []}
    for k in imgs.keys():
        #data文件中存放的是训练数据，下面有两个子文件夹，分别是你的两个分类tiger,cat，当然也可以是其他分类，比如我做的就是车牌和非车牌的二分类
        dir = 'data/' + k
        for file in os.listdir(dir):
            if not file.lower().endswith('.jpg'):
                continue
            try:
                resized_img = load_img(os.path.join(dir, file))
            except OSError:
                continue
            imgs[k].append(resized_img)
            if len(imgs[k]) == 400:
                break
        print('***', k, len(imgs[k]))
    #如果你改了tiger和cat的名字，记着把以下的也改了呀
    tigers_y = np.maximum(36, np.random.randn(len(imgs['Yes']), 1) * 32 + 180)
    cat_y = np.maximum(10, np.random.randn(len(imgs['No']), 1) * 8 + 40)
    return imgs['Yes'], imgs['No'], tigers_y, cat_y


# 这里就是训练的和预测的主体部分了
class Vgg16:
    vgg_mean = [103.939, 116.779, 123.68]

    def __init__(self, vgg16_npy_path=None, restore_from=None, **kwargs):
        try:
            # 这里我们将vgg16.npy里的参数导入
            self.data_dict = np.load(vgg16_npy_path, encoding='latin1').item()

        except FileNotFoundError:
            print('Something error about the vgg_npy')

        self.tfx = tf.placeholder(tf.float32, [None, 224, 224, 3])
        self.tfy = tf.placeholder(tf.float32, [None, 1])

        red, green, blue = tf.split(axis=3, num_or_size_splits=3, value=self.tfx * 255.0)
        # 这里需要将图片的rgb格式转化为bgr格式。

        bgr = tf.concat(axis=3, values=[
            blue - self.vgg_mean[0],
            green - self.vgg_mean[1],
            red - self.vgg_mean[2],
        ])
        # 这里我们都是再用vgg16.npy中的参数来进行前面部分的训练和预测
        conv1_1 = self.conv_layer(bgr, "conv1_1")
        conv1_2 = self.conv_layer(conv1_1, "conv1_2")
        pool1 = self.max_pool(conv1_2, 'pool1')

        conv2_1 = self.conv_layer(pool1, "conv2_1")
        conv2_2 = self.conv_layer(conv2_1, "conv2_2")
        pool2 = self.max_pool(conv2_2, 'pool2')

        conv3_1 = self.conv_layer(pool2, "conv3_1")
        conv3_2 = self.conv_layer(conv3_1, "conv3_2")
        conv3_3 = self.conv_layer(conv3_2, "conv3_3")
        pool3 = self.max_pool(conv3_3, 'pool3')

        conv4_1 = self.conv_layer(pool3, "conv4_1")
        conv4_2 = self.conv_layer(conv4_1, "conv4_2")
        conv4_3 = self.conv_layer(conv4_2, "conv4_3")
        pool4 = self.max_pool(conv4_3, 'pool4')

        conv5_1 = self.conv_layer(pool4, "conv5_1")
        conv5_2 = self.conv_layer(conv5_1, "conv5_2")
        conv5_3 = self.conv_layer(conv5_2, "conv5_3")
        pool5 = self.max_pool(conv5_3, 'pool5')

        # 事实上，我们训练的部分只有一下两个神经网络层，但这两个就足以我们训练自己的模型了
        self.flatten = tf.reshape(pool5, [-1, 7 * 7 * 512])
        self.fc_6 = tf.layers.dense(self.flatten, 256, tf.nn.relu, name='fc6')
        self.out = tf.layers.dense(self.fc_6, 1, name='fc_out')
        self.sess = tf.Session()
        # 这里我们进行判断是 区分训练还是预测，这两个方法都要用到整个类，所以这里通过有无文件路径进行了判断
        if restore_from:
            saver = tf.train.Saver()
            saver.restore(self.sess, restore_from)
        else:  # training graph
            self.loss = tf.losses.mean_squared_error(labels=self.tfy, predictions=self.out)
            self.train_op = tf.train.RMSPropOptimizer(0.001).minimize(self.loss)
            self.sess.run(tf.global_variables_initializer())
        return super().__init__(**kwargs)
    # 这里我们添加我们的卷积层，注意我们里面的参数都是来自于vgg16.npy


    def conv_layer(self, conv_in, name):
        with tf.variable_scope(name):
            conv = tf.nn.conv2d(conv_in, self.data_dict[name][0], [1, 1, 1, 1], padding='SAME')
            l_out = tf.nn.relu(tf.nn.bias_add(conv, self.data_dict[name][1]))
            return l_out
        # 池化层


    def max_pool(self, conv_in, name):
        return tf.nn.max_pool(conv_in, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding='SAME', name=name)
        # 类中的训练函数


    def train(self, x, y):
        loss, _ = self.sess.run([self.loss, self.train_op], {self.tfx: x, self.tfy: y})
        return loss
        # 类中预测函数


    def predict(self, paths):
        #这里是预测方法，我改成了批量预测，你也可以根据自己的需求进行修改。
        count=0 #车牌的数量
        sum=0 #非车牌的数量
        #fig, axs = plt.subplots(1, 2)
        for i, path in enumerate(paths):
            x = load_img("pred_data/No/"+path)
            length = self.sess.run(self.out, {self.tfx: x})
            #axs[i].imshow(x[0])
            if length < 80:
                animal_ = 'This is not a plate!'
                print(path+'不是车牌')
                sum=sum+1
            else:
                print(path+'是车牌')
                animal_ = 'This is a plate!'
                count=count+1
            #axs[i].set_title(animal_ + ' body length: %.1f cm' % length)
            #axs[i].set_xticks(());
            #axs[i].set_yticks(())
        plt.show()
        return count,sum
        #plt.savefig('table.png')
        # 类中保存模型函数


    def save(self, path='model/transform_learning'):
        saver = tf.train.Saver()
        saver.save(self.sess, path, write_meta_graph=False)


# 这里是主函数的训练函数，进行整体训练
def main_train():
    tigers_x, cats_x, tigers_y, cats_y = load_data()

    plt.hist(tigers_y, bins=20, label='Tigers')
    plt.hist(cats_y, bins=10, label='Cats')
    plt.legend()
    plt.xlabel('length')
    plt.ion()
    plt.show()  # 显示猫和老虎的身长直方图

    xs = np.concatenate(tigers_x + cats_x, axis=0)
    ys = np.concatenate((tigers_y, cats_y), axis=0)  # 这里将数据混合，注意顺序

    vgg16 = Vgg16(vgg16_npy_path='vgg16.npy')  # 这里将vgg16.npy添加进来，初始化类

    for i in range(100):
        batch_idx = np.random.randint(0, len(xs), 6)
        loss = vgg16.train(xs[batch_idx], ys[batch_idx])
        print(i, ":get the train loss/", loss)

    vgg16.save('model/transform_learning')  # 这里结束时保存模型


# 这是预测函数，我们采用了两张图片来进行预测
def main_to_pred():
    vgg_pred = Vgg16(vgg16_npy_path='vgg16.npy', restore_from='model/transform_learning')
    #参考博客只是预测两张图片，我做的批量预测，所以改成了以下方法。
    #在pred_data文件夹下新建两个子文件夹，因为我做的是车牌二分类，所以新建了Yes和No两个文件夹，分别
    #存放车牌和非车牌。
    #当测试给非车牌的时候，是以下写法，当预测车牌的时候将No改成Yes即可，记得改预测函数中的Yes和No呀
    for root, dirs, files in os.walk('./pred_data/No/'):
        #vgg_pred.predict(['pred_data/10b76f8-0.jpg', 'pred_data/8aab94-2.jpg'])
        result=[]
        result=vgg_pred.predict(files)
        print("车牌的数量为：")
        print(result[0])
        print("非车牌的数量为：")
        print(result[1])


# 注意主函数，先训练，在预测，分开完成！！！！！！
if __name__ == '__main__':
    main_train()
    main_to_pred()