TensorFlow入门教程(28)车牌识别之使用EAST模型进行车牌检测(四)

本文主要是介绍TensorFlow入门教程(28)车牌识别之使用EAST模型进行车牌检测(四)，希望对大家解决编程问题提供一定的参考价值，需要的开发者们随着小编来一起学习吧！

#
#作者：韦访
#博客：https://blog.csdn.net/rookie_wei
#微信：1007895847
#添加微信的备注一下是CSDN的
#欢迎大家一起学习
#

1、概述

上一讲，我们实现了EAST的代码，但是我们使用的是ICDAR2017数据集，实现的是对自然场景下的文本检测，现在，我们在原来的代码的基础上，来实现针对车牌的检测。

环境配置：

操作系统：Ubuntu 64位

显卡：GTX 1080ti

Python：Python3.7

TensorFlow：2.3.0

2、CCPD2019数据集

下载链接： https://github.com/detectRecog/CCPD

CCPD2019是中科大开源的数据集，用于车牌检测以及识别的任务。数据集包含了远/近距离、水平/倾斜角度、不同光照、不同天气等等的包含车牌的照片。但是它也有很多缺陷，比如，大部分都是皖A车牌，都是7位数的燃油车牌，没有新能源车牌等等，不过作为学习使用也够了。

下载解压数据集后，得到如下所示文件夹，

其中，ccpd_开头的文件夹下都是图片数据，splits文件夹下则是文本文件。splits文件夹如下图所示，

随便打开一个文本看看，如下图所示，

可以看到，ccpd_blur.txt文件里的每一行都指向ccpd_blur文件夹下的一个图片文件，其他的也是类似的。

CCPD车牌坐标等信息是直接写在文件名里，格式如下，

info1-info2-info3-info4-info5-info6-info7.jpg

其含义如下，

info1：地域area

info2：倾斜程度Tilt degree

info3：标注框坐标Bounding box coordinates

info4：四个车牌顶角坐标Four vertices locations

info5：车牌号码License plate number

info6：车牌区域亮度信息Brightness

info7：车牌区域模糊程度Blurriness

我们写个代码来分别根据标注框坐标和顶点坐标画出车牌框，代码如下，

import cv2
import numpy as np
import csv 
import osdef line(image, polys, color=(0,0,255)):    image = cv2.line(image, tuple(polys[0]), tuple(polys[1]), color, thickness=5)image = cv2.line(image, tuple(polys[1]), tuple(polys[2]), color, thickness=5)image = cv2.line(image, tuple(polys[2]), tuple(polys[3]), color, thickness=5)image = cv2.line(image, tuple(polys[3]), tuple(polys[0]), color, thickness=5)    return imagedef get_polys(image_file):polys = []print("image_file:", image_file)if not os.path.exists(image_file):return np.array(polys, dtype=np.float32)parts = image_file.split("-")polys_part = parts[3].split("_")print(polys_part)poly = []for p in polys_part:poly.append(p.split("&"))polys.append(np.asarray(poly).astype(np.int32))print("polys:", polys)return polysdef get_rectangle(image_file):polys = []print("image_file:", image_file)if not os.path.exists(image_file):return np.array(polys, dtype=np.float32)parts = image_file.split("-")rect_part = parts[2].split("_")poly = []for p in rect_part:poly.append(p.split("&"))poly = np.asarray(poly).astype(np.int32)    polys.append([poly[0], [poly[1][0], poly[0][1]], poly[1], [poly[0][0], poly[1][1]]])print(polys)return polysdef show(filename):rects = get_rectangle(filename) polys = get_polys(filename)   image = cv2.imread(filename)for rect in rects:print("rect:", rect)image = line(image, rect, (0,0,255))for poly in polys:print("poly:", poly)image = line(image, poly, (255,0,0))image = cv2.resize(image, (512, 512))cv2.imshow("demo1", image)cv2.imwrite("demo1.jpg", image)cv2.waitKey(0)show("ccpd_blur/0359-5_21-151&285_417&398-417&398_179&377_151&285_389&306-0_0_4_33_32_25_12-59-4.jpg")

运行结果，

3、修改代码

现在，我们基于上一讲的代码来修改。

3.1导入坐标

显然，数据集中使用车牌4个顶点坐标更适合我们的EAST算法，根据文件名获取车牌4个顶点坐标的代码如下，

'''
ccpd数据集的坐标等信息直接在文件名中，所以解析文件名即可拿到坐标信息
'''
def load_ccpd_polys(filename):text_polys = []ignored_label = []parts = filename.split("-")polys_part = parts[3].split("_")poly = []for p in polys_part:poly.append(p.split("&"))text_polys.append(np.asarray(poly).astype(np.int32))  ignored_label.append(False)# print(text_polys)  return np.array(text_polys, dtype=np.float32), np.array(ignored_label, dtype=np.bool)

3.2、导入数据集并创建tf.data.Dataset

数据增强的代码跟上一讲中类似，代码如下，

def parse_func(image_file, FLAGS):[image, score_map, geo_map] = tf.py_function(lambda image_file: preprocess(image_file, FLAGS=FLAGS), [image_file], [tf.float32, tf.float32, tf.float32])return image, score_map, geo_mapclass CCPD2019_Dataset:def __init__(self, FLAGS): self.FLAGS = FLAGSccpd_train_dir = os.path.join(FLAGS.ccpd_dataset_dir, FLAGS.ccpd_train_txtfile)ccpd_valid_dir = os.path.join(FLAGS.ccpd_dataset_dir, FLAGS.ccpd_valid_txtfile)        assert os.path.exists(ccpd_train_dir) and os.path.exists(ccpd_valid_dir)        self.train_filelist = get_iamges_from_txt(FLAGS.ccpd_dataset_dir, ccpd_train_dir)self.valid_filelist = get_iamges_from_txt(FLAGS.ccpd_dataset_dir, ccpd_valid_dir)assert len(self.train_filelist) > 0 and len(self.valid_filelist) > 0# 如果有ICDAR数据集，那么将该数据集也加入训练，让其全部当背景图，这样模型对不是车牌的字符能更好的判断if os.path.exists(FLAGS.icadr_dataset_dir) and os.path.exists(FLAGS.icadr_dataset_dir):icadr_train_dir = os.path.join(FLAGS.icadr_dataset_dir, "ch8_training_images")icadr_valid_dir = os.path.join(FLAGS.icadr_dataset_dir, "ch8_validation_images")assert os.path.exists(icadr_train_dir) and os.path.exists(icadr_valid_dir)        icadr_train_filelist = get_images_from_dir(icadr_train_dir)icadr_valid_filelist = get_images_from_dir(icadr_valid_dir)assert len(icadr_train_filelist) > 0 and len(icadr_valid_filelist) > 0self.train_filelist.extend(icadr_train_filelist)self.valid_filelist.extend(icadr_valid_filelist)self.train_filelist = np.asarray(self.train_filelist)self.valid_filelist = np.asarray(self.valid_filelist)np.random.shuffle(self.train_filelist)np.random.shuffle(self.valid_filelist)def create_dataset(self, subset):if subset == "train":filelist = self.train_filelistelse:filelist = self.valid_filelistds = tf.data.Dataset.from_tensor_slices(filelist)ds = ds.map(lambda image_file: parse_func(image_file, FLAGS=self.FLAGS))ds = ds.batch(self.FLAGS.batch_size)ds = ds.prefetch(AUTOTUNE)return ds

从代码中可以看到，如果存在ICDAR数据集，最好将它也加入训练，这样模型能对非车牌的字符能有更好的判断，只用CCPD数据集的话，对于一些不是车牌的字符，模型可能误认为它是车牌的。

3.2、preprocess

因为同时使用两个数据集，所以preprocess函数也要微改一下，代码如下，

def preprocess(image_file, FLAGS, is_test=False):# start = time.time()# print("------start------------")# print(image_file)if not is_test:# tf.data.Dataset需要这样转换一下才能用image_file = image_file.numpy().decode("utf-8")if FLAGS.task == "plate":_,filename = os.path.split(image_file)if filename.startswith("img_"):# icdar数据集的数据polys = []ignored_labels = []else: # ccpd数据集polys, ignored_labels = load_ccpd_polys(image_file)   else:polys, ignored_labels = load_icdar_polys(image_file)   # print("load_icdar_polys time: ", time.time() - start) image = cv2.imread(image_file)image, polys = random_scale_image(image, polys)# print("random_scale_image time: ", time.time() - start) image, polys, ignored_labels = random_crop_area(FLAGS, image, polys, ignored_labels)# print("random_crop_area time: ", time.time() - start) image, polys = pad_image(image, polys, FLAGS.input_size)# print("pad_image time: ", time.time() - start) image, polys = resize(image, polys, FLAGS.input_size)# print("resize time: ", time.time() - start) score_map, geo_map = map_generator(FLAGS, image, polys, ignored_labels)image = (image / 127.5) - 1# print("map_generator time: ", time.time() - start)return image, score_map[::4, ::4, np.newaxis], geo_map[::4, ::4]

其他部分的代码就基本一致了。然后就开始训练即可。