TensorFlow 2.0 简单车牌识别

TensorFlow 2.0 简单车牌识别

说明：使用tensorflow2 训练模型，然后冻结模型，使用opencv 的 dnn 模块进行推断。这只是个简单的demo,本质上就是个分类模型。主要工作还是在车牌字符的分割。在推断部分只要主意好图片的格式就好。

GitHub

1、tensorflow 训练模型

（1）数据集准备

使用 百度 ai studio 上提供的车牌数据集

部分数据集图片（图片命名还是比较乱的）：

（2）数据集处理

data_path = 'E:\\code\\tf\\proj\\car_num\\data'
character_folders = os.listdir(data_path)
label = 0
LABEL_temp = {}
if(os.path.exists('./train_data.list')):
    os.remove('./train_data.list')
if(os.path.exists('./test_data.list')):
    os.remove('./test_data.list')
for character_folder in character_folders:
    with open('./train_data.list', 'a') as f_train:
        with open('./test_data.list', 'a') as f_test:
            if character_folder == '.DS_Store' or character_folder == '.ipynb_checkpoints' or character_folder == 'data23617':
                continue
            print(character_folder + " " + str(label))
            LABEL_temp[str(label)] = character_folder     #存储一下标签的对应关系
            character_imgs = os.listdir(os.path.join(data_path, character_folder))
            for i in range(len(character_imgs)):
                if i%10 == 0:
                    f_test.write(os.path.join(os.path.join(data_path, character_folder), character_imgs[i]) + "\t" + str(label) + '\n')
                else:
                    f_train.write(os.path.join(os.path.join(data_path, character_folder), character_imgs[i]) + "\t" + str(label) + '\n')
    label = label + 1
print('图像列表已生成')

all_image_paths = []
all_image_labels = []
test_image_paths = []
test_image_labels = []
with open('./train_data.list', 'r') as f:
    lines = f.readlines()
    for line in lines:
        img, label = line.split('\t')
        all_image_paths.append(img)
        all_image_labels.append(int(label))
with open('./test_data.list', 'r') as f:
    lines = f.readlines()
    for line in lines:
        img, label = line.split('\t')
        test_image_paths.append(img)
        test_image_labels.append(int(label))
def preprocess_image(image):
    image = tf.image.decode_jpeg(image, channels=3)
    image = tf.cast(image,dtype=tf.float32)
    image = tf.image.resize(image, [20, 20])
    image /= 255.0  # normalize to [0,1] range
    return image

def load_and_preprocess_image(path,label):
    image = tf.io.read_file(path)
    return preprocess_image(image), label


ds = tf.data.Dataset.from_tensor_slices((all_image_paths, all_image_labels))
train_data = ds.map(load_and_preprocess_image).batch(64)
db = tf.data.Dataset.from_tensor_slices((test_image_paths, test_image_labels))
test_data = db.map(load_and_preprocess_image).batch(64)

（3）搭建网络训练模型

def train_model(train_data,test_data):
    #构建模型
    network = keras.Sequential([
        keras.layers.Conv2D(32, kernel_size=[3, 3], padding="same", activation=tf.nn.relu),
        keras.layers.BatchNormalization(),
        keras.layers.MaxPool2D(pool_size=[2, 2], strides=2, padding='same'),
        keras.layers.Conv2D(64, kernel_size=[3, 3], padding="same", activation=tf.nn.relu),
        keras.layers.BatchNormalization(),
        keras.layers.MaxPool2D(pool_size=[2, 2], strides=2, padding='same'),
        keras.layers.Conv2D(64, kernel_size=[3, 3], padding="same", activation=tf.nn.relu),
        keras.layers.BatchNormalization(),
        keras.layers.Flatten(),
        keras.layers.Dense(512, activation='relu'),
        keras.layers.Dropout(0.5),
        keras.layers.Dense(128, activation='relu'),
        keras.layers.Dense(65)])
    network.build(input_shape=(None, 20, 20, 3))
    network.summary()

    reduce_lr = ReduceLROnPlateau(monitor='val_loss', patience=10, mode='auto')
    network.compile(optimizer=optimizers.SGD(lr=0.001),
                    loss=tf.losses.SparseCategoricalCrossentropy(from_logits=True),
                    metrics=['accuracy'])
    network.fit(train_data, epochs=100, validation_data=test_data, callbacks=[reduce_lr])
    network.evaluate(test_data)
    tf.saved_model.save(network, 'E:\\code\\tf\proj\\car_num\\model\\')
train_model(train_data,test_data)

模型测试集上的准确率大概 97%，忘了截图了。

2、模型转换

import os
import numpy as np
import tensorflow as tf
from tensorflow import keras
from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2
import matplotlib.pyplot as plt
os.environ["TF_CPP_MIN_LOG_LEVEL"] = '2'
print('tf版本：',tf.__version__)

DEFAULT_FUNCTION_KEY = "serving_default"
loaded = tf.saved_model.load('.\\model\\')
network = loaded.signatures[DEFAULT_FUNCTION_KEY]
print(list(loaded.signatures.keys()))
print('加载 weights 成功')

# Convert Keras model to ConcreteFunction
full_model = tf.function(lambda x: network(x))
full_model = full_model.get_concrete_function(
tf.TensorSpec(network.inputs[0].shape, network.inputs[0].dtype))

# Get frozen ConcreteFunction
frozen_func = convert_variables_to_constants_v2(full_model)
frozen_func.graph.as_graph_def()

layers = [op.name for op in frozen_func.graph.get_operations()]
print("-" * 50)
print("Frozen model layers: ")
for layer in layers:
    print(layer)
print("Frozen model inputs: ")
print(frozen_func.inputs)
print("Frozen model outputs: ")
print(frozen_func.outputs)
# Save frozen graph from frozen ConcreteFunction to hard drive
tf.io.write_graph(graph_or_graph_def=frozen_func.graph,
		logdir="./frozen_models",
		name="frozen_graph.pb",
		as_text=False)

3、OpenCV C++ 推断。

#include <iostream>
#include <opencv2/opencv.hpp>
#include <opencv2/dnn.hpp>

#define DEBUG
using namespace std;
using namespace cv;
using namespace cv::dnn;

string label_list[] = {"0", "1", "2", "3", "4", "5", "6", "7", "8", "9", "A", "B",
              "C", "川",
              "D", "E", "鄂", "F",
              "G", "赣", "甘", "贵", "桂",
              "H", "黑", "沪",
              "J", "冀", "津", "京", "吉",
              "K", "L", "辽", "鲁", "M", "蒙", "闽",
              "N", "宁",
              "P", "Q", "青", "琼",
              "R", "S", "陕", "苏", "晋",
              "T", "U", "V", "W", "皖",
              "X", "湘", "新",
              "Y", "豫", "渝", "粤", "云",
              "Z", "藏", "浙"};
int main()
{
    string  model_path = "../infer/frozen_graph.pb";
    Net net = readNetFromTensorflow(model_path);
    Mat license_plate = imread("E:\\code\\tf\\proj\\car_num\\zf.png",1);
    int image_h = license_plate.rows;
//    cout<<image_h<<endl;
    Mat gray_plate;
    Mat binary_plate;
    cvtColor(license_plate,gray_plate,COLOR_BGR2GRAY);
    threshold(gray_plate,binary_plate,175,255,THRESH_OTSU);


    vector<uint> pix(binary_plate.cols,0);

    for(int i = 0; i < binary_plate.cols;i++)
    {
        pix.push_back(0);
        for(int j = 0;j<binary_plate.rows;j++)
        {
           pix[i] = pix[i] +  binary_plate.at<uchar>(j,i);
        }
//        printf("%d ",pix[i]);
//        cout<<endl;
    }

    uint i = 0;
    uint num = 0;
    vector<Point> index_range;
    while(i < pix.size())
    {
        if (pix[i] == 0)
        {
            i +=1 ;
        }
        else
        {
            uint index = i + 1;
            while(pix[index] != 0)
            {
                index += 1;
            }
            index_range.push_back(Point(i,index-1));
            num += 1;
            i = index;
        }
    }
//    cout<<index_range.size();
    vector<Mat> seg_img;
    for(uint i = 0,num = 0;i<index_range.size();i++)
    {
        if(i == 2)
            continue;
        Mat img(binary_plate,Rect(index_range[i].x,0,index_range[i].y-index_range[i].x,binary_plate.rows));
        Mat temp_img = Mat::zeros(img.size(),CV_8UC1);
//        cout<<index_range[i].y<<","<<index_range[i].x<<endl;
        int pad = (binary_plate.rows -(index_range[i].y - index_range[i].x))/2;
        img.copyTo(temp_img);
        copyMakeBorder(temp_img,temp_img,0,0,pad,pad, cv::BORDER_CONSTANT,Scalar(0,0,0));

        imwrite(to_string(num++)+".png",temp_img);
        imshow(to_string(i+10),temp_img);
    }
    cout <<"车牌识别结果:";
    for(int i = 0;i<7;i++)
    {
        Mat frame = imread(to_string(i)+".png",1);
        Mat frame_32F;
        frame.convertTo(frame_32F,CV_32FC1);

        Mat blob = blobFromImage(frame_32F/255.0,
                                      1.0,
                                      Size(20,20),
                                      Scalar(0,0,0));

//        cout<<(blob.size);
        net.setInput(blob);
        Mat out = net.forward();
        Point maxclass;
        minMaxLoc(out, NULL, NULL, NULL, &maxclass);
        cout <<label_list[maxclass.x];
    }
    cout<<endl;

#ifdef DEBUG
    imshow("1",license_plate);
    imshow("2",gray_plate);
    imshow("3",binary_plate);
    while(1)
        if(waitKey(0) == '1')
            break;
#endif
    return 0;
}