TensorFlow入门教程(26)车牌识别之文本检测模型EAST代码实现(二)
#
#作者：韦访
#博客：https://blog.csdn.net/rookie_wei
#微信：1007895847
#添加微信的备注一下是CSDN的
#欢迎大家一起学习
#
1、概述
上一讲，我们简单是介绍了EAST的论文，有了理论依据以后，接下来我们来一步一步实现代码。为了照顾不做车牌检测的网友，我们先来实现通用的自然场景下的文本检测，再基于此实现车牌检测。
环境配置：
操作系统：Ubuntu 64位
显卡：GTX 1080ti
Python：Python3.7
TensorFlow：2.3.0
2、ICDAR2017数据集
文字检测有很多公开的数据集，我这里选择了ICDAR2017，因为这个数据集支持的语言种类比较多，而且数据集大小也不是几百G的那种巨无霸。
官网链接：https://rrc.cvc.uab.es/?ch=8&com=downloads
百度网盘：https://pan.baidu.com/s/1S0a8cL743ZjvMzs6IZ_vrA 密码: k6oj
数据集一共由11个压缩包组成，包含了训练集和验证集的数据，我们将ch8_training开头的压缩包都解压到ch8_training_images文件夹下，将ch8_validation开头的压缩包解压到ch8_validation_images文件夹下，这样比较方便我们操作。
上图是ch8_training_images文件夹下的文件总数，可以看到一共有14400个文件，其中有7200个TXT文本文件，和7200个jpg或png图片文件，他们通过文件名来一一对应。比如，图片img_1.png对应的文件是gt_img_1.txt。gt_img_1.txt文件的内容如下图所示，
上图中，每一行代表一个文本框信息，以逗号为分隔符，其中前8个字段代表的是文本框的四个顶点的坐标，分别为左上、右上、右下和左下。第9个字段表示文本框内的文字属于什么语言。最后一个字段表示文本框内的文字，”###”表示无法识别文本框内的文字内容，我们一般选择忽略这种文本框。
3、数据增强
3.1、读取文本框坐标
首先，我们要根据图片的文件名找到其对应的TXT文本文件（TXT文件名只是比图片文件名多了个“gt_”前缀和后缀为“.txt”），然后再解析其中所有的文本框的坐标信息。由于”###”的表示不知道文本框内的文字内容，所以这种文本框我们选择忽略，将它们在ignored_label列表中的值置为“True”。代码如下，
'''
获取ICDAR数据集的图片的文件名所对应的标签文本文件（包含文本框坐标等信息）
'''
def get_icdar_text_file(image_file):
# 文本文件名跟图片文件名一样，只是多了个gt_前缀
txt_file = image_file.replace(os.path.basename(image_file).split('.')[1], 'txt')
txt_file_name = os.path.basename(txt_file)
txt_file = txt_file.replace(txt_file_name, 'gt_' + txt_file_name)
return txt_file
'''
通过txt导入对应图片的文本框坐标等信息
'''
def load_icdar_polys(image_file):
polys = []
ignored_label = []
# 找到对应的文本文件
text_file = get_icdar_text_file(image_file)
if not os.path.exists(text_file):
return np.array(polys, dtype=np.float32)
with open(text_file, 'r', encoding="utf-8") as fd:
reader = csv.reader(fd)
for line in reader:
# strip BOM. \ufeff for python3, \xef\xbb\bf for python2
line = [i.strip('\ufeff').strip('\xef\xbb\xbf') for i in line]
# 获取每行的文本框坐标
x1, y1, x2, y2, x3, y3, x4, y4 = list(map(float, line[:8]))
poly = np.asarray([[x1, y1], [x2, y2], [x3, y3], [x4, y4]])
polys.append(poly)
# 每行的最后一个属性，即文本框内的文字
label = line[-1]
# 如果文字是###，表示该文本框内的文字不清楚，我们忽略这种文本框
if label == '*' or label == "###":
ignored_label.append(True)
else:
ignored_label.append(False)
return np.array(polys, dtype=np.float32), np.array(ignored_label, dtype=np.bool)3.2、随机缩放图片
随机缩放是数据增强中常用的手段，我们随机缩放图片的宽和高，但是每次缩放的宽高比例不能相差太大，否则就失真了。代码如下，
'''
随机缩放图片和文本框坐标
'''
def random_scale_image(image, polys):
random_scale = np.array([0.5, 0.75, 1., 1.25, 1.5])
rd_scale = np.random.choice(random_scale)
x_scale_variation = np.random.randint(-10, 10) / 100.
y_scale_variation = np.random.randint(-10, 10) / 100.
x_scale = rd_scale + x_scale_variation
y_scale = rd_scale + y_scale_variation
image = cv2.resize(image, dsize=None, fx=x_scale, fy=y_scale)
if len(polys) > 0:
polys[:, :, 0] *= x_scale
polys[:, :, 1] *= y_scale
return image, polys3.3、随机裁剪
接下来是随机裁剪图片了，分两种情况，
第一种是裁剪后的图片只有背景，没有文本框，让模型学会识别背景图。
第二种是裁剪后的图片至少包含一个文本框，让模型学会识别文本框。需要注意的是，裁剪后，如果是带文本框的，那么，文本框的坐标也要跟裁剪后的图片的坐标对应得上，文本框是否是应该忽略的标签信息也不能丢。
先来看整体的代码，再具体看应该怎么裁剪，整体代码如下，
'''
随机截取图片中的一个区域
'''
def random_crop_area(FLAGS, image, polys, ignored_labels):
# DEBUG = True
h, w, _ = image.shape
# 计算最小截取宽度和高度
min_crop_w = np.round(FLAGS.min_crop_side_ratio * w).astype(np.int32)
min_crop_h = np.round(FLAGS.min_crop_side_ratio * h).astype(np.int32)
# 如果该图片没有文本框信息，则直接随机截取
if len(polys) < 1:
return random_crop_backgroup_area(FLAGS, image, min_crop_w, min_crop_h)
rectangle_polys = []
crop_image = []
crop_polys = []
crop_ignored_labels = []
# 将文本框变换成矩形的形式
for poly in polys:
# round
poly = np.round(poly, decimals=0).astype(np.int32)
min_x = np.min(poly[:, 0])
max_x = np.max(poly[:, 0])
min_y = np.min(poly[:, 1])
max_y = np.max(poly[:, 1])
rectangle_polys.append([[min_x, min_y], [max_x, min_y], [max_x, max_y], [min_x, max_y]])
rectangle_polys = np.asarray(rectangle_polys)
# 随机获取背景截图或带文本框的截图
if np.random.rand() < FLAGS.background_ratio:
crop_image, crop_polys, crop_ignored_labels = random_crop_backgroup_area_with_polys(image, rectangle_polys, min_crop_w, min_crop_h)
# print("background")
else:
crop_image, crop_polys, crop_ignored_labels = random_crop_text_area(image, polys, rectangle_polys, ignored_labels, min_crop_w, min_crop_h)
# print("text")
# 如果文本框坐标长度和截图的长度都为0，则表示截取失败，则直接返回原图和原坐标
if len(crop_image) < 1 and len(crop_polys) < 1:
crop_image = image
crop_polys = polys
crop_ignored_labels = ignored_labels
if DEBUG:
for poly in crop_polys:
crop_image = draw_line(crop_image, poly)
if len(crop_image) > 0:
crop_image = cv2.resize(crop_image, (512, 512))
image = cv2.resize(image, (800, 800))
cv2.imshow("crop_image", crop_image)
cv2.imshow("image", image)
cv2.waitKey(0)
# show(image)
return crop_image, crop_polys, crop_ignored_labels上面代码中，如果送进来的图片没有文本框信息，则随机截取，然后返回。如果送进来的图片有文本框，那么，根据设置的FLAGS.background_ratio随机选择这次是裁剪背景图还是裁剪包含文本框的图，然后返回裁剪后的图片信息、文本框坐标和忽略标签即可。
3.3.1、随机裁剪背景图
先来看看怎么随机裁剪背景图。函数名为random_crop_backgroup_area_with_polys，代码如下，
'''
随机截取没有文字的背景区域
'''
def random_crop_backgroup_area_with_polys(image, rectangle_polys, min_crop_w, min_crop_h):
# DEBUG = True
crop_image = []
crop_polys = []
crop_ignored_labels = []
h, w, _ = image.shape
# 随机生成要截取的图片的x轴的起始坐标
crop_x = np.random.randint(0, w - min_crop_w - 1)
if DEBUG:
cv2.circle(image, (crop_x, 0), 2, (0,255,0), 4)
cv2.imshow("image", image)
cv2.waitKey(0)
# 随机生成要截取的图片的x轴的x轴宽度
crop_w = np.random.randint(min_crop_w, w - crop_x - 1)
if DEBUG:
cv2.line(image, (crop_x, 0), (crop_x + crop_w, 0), (255,0,0), 4)
cv2.imshow("image", image)
cv2.waitKey(0)
# print("crop_x:", crop_x, " crop_w:", crop_w)
# print("len polys:", len(polys))
# 找到x轴跟点crop_x到crop_x+crop_w有交集的文本框
relevant_rectangle_polys = []
for poly in rectangle_polys:
if (crop_x >= poly[0][0] and crop_x <= poly[1][0]) or (crop_x + crop_w >= poly[0][0] and crop_x + crop_w <= poly[1][0]) or (crop_x <= poly[0][0] and crop_x + crop_w >= poly[1][0]):
relevant_rectangle_polys.append(poly)
# print("len relevant_rectangle_polys:", len(relevant_rectangle_polys))
# 将与截取图相关的文本框的y轴标记，被标记的区域是不能选的
h_array = np.zeros((h), dtype=np.int32)
for poly in relevant_rectangle_polys:
# print(poly)
min_h = np.min(poly[:, 1])
max_h = np.max(poly[:, 1])
# print("min_h:", min_h, " max_h:", max_h)
h_start = np.where(min_h - min_crop_h > 0, min_h - min_crop_h, 0)
h_end = np.where(max_h + min_crop_h < h, max_h + min_crop_h, h)
# print("h_start:", h_start, " h_end:", h_end)
h_array[h_start : h_end] = 1
# print("h_array:", h_array)
# 将y轴中自底向上的min_crop_h长度的区域标记
h_array[h-min_crop_h : ] = 1
# 算出未被标记的y轴坐标，要截取的图片的y轴起始坐标可以在这个区域随机生成
h_axis = np.where(h_array == 0)[0]
# print("h_axis:", h_axis)
if len(h_axis) > 0:
# print("h_axis:", h_axis)
# 随机获取截取图的y轴起始坐标
crop_y = np.random.choice(h_axis, size=1)[0]
if DEBUG:
cv2.circle(image, (0, crop_y), 2, (0,255,0), 4)
cv2.imshow("image", image)
cv2.waitKey(0)
# print("h_axis:", h_axis, " crop_y:", crop_y)
# 找到h_axis中，crop_y往上的第一个不连续的点的坐标，用于限定随机生成的截取高度
len_h_axis = len(h_axis)
# print("h_axis.index(crop_y):", np.argwhere(h_axis == crop_y), " crop_y:", crop_y)
discontinuity = 0
for i in range(np.argwhere(h_axis == crop_y)[0][0], len_h_axis, 1):
# print("i:", i, " h_axis[i]:", h_axis[i], " h_axis[i]+1:", h_axis[i+1] - 1)
if i < len_h_axis - 1 and h_axis[i] != h_axis[i+1] - 1:
discontinuity = h_axis[i]
break
if i == len_h_axis - 1:
discontinuity = h_axis[i]
# print("crop_y:", crop_y, "discontinuity:", discontinuity)
if discontinuity != 0:
# 随机生成高度
crop_h = np.random.randint(min_crop_h, discontinuity + min_crop_h - crop_y + 1)
if DEBUG:
cv2.line(image, (0, crop_y), (0, crop_y + crop_h), (255,0,0), 4)
cv2.imshow("image", image)
print("crop_x:", crop_x, " crop_w:", crop_w)
print("crop_y:", crop_y, " crop_h:", crop_h)
image = cv2.line(image, (crop_x, crop_y), (crop_x + crop_w, crop_y), (255,0,0), thickness=2)
image = cv2.line(image, (crop_x + crop_w, crop_y), (crop_x + crop_w, crop_y + crop_h), (255,0,0), thickness=2)
image = cv2.line(image, (crop_x + crop_w, crop_y + crop_h), (crop_x, crop_y + crop_h), (255,0,0), thickness=2)
image = cv2.line(image, (crop_x, crop_y + crop_h), (crop_x, crop_y), (255,0,0), thickness=2)
cv2.waitKey(0)
# 截取图像
crop_image = image[crop_y:crop_y+crop_h, crop_x:crop_x+crop_w, :]
return crop_image, crop_polys, crop_ignored_labels这部分代码可能有点难理解，我看了其他开源代码，都是采用“碰运气式”的裁剪，也就是说，先把所有文本框的x和y轴映射出来，这部分区域都不能选，然后再随机截取其他区域的，如果截取的区域包含了文本框，就再随机截取，直到不包含文本框为止。这种方法比较简单，但是效率比较低。我上面裁剪代码的思路是，
- 先将所有的文本框变成与x轴平行的矩形的形式，即找到四个顶点中的最大和最小的x坐标和y坐标组成的矩形。
- 然后随机生成要截取的图片的x轴的起始坐标，再随机生成要截取图片的宽度，如下图示。
- 找到所有x轴与步骤“2”中的直线有交集的文本框，称为相关文本框。
- 将所有相关文本框的y轴进行映射，对h-min_crop_h区域也进行映射。
- 在步骤“4”中，在未被映射的区域中随机选择截取图的y轴的起始坐标，如下图所示。
- 从步骤”5”中的y轴起始坐标往上找（这里的往上对应于图片是往下，因为计算机中，原点的坐标一般都是左上角），在未被映射列表中找到第一个不连续点的坐标，可以在这中间随机生成截取图片的高度，如下图所示。
- 最后，根据x轴、y轴的起点坐标，以及宽和高，就得到要截取的矩形框坐标，如下图所示。
3.3.2、随机裁剪包含文本框的截图
接下来看看随机裁剪带文本框的截图，代码如下，
'''
随机截取包含文本框的区域
'''
def random_crop_text_area(image, polys, rectangle_polys, ignored_labels, min_crop_w, min_crop_h):
# DEBUG = True
crop_image = []
crop_polys = []
crop_ignored_labels = []
h, w, _ = image.shape
# print("rectangle_polys:", rectangle_polys)
# 标记x轴和y轴中所有文本框映射的区域，该区域不能为起始坐标
w_array = np.zeros((w), dtype=np.int32)
h_array = np.zeros((h), dtype=np.int32)
padding = 1
for poly in rectangle_polys:
# 求该文本坐标中的x轴的最大和最小点
minx = np.where(np.min(poly[:, 0]) - padding > 0, np.min(poly[:, 0]) - padding, 0)
maxx = np.where(np.max(poly[:, 0]) + padding > w, w, np.max(poly[:, 0]) + padding)
# 将w_array中对应的文本坐标x轴往外扩展padding设置为1
w_array[minx:maxx] = 1
# 求该文本坐标中的y轴的最大和最小点
miny = np.where(np.min(poly[:, 1]) - padding > 0, np.min(poly[:, 1]) - padding, 0)
maxy = np.where(np.max(poly[:, 1]) + padding > h, h, np.max(poly[:, 1]) + padding)
# 将h_array中对应的文本坐标y轴往外扩展padding设置为1
h_array[miny:maxy] = 1
# 找到x轴中，最右的文本框左上角的x坐标,这个点往后的都标记为1，这些区域不能作为截取点的左上角顶点
txt_rect_max_x = np.max(rectangle_polys[:,:,0])
w_array[txt_rect_max_x:] = 1
# print("txt_rect_max_x:", w_array)
# 找到y轴中，最底部的文本框的左上角的y坐标,这个点往下的都标记为1，这些区域不能作为截取点的左上角顶点
txt_rect_max_y = np.max(rectangle_polys[:,:,1])
h_array[txt_rect_max_y:] = 1
# print("txt_rect_max_y:", h_array)
# 求未被标记的x轴和y轴坐标
w_axis = np.where(w_array == 0)[0]
h_axis = np.where(h_array == 0)[0]
# 如果都被标记了，就没法裁剪了，直接返回空
if len(w_axis) < 1 or len(h_axis) < 1:
return crop_image, crop_polys, crop_ignored_labels
# 随机生成截取图左上角的坐标x和y
crop_x = np.random.choice(w_axis, size=1)[0]
crop_y = np.random.choice(h_axis, size=1)[0]
if DEBUG:
cv2.circle(image, (crop_x, crop_y), 2, (0,255,0), 4)
cv2.imshow("image", image)
cv2.waitKey(0)
# 将坐标x和y往右的所有文本框找出来，这些文本框为相关框
relevant_rectangle_polys = []
for poly in rectangle_polys:
if crop_x <= poly[0][0] and crop_y <= poly[0][1]:
relevant_rectangle_polys.append(poly)
relevant_rectangle_polys = np.asarray(relevant_rectangle_polys)
# 如果没有包含相关框，表示没裁剪到文本框，直接返回空
if len(relevant_rectangle_polys) < 1:
return crop_image, crop_polys, crop_ignored_labels
# print("relevant_rectangle_polys:", relevant_rectangle_polys)
# 将相关框的x轴和y轴进行标记
w_array_relevant = np.zeros((w), dtype=np.int32)
h_array_relevant = np.zeros((h), dtype=np.int32)
for poly in relevant_rectangle_polys:
# 求该文本坐标中的x轴的最大和最小点
minx = np.where(np.min(poly[:, 0]) - padding > 0, np.min(poly[:, 0]) - padding, 0)
maxx = np.where(np.max(poly[:, 0]) + padding > w, w, np.max(poly[:, 0]) + padding)
# 将w_array_relevant中对应的文本坐标x轴往外扩展padding设置为1
w_array_relevant[minx:maxx] = 1
# 求该文本坐标中的y轴的最大和最小点
miny = np.where(np.min(poly[:, 1]) - padding > 0, np.min(poly[:, 1]) - padding, 0)
maxy = np.where(np.max(poly[:, 1]) + padding > h, h, np.max(poly[:, 1]) + padding)
# 将h_array_relevant中对应的文本坐标y轴往外扩展padding设置为1
h_array_relevant[miny:maxy] = 1
# 找到x轴中，最左的文本框左上角的x坐标,这个点往前的都标记为1,如果右下角顶点在这个区域就框不到文本框了
txt_rect_min_x = np.max(relevant_rectangle_polys[:,:,0])
w_array_relevant[:txt_rect_min_x] = 1
# print("w_array_relevant:", w_array_relevant)
# 找到y轴中，最底部的文本框的左上角的y坐标,这个点往上的都标记为1,如果右下角顶点在这个区域就框不到文本框了
txt_rect_min_y = np.max(relevant_rectangle_polys[:,:,1])
h_array_relevant[:txt_rect_min_y] = 1
# print("h_array_relevant:", h_array_relevant)
# x轴从crop_x到crop_x+min_crop_w都标记为1，否则截取的宽度达不到要求
w_array_relevant[crop_x : crop_x+min_crop_w] = 1
# y轴从crop_y到crop_y+min_crop_y都标记为1，否则截取的高度达不到要求
h_array_relevant[crop_y : crop_y+min_crop_h] = 1
# 求未被标记的x轴和y轴坐标
w_axis_relevant = np.where(w_array_relevant == 0)[0]
h_axis_relevant = np.where(h_array_relevant == 0)[0]
# print("w_axis:", w_axis_relevant)
# print("h_axis:", h_axis_relevant)
# 如果都被标记了，表示没法裁剪，直接返回空
if len(w_axis_relevant) < 1 or len(h_axis_relevant) < 1:
return crop_image, crop_polys, crop_ignored_labels
# 随机选择截取图的宽高
crop_w = np.random.choice(w_axis_relevant, size=1)[0]
crop_h = np.random.choice(h_axis_relevant, size=1)[0]
crop_w -= crop_x
crop_h -= crop_y
if DEBUG:
image = cv2.line(image, (crop_x, crop_y), (crop_x + crop_w, crop_y), (255,0,0), thickness=2)
image = cv2.line(image, (crop_x + crop_w, crop_y), (crop_x + crop_w, crop_y + crop_h), (255,0,0), thickness=2)
image = cv2.line(image, (crop_x + crop_w, crop_y + crop_h), (crop_x, crop_y + crop_h), (255,0,0), thickness=2)
image = cv2.line(image, (crop_x, crop_y + crop_h), (crop_x, crop_y), (255,0,0), thickness=2)
cv2.imshow("image", image)
cv2.waitKey(0)
# 截取图像
crop_image = image[crop_y:crop_y+crop_h, crop_x:crop_x+crop_w, :]
# 找到原文本框中的相关框
for poly, label in zip(polys, ignored_labels):
if (crop_x <= poly[0][0] and crop_y <= poly[0][1] and (crop_x + crop_w) >= poly[0][0] and (crop_y + crop_h) >= poly[0][1]):
crop_polys.append(poly)
crop_ignored_labels.append(label)
crop_polys = np.asarray(crop_polys)
crop_ignored_labels = np.asarray(crop_ignored_labels)
# print("crop_x:", crop_x, "crop_y:", crop_y)
# print("crop_polys:", crop_polys)
crop_polys[:,:,0] -= crop_x
crop_polys[:,:,1] -= crop_y
# print("crop_polys after:", crop_polys)
return crop_image, crop_polys, crop_ignored_labels上面代码的思路是：
- 先映射所有文本框的x轴和y轴，这些区域不能被选为起始坐标。
- 找到最右边和最下边的文本框，这个文本框往右和往下的区域都标记为不能选为起始坐标的区域。
- 在未被标记的区域中随机生成x轴和y轴的起始坐标，如下图所示。
- 找到起始坐标往右和往下的所有文本框，称为相关文本框。
- 将所有相关文本框的x和y轴进行映射，并且将最左文本框往左的区域和最上的文本框往上的区域都进行映射，这个区域不能选为截取图的右下角顶点。
- 然后在未被映射的区域中随机生成截取图的宽和高（即截取框右下角顶点坐标），如下图所示。
- 最后重新计算截取图中的文本框相对于截取图的坐标，并返回。
3.4、填充
上面进行随机裁剪后，得到的裁剪图大小不一，如果直接进行缩放，那么就会导致严重的失真，所以先对裁剪后的图像进行填充。填充图的大小取裁剪图的宽、高和我们预设的模型输入大小中最大的一个，代码如下，
'''
为了不让原图过度变形，对截取后的图片进行填充
'''
def pad_image(image, polys, input_size):
# DEBUG = True
h, w, _ = image.shape
max_h_w_i = np.max([h, w, input_size])
img_padded = np.zeros((max_h_w_i, max_h_w_i, 3), dtype=np.uint8)
shift_h = np.random.randint(max_h_w_i - h + 1)
shift_w = np.random.randint(max_h_w_i - w + 1)
img_padded[shift_h:h+shift_h, shift_w:w+shift_w, :] = image.copy()
if DEBUG:
cv2.imshow("pad", img_padded)
cv2.waitKey(0)
if len(polys) > 0:
polys[:, :, 0] += shift_w
polys[:, :, 1] += shift_h
return img_padded, polys运行结果，
3.5、缩放成固定大小图片
最后对图片进行缩放，缩放至我们预设的模型输入大小。虽然模型并不会要求输入图像的宽高，但是在训练中，我们还是会指定输入图像的宽高的，这样才能进行批量训练。代码如下，
'''
将图片缩放成固定大小
'''
def resize(image, polys, input_size):
h, w, _ = image.shape
image = cv2.resize(image, dsize=(input_size, input_size))
resize_ratio_x = input_size/float(w)
resize_ratio_y = input_size/float(h)
if len(polys) > 0:
polys[:, :, 0] *= resize_ratio_x
polys[:, :, 1] *= resize_ratio_y
return image, polys