其他分享
首页 > 其他分享> > TensorFlow入门教程(26)车牌识别之文本检测模型EAST代码实现(二)

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, polys

3.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轴映射出来,这部分区域都不能选,然后再随机截取其他区域的,如果截取的区域包含了文本框,就再随机截取,直到不包含文本框为止。这种方法比较简单,但是效率比较低。我上面裁剪代码的思路是,

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

上面代码的思路是:

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

 

标签:26,min,polys,image,入门教程,文本框,crop,np,EAST
来源: https://blog.csdn.net/rookie_wei/article/details/112130227