OCR数据集生成利器—TextRecognitionDataGenerator

本文主要介绍通过开源项目 TextRecognitionDataGenerator 生成定制OCR数据集，用于训练模型和测试模型效果，达到特定场景下识别的优化。

1、现有数据集的问题

在自然场景（如照片，视频）的文字OCR识别中，由于文字的字体和颜色不同，长短和布局不定，同时文字所处背景的复杂性，模型训练后在特定场景下的文字识别效果依赖于数据集是否覆盖类似风格的文字（训练-测试同分布假设）。

当前虽然有不少开放的大规模OCR数据集，比如牛津大学的 Synthetic Word Dataset 和 SynthText in the Wild Dataset，作为业界的标准数据集，两者大小均在100g以上，但是覆盖的场景仍然非常有限，比如训练集中不包含标点符号，很多字体没有覆盖，另外还有各大竞赛的数据集 IC13, IC15, IIIT, SVT. IIIT, SVT, IC03, IC13, IC15, SVTP 和 CUTE，也多是特定场景下的数据集。

在实际测试中发现，即使在这些大规模数据集（如 Synthetic Word Dataset ）达到非常高的测试精度，在特定场景下的识别精度仍然不理想（比如身份证文字的识别），原因可能有两点，当前的 crnn+ctc 或者 attention+ctc（目前识别模型就这两种）难以学习到较好泛化性能的权值，二是尽管数据集足够大，但是由于自然场景中文字的多样性，数据集很难覆盖到特定场景下的文字风格。

所以从数据集的角度考虑问题，针对特定场景生成相应的数据集是解决特定场景下识别问题的一个思路。

2、数据集生成方法

从训练集和测试集同分布假设考虑，最好的数据集当然是从实际的样本标注而来。但是文字的标注比较耗时，ocr数据集通常至少需要包含几十万张图片，所以实际样本的标注是一个大工程。

本文主要介绍另外一种方法：合成法

所谓合成法，是通过已有的单词或者字符进行随机组合，然后通过模糊、倾斜、扭曲、仿射变换和加背景等方法，生成接近真实场景中的图片。由于单词和已有的字符已知，所以省去了标注的过程。

实际上工业上普遍采用的Synthetic Word Dataset 和 SynthText in the Wild Dataset就是合成而来，下图是两个数据集的合成方式，实践证明，用这两个数据集训练的模型在街景文字识别中有较好的性能。

本文介绍一个开源OCR数据集生成项目TextRecognitionDataGenerator ，可以实现类似的合成效果，但是可以根据具体的场景定制自己的数据集。

3、开源OCR数据集生成项目TextRecognitionDataGenerator 

该项目通过 Python实现，可以通过 pip 安装：

终端： pip install trdg

然后在终端中输入以下命令：

终端：trdg -c 1000 -w 5

就可以生成如下图片，其中 -c 参数表示生成图片的数量， -w 表示图片中单词的个数。

但是用 pip 安装和 shell 命令生成的方式不够灵活，本文主要介绍通过源码的方式生成数据集。

 终端：git clone git@github.com:Belval/TextRecognitionDataGenerator.git

进入项目后，修改 run.py 文件中的默认参数即可，该项目的强大之处在于作者提供了非常多的定制参数，可以根据自己的需求来生成数据集。

（1）参数说明

--output_dir

-i/--input_file

-l/--language

-c/--count

-rs/--random_sequences

-let/--include_letters

-num/--include_numbers

-sym/--include_symbols

-w/--length

-r/--random

-f/--format

-t/--thread_count

-e/--extension

-k/--skew_angle

-rk/--random_skew

-wk/--use_wikipedia

-bl/--blur

-rbl/--random_blur

-b/--background

-id/--image_dir

-hw/--handwritten

-na/--name_format

-om/--output_mask

-d/--distorsion

-do/--distorsion_orientation

-wd/--width

-al/--alignment

-or/--orientation

-tc/--text_color

-sw/--space_width

-cs/--character_spacing

-m/--margins

-fi/--fit

-ft/--font

-fd/--font_dir

-ca/--case

-dt/--dict

-ws/--word_split

对应的代码如下：

def parse_arguments():
    """
        Parse the command line arguments of the program.
    """
 
    parser = argparse.ArgumentParser(
        description="Generate synthetic text data for text recognition."
    )
    parser.add_argument(
        "--output_dir", type=str, nargs="?", help="The output directory", default="out/"
    )
    parser.add_argument(
        "-i",
        "--input_file",
        type=str,
        nargs="?",
        help="When set, this argument uses a specified text file as source for the text",
        default="",
    )
    parser.add_argument(
        "-l",
        "--language",
        type=str,
        nargs="?",
        help="The language to use, should be fr (French), en (English), es (Spanish), de (German), ar (Arabic), cn (Chinese), or hi (Hindi)",
        default="en"
    )
    parser.add_argument(
        "-c",
        "--count",
        type=int,
        nargs="?",
        help="The number of images to be created.",
        default=10
    )
    parser.add_argument(
        "-rs",
        "--random_sequences",
        action="store_true",
        help="Use random sequences as the source text for the generation. Set '-let','-num','-sym' to use letters/numbers/symbols. If none specified, using all three.",
        default=True,
    )
    parser.add_argument(
        "-let",
        "--include_letters",
        action="store_true",
        help="Define if random sequences should contain letters. Only works with -rs",
        default=True,
    )
    parser.add_argument(
        "-num",
        "--include_numbers",
        action="store_true",
        help="Define if random sequences should contain numbers. Only works with -rs",
        default=True,
    )
    parser.add_argument(
        "-sym",
        "--include_symbols",
        action="store_true",
        help="Define if random sequences should contain symbols. Only works with -rs",
        default=False,
    )
    parser.add_argument(
        "-w",
        "--length",
        type=int,
        nargs="?",
        help="Define how many words should be included in each generated sample. If the text source is Wikipedia, this is the MINIMUM length",
        default=2,
    )
    parser.add_argument(
        "-r",
        "--random",
        action="store_true",
        help="Define if the produced string will have variable word count (with --length being the maximum)",
        default=False,
    )
    parser.add_argument(
        "-f",
        "--format",
        type=int,
        nargs="?",
        help="Define the height of the produced images if horizontal, else the width",
        default=32,
    )
    parser.add_argument(
        "-t",
        "--thread_count",
        type=int,
        nargs="?",
        help="Define the number of thread to use for image generation",
        default=8,
    )
    parser.add_argument(
        "-e",
        "--extension",
        type=str,
        nargs="?",
        help="Define the extension to save the image with",
        default="jpg",
    )
    parser.add_argument(
        "-k",
        "--skew_angle",
        type=int,
        nargs="?",
        help="Define skewing angle of the generated text. In positive degrees",
        default=0,
    )
    parser.add_argument(
        "-rk",
        "--random_skew",
        action="store_true",
        help="When set, the skew angle will be randomized between the value set with -k and it's opposite",
        default=False,
    )
    parser.add_argument(
        "-wk",
        "--use_wikipedia",
        action="store_true",
        help="Use Wikipedia as the source text for the generation, using this paremeter ignores -r, -n, -s",
        default=False,
    )
    parser.add_argument(
        "-bl",
        "--blur",
        type=int,
        nargs="?",
        help="Apply gaussian blur to the resulting sample. Should be an integer defining the blur radius",
        default=0,
    )
    parser.add_argument(
        "-rbl",
        "--random_blur",
        action="store_true",
        help="When set, the blur radius will be randomized between 0 and -bl.",
        default=False,
    )
    parser.add_argument(
        "-b",
        "--background",
        type=int,
        nargs="?",
        help="Define what kind of background to use. 0: Gaussian Noise, 1: Plain white, 2: Quasicrystal, 3: Image",
        default=1,
    )
    parser.add_argument(
        "-hw",
        "--handwritten",
        action="store_true",
        help='Define if the data will be "handwritten" by an RNN',
    )
    parser.add_argument(
        "-na",
        "--name_format",
        type=int,
        help="Define how the produced files will be named. 0: [TEXT]_[ID].[EXT], 1: [ID]_[TEXT].[EXT] 2: [ID].[EXT] + one file labels.txt containing id-to-label mappings",
        default=0,
    )
    parser.add_argument(
        "-om",
        "--output_mask",
        type=int,
        help="Define if the generator will return masks for the text",
        default=0,
    )
    parser.add_argument(
        "-d",
        "--distorsion",
        type=int,
        nargs="?",
        help="Define a distorsion applied to the resulting image. 0: None (Default), 1: Sine wave, 2: Cosine wave, 3: Random",
        default=0,
    )
    parser.add_argument(
        "-do",
        "--distorsion_orientation",
        type=int,
        nargs="?",
        help="Define the distorsion's orientation. Only used if -d is specified. 0: Vertical (Up and down), 1: Horizontal (Left and Right), 2: Both",
        default=0,
    )
    parser.add_argument(
        "-wd",
        "--width",
        type=int,
        nargs="?",
        help="Define the width of the resulting image. If not set it will be the width of the text + 10. If the width of the generated text is bigger that number will be used",
        default=-1,
    )
    parser.add_argument(
        "-al",
        "--alignment",
        type=int,
        nargs="?",
        help="Define the alignment of the text in the image. Only used if the width parameter is set. 0: left, 1: center, 2: right",
        default=1,
    )
    parser.add_argument(
        "-or",
        "--orientation",
        type=int,
        nargs="?",
        help="Define the orientation of the text. 0: Horizontal, 1: Vertical",
        default=0,
    )
    parser.add_argument(
        "-tc",
        "--text_color",
        type=str,
        nargs="?",
        help="Define the text's color, should be either a single hex color or a range in the ?,? format.",
        default="#282828",
    )
    parser.add_argument(
        "-sw",
        "--space_width",
        type=float,
        nargs="?",
        help="Define the width of the spaces between words. 2.0 means twice the normal space width",
        default=1.0,
    )
    parser.add_argument(
        "-cs",
        "--character_spacing",
        type=int,
        nargs="?",
        help="Define the width of the spaces between characters. 2 means two pixels",
        default=0,
    )
    parser.add_argument(
        "-m",
        "--margins",
        type=margins,
        nargs="?",
        help="Define the margins around the text when rendered. In pixels",
        default=(0, 0, 0, 0),
    )
    parser.add_argument(
        "-fi",
        "--fit",
        action="store_true",
        help="Apply a tight crop around the rendered text",
        default=False,
    )
    parser.add_argument(
        "-ft", "--font", type=str, nargs="?", help="Define font to be used"
    )
    parser.add_argument(
        "-fd",
        "--font_dir",
        type=str,
        nargs="?",
        help="Define a font directory to be used",
    )
    parser.add_argument(
        "-id",
        "--image_dir",
        type=str,
        nargs="?",
        help="Define an image directory to use when background is set to image",
        default=os.path.join(os.path.split(os.path.realpath(__file__))[0], "images")
    )
    parser.add_argument(
        "-ca",
        "--case",
        type=str,
        nargs="?",
        help="Generate upper or lowercase only. arguments: upper or lower. Example: --case upper",
        default="upper"
    )
    parser.add_argument(
        "-dt", "--dict", type=str, nargs="?", help="Define the dictionary to be used"
    )
    parser.add_argument(
        "-ws", "--word_split",
        action="store_true",
        help="Split on words instead of on characters (preserves ligatures, no character spacing)",
        default=False,
    )
    return parser.parse_args()

（2）实例应用

对于卡证类的识别，由于卡证的背景和字体都是固定的，所以可以通过指定字体和背景生成数据集。

以一个卡片的场景为例，

PAN 卡中包含两种字体：Mangal Bold 和 time new roman，所以可以下载对应字体的 ttf  文件，放到项目的font文件夹中，然后移除其它的字体文件，让程序仅根据这两种字体生成图片。

同时识别的字段中包含数字，字母和符号，同时人名没有明显的规律，所以训练集最好通过字符来生成。

所以可以设定如下参数：

1) 指定生成文本的单词由字母、数字和符号组合生成

--random_sequences: True

--include_letters: True

--include_numbers: True

--include_symbols: True

2) 设置一个图片的单词数在1~3之间（考虑到卡片中的字段不会超过3个词）

--length：3

--random：True

3) 照片有的时候由于拍照原因会倾斜，所以设定倾斜角度在（-10度～10度之间）

--skew_angle：10

--random_skew：True

4) 照片有的时候会拍得很模块，所以设置高斯模糊（0~0.1）

--blur:0

--random_blur:True

5) 背景基本上为纯色，采用高斯噪声作为背景

--background：0

其余的再设置一下字符间隔，图片的上下间隔和单词之间的间隔，得到的数据集如下所示：

  从字体和排版上以及比较接近卡片上的文字了。作为改进，可以随机截掉文字的上下一小部分，这个在现实中也经常出现（文字检测不够准确），替换更加相似的背景等等。