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 os
def 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 image
def 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 polys
def 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 polys
def 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_map
class CCPD2019_Dataset:
def __init__(self, FLAGS):
self.FLAGS = FLAGS
ccpd_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) > 0
self.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_filelist
else:
filelist = self.valid_filelist
ds = 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]其他部分的代码就基本一致了。然后就开始训练即可。
4、验证模型
运行eval.py文件，可以看到我们训练好的模型的运行效果，如下图所示，
可以看到，不管是水平的还是斜的车牌，都能准确识别出来了，而且，对于不是车牌的文字，它也不会误当车牌。
5、完整代码
https://mianbaoduo.com/o/bread/YZWcl5tw