用Transformer实现OCR字符识别！

最新推荐文章于 2024-09-02 08:35:43 发布

Datawhale

最新推荐文章于 2024-09-02 08:35:43 发布

阅读量4.3k

点赞数 8

文章标签：大数据数据挖掘编程语言 python 计算机视觉

本文链接：https://blog.csdn.net/Datawhale/article/details/121646034

版权

本文通过ICDAR2015数据集，详细讲解如何利用Transformer实现OCR字符识别任务，包括数据集分析、transformer模型构建和训练，以及贪心解码策略。文章提供了数据集获取、数据预处理、模型训练等关键步骤的代码实现，帮助读者理解如何将Transformer应用到计算机视觉任务中。

摘要由CSDN通过智能技术生成

Datawhale干货

作者：安晟、袁明坤，Datawhale成员

在CV领域中，transformer除了分类还能做什么？本文将采用一个单词识别任务数据集，讲解如何使用transformer实现一个简单的OCR文字识别任务，并从中体会transformer是如何应用到除分类以外更复杂的CV任务中的。全文分为四部分：

一、数据集简介与获取

二、数据分析与关系构建

三、如何将transformer引入OCR

四、训练框架代码讲解

注：本文围绕如何设计模型和训练架构来解决OCR任务，文章含完整实践，代码很长建议收藏。不熟悉transformer的小伙伴可以点击这里回顾。

整个文字识别任务中，主要包括以下几个文件：

- analysis_recognition_dataset.py (数据集分析脚本)

- ocr_by_transformer.py (OCR任务训练脚本)

- transformer.py (transformer模型文件)

- train_utils.py (训练相关辅助函数，loss、optimizer等)

其中 ocr_by_transformer.py 为主要的训练脚本，其依托 train_utils.py 和 transformer.py 两个文件构建 transformer 来完成字符识别模型的训练

一、数据集简介与获取

本文使用的数据集基于ICDAR2015 Incidental Scene Text 中的 Task 4.3: Word Recognition，这是一个著名的自然场景下文本识别数据集，本次用来进行单词识别任务，我们去掉了其中一些图片，来简化这个实验的难度，因此本文的数据集与原始数据集略有差别。

为了能够更好的进行数据共享和版本管控，我们选择在线调用数据集，将简化后的数据集存放在专门的数据共享平台，数据开源地址： https://gas.graviti.cn/dataset/datawhale/ICDAR2015 ，有相关问题可以直接在数据集讨论区交流。

该数据集包含了众多自然场景图像中出现的文字区域，数据中训练集含有4326张图像，测试集含有1992张图像，他们都是从原始大图中依据文字区域的bounding box裁剪出来的，图像中的文字基本处于图片中心位置。

数据集中图像类似如下样式：

word_104.png, "Optical"

数据本身以图像展示，对应的标签存储在CLASSIFICATION中，后文代码中标签获取，将直接得到一个包含所有字符的列表，这也是为了方便标签易用性进行的存储选择。

下面简单介绍数据集的快速使用：

本地下载安装 tensorbay

pip3 install tensorbay

打开本文数据集链接：https://gas.graviti.cn/dataset/datawhale/ICDAR2015
将数据集fork到自己账户下
点击网页上方开发者工具 --> AccessKey --> 新建一个AccessKey --> 复制这个Key

from tensorbay import GAS
from tensorbay.dataset import Dataset

# GAS凭证
KEY = 'Accesskey-***************80a'  # 添加自己的AccessKey
gas = GAS(KEY)

# 获取数据集
dataset = Dataset("ICDAR2015", gas)
# dataset.enable_cache('./data')  # 开启本语句，选择将数据建立本地缓存

# 训练集和验证集
train_segment = dataset["train"]valid_segment = dataset['valid']

# 数据及标签
for data in train_segment:
    # 图像数据
    img = data.open()
    # 图像标签
    label = data.label.classification.category
    break

通过以上代码获取的图像及标签形式如下：

img	label
	['C', 'A', 'U', 'T', 'I', 'O', 'N']

通过以上简单代码便可以快速获取图像数据及标签，但程序每次运行都会自动去平台下载数据，因此耗时较长，建议开启本地缓存，一次下载多次使用，当不再使用数据时便可以将数据删除。

二、数据分析与关系构建

开始实验前，我们先对数据进行简单分析，只有对数据的特性足够了解，才能够更好的搭建出baseline，在训练中少走弯路。

运行下面代码，即可一键完成对于数据集的简单分析：

python analysis_recognition_dataset.py

具体地，这个脚本所做的工作包括：对数据进行标签字符统计(有哪些字符、每个字符出现次数多少)、最长标签长度统计，图像尺寸分析等，并且构建字符标签的映射关系文件 lbl2id_map.txt。

下面我们来一点点看代码：

注：本文代码开源地址：

https://github.com/datawhalechina/dive-into-cv-pytorch/tree/master/code/chapter06_transformer/6.2_recognition_by_transformer(online_dataset)

首先完成准备工作，导入需要的库，并设置好相关目录或文件的路径

import os
from PIL import Image
import tqdm

from tensorbay import GAS
from tensorbay.dataset import Dataset

# GAS凭证
KEY = 'Accesskey-************************480a'  # 添加自己的AccessKey
gas = GAS(KEY)
# 数据集获取并本地缓存
dataset = Dataset("ICDAR2015", gas)
dataset.enable_cache('./data')  # 数据缓存地址

# 获取训练集和验证集
train_segment = dataset["train"]
valid_segment = dataset['valid']

# 中间文件存储路径，存储标签字符与其id的映射关系
base_data_dir = './'
lbl2id_map_path = os.path.join(base_data_dir, 'lbl2id_map.txt')

2.1 标签最长字符个数统计

首先统计数据集最长label中包含的字符数量，此处要将训练集和验证集中的最长标签都进行统计，进而得到最长标签所含字符。

def statistics_max_len_label(segment):
    """
    统计标签中最长的label所包含的字符数
    """
    max_len = -1
    for data in segment:
        lbl_str = data.label.classification.category  # 获取标签
        lbl_len = len(lbl_str)
        max_len = max_len if max_len > lbl_len else lbl_len
    return max_len

train_max_label_len = statistics_max_len_label(train_segment)  # 训练集最长label
valid_max_label_len = statistics_max_len_label(valid_segment)  # 验证集最长label
max_label_len = max(train_max_label_len, valid_max_label_len)  # 全数据集最长label
print(f"数据集中包含字符最多的label长度为{max_label_len}")

数据集中最长label含有21个字符，这将为后面transformer模型搭建时的时间步长度的设置提供参考。

2.2 标签所含字符统计

下面代码查看数据集中出现过的所有字符：

def statistics_label_cnt(segment, lbl_cnt_map):
    """
    统计标签文件中label都包含哪些字符以及各自出现的次数
    lbl_cnt_map : 记录标签中字符出现次数的字典
    """
    for data in segment:
        lbl_str = data.label.classification.category  # 获取标签
        for lbl in lbl_str:
                if lbl not in lbl_cnt_map.keys():
                    lbl_cnt_map[lbl] = 1
                else:
                    lbl_cnt_map[lbl] += 1

lbl_cnt_map = dict()  # 用于存储字符出现次数的字典
statistics_label_cnt(train_segment, lbl_cnt_map)  # 训练集中字符出现次数统计
print("训练集中label中出现的字符:")
print(lbl_cnt_map)
statistics_label_cnt(valid_segment, lbl_cnt_map)  # 训练集和验证集中字符出现次数统计
print("训练集+验证集label中出现的字符:")
print(lbl_cnt_map)

输出结果为：

训练集中label中出现的字符:
{'C': 593, 'A': 1189, 'U': 319, 'T': 896, 'I': 861, 'O': 965, 'N': 785, 'D': 383, 'W': 179, 'M': 367, 'E': 1423, 'X': 110, '$': 46, '2': 121, '4': 44, 'L': 745, 'F': 259, 'P': 389, 'R': 836, 'S': 1164, 'a': 843, 'v': 123, 'e': 1057, 'G': 345, "'": 51, 'r': 655, 'k': 96, 's': 557, 'i': 651, 'c': 318, 'V': 158, 'H': 391, '3': 50, '.': 95, '"': 8, '-': 68, ',': 19, 'Y': 229, 't': 563, 'y': 161, 'B': 332, 'u': 293, 'x': 27, 'n': 605, 'g': 171, 'o': 659, 'l': 408, 'd': 258, 'b': 88, 'p': 197, 'K': 163, 'J': 72, '5': 80, '0': 203, '1': 186, 'h': 299, '!': 51, ':': 19, 'f': 133, 'm': 202, '9': 66, '7': 45, 'j': 15, 'z': 12, '´': 3, 'Q': 19, 'Z': 29, '&': 9, ' ': 50, '8': 47, '/': 24, '#': 16, 'w': 97, '?': 5, '6': 40, '[': 2, ']': 2, 'É': 1, 'q': 3, ';': 3, '@': 4, '%': 28, '=': 1, '(': 6, ')': 5, '+': 1}
训练集+验证集label中出现的字符:
{'C': 893, 'A': 1827, 'U': 467, 'T': 1315, 'I': 1241, 'O': 1440, 'N': 1158, 'D': 548, 'W': 288, 'M': 536, 'E': 2215, 'X': 181, '$': 57, '2': 141, '4': 53, 'L': 1120, 'F': 402, 'P': 582, 'R': 1262, 'S': 1752,