【数据标注】优雅的将文本检测数据集转换为端到端文本提取数据集（上）

背景

任务背景

出于自身需求，我需要将一个有文本框的文本检测中文数据集转换成一个文本提取数据集，并且我希望标签格式为 MMOCR 格式，以便于我使用 MMOCR 提供的工具进行可视化和模型训练。然而直接手动的为每个文本实例打标签工作量非常大，因此想到可以借助 PaddleOCR 训练好的强大文本识别模型来为我的每个框生成伪标签，随后再使用 PPOCRLabel 对数据集进行校验，修改其中错误的标签以得到我想要的文本提取数据集。

一句话总结： 本文旨在提供一种首先使用 PaddleOCR 为矩形文本检测数据集生成识别伪标签，随后使用 PPOCRLabel 进行标签纠错 从而得到一个文本提取数据集 的方法。

如果你有以下需求，本文将可能对你有所帮助：

• 对一个文本数据集进行标注（PPOCRLabel）
• 将一个大的数据集标注任务转换为很多个小的数据集标注任务（本文）
• 将一个文本检测数据集转换为文本提取数据集（PaddleOCR + PPOCRLabel + 本文）
• 将多个小的数据集标注结果合并为一个数据集，并转化为 MMOCR 支持的标签（本文）

需求背景

端到端文本提取（Text Spotting） 技术是近两年研究的热点，现有的公开数据集如 Total Text、ICDAR 2015、CTW1500 等大多为英文数据集，中文数据集较少，对于一些中文文本提取场景，可能仍需要进行数据集标注。

在开源 OCR 方面，PaddleOCR 已经做的很出色了，其开源的 PPOCRLabel 极大的简化了包括水平矩形框、四边形框（可旋转、可以不规则）和任意形状文本框的标注（虽然使用起来可能存在很多的坑）。最值得称赞的是：对于矩形文本数据集（水平和旋转均可），可以使用百度训练好的文本检测与识别模型进行半自动化标注。 PS：自动化标注的原理其实就是用一个百度 Train 好的模型对你的图像进行检测与识别，所以可能有错，还需要你手动校验纠错一遍才能确保正确性，所以叫半自动化。

另一方面，在开源 OCR 方面，另一个顶顶有名的是 MMOCR。MMOCR 提供了一个非常易于使用的模型搭建框架，统一了文本的三个任务（Text Detection、Text Recognition & Text Spotting）的标签格式，可以说 CV 做文本的没用过应该也都知道这个框架。既是不用他搭建模型，用他提供的一些工具来对数据集进行预处理、可视化也是极好的。

因此，可以结合各个工具的优势，来为我们标准文本提取数据集并训练模型提供帮助。

第一步：为文本检测数据集的每个文本框生成识别伪标签

任务目标

前言：本节旨在将矩形框文本检测数据集转换为文本提取数据集，若您的数据集具有一下几种特点，可能并不适用本节内容，可跳至下一节：

• 没有检测框标签（可以直接用 PPOCRLabel 进行全流程的半自动化标注）
• 检测框为弯曲文本（PPOCR 自带模型对弯曲文本识别效果很差）

本文将以本人的数据集情况进行介绍，如有不同，可将自身数据集转换为和我一样的格式，也可自行修改代码逻辑来适应你的数据集。

首先，数据集名称假设成为 textRoot，目录下将包含images_train、label_det_train、images_test和label_det_test，分别是训练集图像、训练集标签、检测集图像和检测集标签。目录结构如下：

- textRoot
  - images_train     # 训练集图像文件夹
    - xxx1.jpg
    - xxx2.jpg
    ...
  - images_test      # 测试集集图像文件夹
    - yyy1.jpg          # 文件名和图像相同，进后缀不同
    - yyy2.jpg
    ...
  - label_det_train  # 训练集标签文件夹
    - xxx1.txt
    - xxx2.txt
    ...
  - label_det_test   # 测试集标签文件夹
    - yyy1.txt
    - yyy2.txt
    ...

其中，每个标签文件如 xxx1.txt 每一行表示一个文本实例，前 8 个数为框的四点坐标（x1, y1,x2,y2,x3,y3,x4,y4），后面固定跟随 0 和 text

260,29,371,29,371,66,260,66,0,text
123,82,540,82,540,109,123,109,0,text
72,114,434,114,434,145,72,145,0,text
241,198,388,198,388,237,241,237,0,text
125,246,547,246,547,286,125,286,0,text
71,285,364,285,364,324,71,324,0,text

本节内容基于 PPOCR，主要原理：借助 PPOCR 提供的端到端文本提取模型（由文本检测模型 + 文本识别模型级联处理组成），通过修改代码，舍弃其中使用检测模型获取文本框的步骤，改为直接读取文本框标签，随后送入到识别模型中获取文本伪标签。该过程复用了 PPOCR 对图像的前中后处理过程，只需进行少量代码修改即可。

环境配置

由于该步骤基于 PPOCR，因此需要配置 Python 环境和 Paddle 环境，本文使用 conda 进行环境管理。

1. 配置 python 环境、paddle 环境

conda create -n paddle python=3.8

conda activate paddle

pip install paddlepaddle  # cpu 版 paddle，文本使用的版本
# pip install paddlepaddle-gpu  # gpu 版 paddle
pip install "paddocr>=2.0.1"

2. 克隆 PaddleOCR 代码到本地

git clone https://github.com/PaddlePaddle/PaddleOCR.git

修改 PPOCR 代码

接下来将修改 PPOCR 的代码，使其直接读取图像的文本框标签，而跳过使用自带的文本检测模型进行文本框的检测。

1. 步骤一：修改 paddleocr.py，在最后一行加一行 main()，否则运行这个脚本的时候不会做任何事情。
2. 步骤二：直接读取检测框标签。修改 tools\infere\predict_system.py 的第 76 行左右，这一行调用了检测模型进行文本框的检测，可以在这里替换成直接读取检测框标签，进行后序的处理。

ori_im = img.copy()
dt_boxes, elapse = self.text_detector(img)  # 这一行
time_dict['det'] = elapse

3. 首先，在 paddleocr.py 647 行附近将图像路径存下来，然后在 669 行附近，将图像路径传进去（用于从图像路径解析对应标签的路径）

...
img_path = img   # 新增这一行
img = check_img(img)
...
for idx, img in enumerate(imgs):
    img = preprocess_image(img)
    dt_boxes, rec_res, _ = self.__call__(img, cls, img_path)  # 修改这一行
    if not dt_boxes and not rec_res:
        ocr_res.append(None)
        continue
...

4. 然后修改 tools\infere\predict_system.py 第 67 行左右，接受图像路径。并在 76 行附近添加判断逻辑，若传入了图像路径，则直接读取对应标签。

...
def __call__(self, img, cls=True, img_path=None):  # 修改这一行
    time_dict = {'det': 0, 'rec': 0, 'cls': 0, 'all': 0}
...

start = time.time()
ori_im = img.copy()
dt_boxes, elapse = self.text_detector(img) # dt_boxes 为四边形框，np.array 类型，n * 4 * 2
time_dict['det'] = elapse    # 检测部分推理时间，约 0.3

# 改为

start = time.time()
ori_im = img.copy()
# 添加了这部分逻辑
if img_path is not None:
    label_path = img_path.replace('images_', 'label_det_').replace('.jpg', '.txt')
    boxes = []
    if os.path.exists(label_path):
        with open(label_path, 'r', encoding='utf-8') as f:
            lines = f.readlines()
            for line in lines:
                line = line.strip().split(',')
                box = line[:8]
                boxes.append(box)
    dt_boxes = np.array(boxes, dtype=np.float32).reshape(-1, 4, 2)
    elapse = 0.308
else:
# 修改结束
    dt_boxes, elapse = self.text_detector(img)
time_dict['det'] = elapse

5. 步骤三：修改 paddleocr.py 约 794 行的代码，此处的 result 变量为一张图的推理的结果，可以在这里修改代码将推理结果保存成文件

# paddleocr.py 部分代码
...
for img_path in image_file_list:  # 将对这个 for 循环进行修改，修改结果见后面
    img_name = os.path.basename(img_path).split('.')[0]
    logger.info('{}{}{}'.format('*' * 10, img_path, '*' * 10))
    if args.type == 'ocr':
        result = engine.ocr(img_path,  # 注意这里
                            det=args.det,
                            rec=args.rec,
                            cls=args.use_angle_cls,
                            bin=args.binarize,
                            inv=args.invert,
                            alpha_color=args.alphacolor)
        if result is not None:
            for idx in range(len(result)):
                res = result[idx]
                for line in res:
                    logger.info(line)
...

# result 结果示例
result = [
    [
        [[[315.0, 18.0], [426.0, 26.0], [423.0, 65.0], [312.0, 57.0]], ('谭天锦', 0.9977920651435852)],
        [[[175.0, 77.0], [599.0, 93.0], [598.0, 120.0], [174.0, 104.0]], ('男，81岁，汉族，泸州市江阳区关工委', 0.9844688773155212)],
        [[[115.0, 109.0], [364.0, 117.0], [363.0, 144.0], [114.0, 136.0]], ('五老志愿者，老干部。', 0.9674175381660461)],
        [[[308.0, 205.0], [417.0, 205.0], [417.0, 241.0], [308.0, 241.0]], ('陈恩薄', 0.9729852676391602)],
        ...
    ],
]

修改结果参考如下（可中断后任务，下次继续原来进度版）

for img_path in image_file_list:
    img_name = os.path.basename(img_path).split('.')[0]
    logger.info('{}{}{}'.format('*' * 10, img_path, '*' * 10))
    # 读取图片所在最后一级目录的名称
    dirName = os.path.dirname(img_path).replace("images_", "label_spotting_")
    labelName = img_name + ".txt"
    if not os.path.exists(dirName):
        os.makedirs(dirName)
    labelPath = os.path.join(dirName, labelName)
    if os.path.exists(labelPath):
        continue
    if args.type == 'ocr':
        result = engine.ocr(img_path,
                            det=args.det,
                            rec=args.rec,
                            cls=args.use_angle_cls,
                            bin=args.binarize,
                            inv=args.invert,
                            alpha_color=args.alphacolor)
        if result is not None:
            for idx in range(len(result)):
                with open(labelPath, 'w', encoding='utf-8') as f:
                    res = result[idx]
                    if res is not None:
                        for lineIdx in range(len(res)):
                            line = res[lineIdx]
                            tmpArray = np.array(line[0]).reshape((-1,)).tolist()
                            f.write(",".join(str(int(i)) for i in tmpArray) + '||||' + line[1][0])
                            if lineIdx != len(res) - 1:
                                f.write('\n')

6. 步骤四：执行推理脚本，--image_dir 为数据集目录参数，--use_angle为是否使用反转检测器，选为 true 即可，--use_gpu 表示是否使用 gpu 进行推理，可以根据自身电脑情况决定，--use_mp 表示是否使用多线程。

python paddleocr.py --image_dir your\path\textRoot\images_test --use_angle_cls true --use_gpu false --use_mp true

7. 生成伪标签将放在 textRoot 目录下的 label_spotting_test 或 label_spotting_train 里：

- textRoot
  - images_train
    - xxx1.jpg
    - xxx2.jpg
    ...
  - label_det_train
    - xxx1.txt
    - xxx2.txt
    ...
  - label_spotting_train    # 生成的文本提取标签文件夹
    - xxx1.txt
    - xxx2.txt
    ...
  ...

标签格式为，前面是原来的文本框的点，中间用 ||||分割，后面是该文本实例生成的伪标签：

260,29,371,29,371,66,260,66||||黄世德
123,82,540,82,540,109,123,109||||男.74岁，汉族，中共党员，雅安市天
72,114,434,114,434,145,72,145||||全县关工委五老志愿者，老干部
241,198,388,198,388,237,241,237||||四龙泽翁
125,246,547,246,547,286,125,286||||男.72岁，藏族，中共党员，甘孜州新
71,285,364,285,364,324,71,324||||龙县五老志愿者，老干部。

文本总结

以上本次分享的上半部分内容，介绍了如何使用 PaddleOCR 为文本检测数据集生成伪标签，后续如何使用 PPOCRLabel 对伪标签进行检查将放在下一篇文章，感兴趣的朋友可以持续关注。点赞、评论可以催更呦！

联系方式

欢迎大家点赞、评论我的文章，你们的支持将是我继续分享的最大动力！也可以联系我获得技术支持，联系方式包括：

• 评论（可能看不见）
• 站内私信（可能看不见）
• 邮箱：cygong@foxmail.com（反馈会更快哦）