CNOCR重训练

前不久在实验室接手一个项目，与甲方几经周旋后给了一个接口，核心部分是yolo3的文字检测与cnocr的文本识别。在文本识别中，由于给的预训练的模型的训练数据集与项目应用的数据分布差距较大（最明显的是识别字符的范围不同），可能需要对模型重新训练。为应对甲方朋友的一时兴起，特意花了一个下午搞清楚如何重新训练。特撰写此博客以作记录。

一、重训练cnocr的理由

"""
识别范围不同，cnocr提供的模型只能够识别：中文字符，英文字符与若干特殊字符。
但实际应用中，希腊字母也是较为常见的。
经排查，cnocr的模型仅能识别包括：α与β两个希腊字母！

这对于一个理工科男是绝对不能忍的，重训，必须重训！
"""

二、准备工作

下载cnocr工程文件。

git clone https://github.com/breezedeus/cnocr.git
cd cnocr

三、数据集准备

尽管在/cnocr/scripts/cnocr_train.py中提到了两个数据集captcha与cn_ocr，读者们可以从项目的issues中去找相关的链接。下面主要讲我们如何准备自己的数据。

1）新建文件夹datasets，文件夹中新建子目录self_dataset_001，如图，我们将在这里边创建我们自己的数据集。

2）在自己的电脑上找到微软雅黑（msyh）字体，以windows10为例，在“C:\Windows\Fonts”下查找到msyh.ttc，复制过来即可（当然如果想要识别其他印刷体也可以换用其他字体）。

3）在/cnocr/examples下找到label_cn.txt文件，打开可以看到：总共有6000+不重复字符，基本涵盖了所有常见的汉字与字母与特殊字符（偏偏就是没有希腊字母 = =）；复制一份到刚刚新建的datasets下。

4）在self_dataset_001下新准备两个 .txt 文件，分别为：train.txt 和 test.txt ，内容如下：

，其中，两个文件夹的内容是一样的，都是一行一行的文本，格式为：

xxxxxxxxxx.jpg

"""
根据模型本身的训练设置需要，每一行都是 10 个字符，这里小编是为了方便随便找了几段文本拼凑的。
"""

这样的命名是与小编自己写的生成代码有关的。由于小编是一边摸索一边修改的代码，所以过程与代码可能不大优雅，见笑了！

5）执行以下的生成文件：

"""
/cnocr/datasets/data_generator.py
"""

import os
import matplotlib.pyplot as plt 
import numpy as np 
import pygame

pygame.init()
font = pygame.font.Font('msyh.ttc', 64)

def writing(txt, pth):
    r_txt = font.render(txt, True, (0, 0, 0), (255, 255, 255))
    pygame.image.save(r_txt, os.path.join(pth, txt+'.jpg'))


def indexing(txt):
    res = []
    for i in range(len(txt)):
        try:
            res.append(standards.index(txt[i]+'\n')+1)
        except:
            new_chrs.append(txt[i])
            res.append(len(standards)+len(new_chrs)+1)
    return res

if __name__ == '__main__':
    pth = 'self_dataset_001'

    f_tr = open(os.path.join(pth, 'train.txt'), encoding='utf-8-sig')
    f_ts = open(os.path.join(pth, 'test.txt' ), encoding='utf-8-sig')
    f_st = open('label_cn.txt', 'r', encoding='utf-8-sig')

    f_tr2 = open(os.path.join(pth, 'train2.txt'), 'w', encoding='utf-8-sig')
    f_ts2 = open(os.path.join(pth, 'test2.txt' ), 'w', encoding='utf-8-sig')

    train_items = f_tr.readlines()
    test_items  = f_ts.readlines()
    standards   = f_st.readlines()
    new_chrs = []

    f_tr.close()
    f_ts.close()
    f_st.close()

    for i in range(len(train_items)):
        txt = train_items[i].split(".")[0]
        idxes = indexing(txt)
        img = writing(txt, pth)

        cnt = "train_%06d.jpg"%i
        os.rename(os.path.join(pth, txt+'.jpg'), os.path.join(pth, cnt))
        for idx in idxes:
            cnt = cnt + " {}".format(idx)
        f_tr2.write(cnt+'\n')


    for i in range(len(train_items)):
        txt = test_items[i].split(".")[0]
        idxes = indexing(txt)
        img = writing(txt, pth)

        cnt = "test_%06d.jpg"%i
        os.rename(os.path.join(pth, txt+'.jpg'), os.path.join(pth, cnt))
        for idx in idxes:
            cnt = cnt + " {}".format(idx)
        f_ts2.write(cnt+'\n')

    fh = open('label_cn.txt', 'a', encoding='utf-8-sig')
    for nw in new_chrs:
        fh.write(nw+'\n')
    fh.close()

结果是对应 train.txt 和 test.txt 的每一行，渲染了一张图像 “train_xxxxxx.jpg”或者“test_xxxxxx.jpg”的图像，保存在 self_dataset_001 下，

同时在该子目录下还生成了另外两个 .txt 文件：train2.txt 和 test2.txt。形式如下：

即，每一行的格式为：

train_xxxxxx.jpg idx1 idx2 idx3 idx4 idx5 idx6 idx7 idx8 idx9 idx10
"""
11个字段，
第一个对应文件名，方便代码读取对应的文字图像；
后面跟着的10个数字，按顺序对应文字图像里每个字符在 label_cn.txt 文件中的索引（从 1 开始）
"""

值得注意的是，小编在 data_generator.py 文件中增加了对于新字符自动添加至 label_cn.txt 末尾的功能。

四、开始训练

修改 /cnocr/scripts/cnocr_train.py 的 parser 的定义，如下图，即修改数据集相关的三个参数的默认值。

返回到 /cnocr 目录，执行：

python scripts/cnocr_train.py --cpu 2 --num_proc 4 --loss ctc

训练需要花蛮长时间的，完了后应该也是像cnocr提到的替换模型就可以了吧，这一点小编还没尝试，目前只是初步弄懂了如何训练，希望热情的读者如果有成功的经验可以不吝指教哈！