前不久在實驗室接手一個項目,與甲方几經周旋後給了一個接口,核心部分是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提到的替換模型就可以了吧,這一點小編還沒嘗試,目前只是初步弄懂了如何訓練,希望熱情的讀者如果有成功的經驗可以不吝指教哈!