【OCR技術】大批量生成文字訓練集

七月線上實驗室發表於2019-03-12

640?wx_fmt=png

來源 | 部落格園

作者 | Madcola

如果是想訓練一個手寫體識別的模型，用一些前人收集好的手寫文字集就好了，比如中科院的這些資料集。

http://www.nlpr.ia.ac.cn/databases/handwriting/Offline_database.html

但是如果我們只是想要訓練一個專門用於識別印刷漢字的模型，那麼我們就需要各種印刷字型的訓練集，那怎麼獲取呢？藉助強大的影象庫，自己生成就行了！

GitHub地址：

https://github.com/AstarLight/CPS-OCR-Engine/blob/master/ocr/gen_printed_char.py

先捋一捋思路，生成文字集需要什麼步驟：

確定你要生成多少字型，生成一個記錄著漢字與label的對應表。
確定和收集需要用到的字型檔案。
生成字型影象，儲存在規定的目錄下。
適當的資料增強。

第三步的生成字型影象最為重要，如果僅僅是生成很正規的文字，那麼用這個正規文字集去訓練模型，第一影象數目有點少，第二模型泛化能力比較差，所以我們需要對字型影象做大量的影象處理工作，以增大我們的印刷體文字資料集。

我總結了一下，我們可以做的一些影象增強工作有這些：

文字扭曲
背景噪聲（椒鹽）
文字位置（設定文字的中心點）
筆畫粘連（膨脹來模擬）
筆畫斷裂（腐蝕來模擬）
文字傾斜（文字旋轉）
多種字型

做完以上增強後，我們得到的資料集已經非常龐大了。

現在開始一步一步生成我們的3755個漢字的印刷體文字資料集。

一、生成漢字與label的對應表

這裡的漢字、label對映表的生成我使用了pickel模組，藉助它生成一個id:漢字的對映檔案儲存下來。
這裡舉個小例子說明怎麼生成這個“漢字：id”對映表。

首先在一個txt檔案裡寫入你想要的漢字，如果對漢字對應的ID沒有要求的話，我們不妨使用該漢字的排位作為其ID，比如“一二三四五”中，五的ID就是00005。如此類推，把漢字讀入記憶體，建立一個字典，把這個關係記錄下來，再使用pickle.dump存入檔案儲存。

二、收集字型檔案

字型檔案上網收集就好了，但是值得注意的是，不是每一種字型都支援漢字，所以我們需要篩選出真正適合漢字生成的字型檔案才可以。我一共使用了十三種漢字字型作為我們接下來漢字資料集用到的字型，具體如下

圖：

640?wx_fmt=png

當然，如果需要進一步擴大資料集來增強訓練得到的模型的泛化能力，可以花更多的時間去收集各類漢字字型，那麼模型在面對各種字型時也能從容應對，給出準確的預測。

三、文字影象生成

首先是定義好輸入引數，其中包括輸出目錄、字型目錄、測試集大小、影象尺寸、影象旋轉幅度等等。

640?wx_fmt=png

接下來需要將我們第一步得到的對應表讀入記憶體，因為這個表示ID到漢字的對映，我們在做一下轉換，改成漢字到ID的對映，用於後面的字型生成。

640?wx_fmt=png

我們對旋轉的角度儲存到列表中，旋轉角度的範圍是[-rotate,rotate].

640?wx_fmt=png

現在說一下字型影象是怎麼生成的，首先我們使用的工具是PIL。PIL裡面有很好用的漢字生成函式，我們用這個函式再結合我們提供的字型檔案，就可以生成我們想要的數字化的漢字了。我們先設定好我們生成的字型顏色為黑底白色，字型尺寸由輸入引數來動態設定。

640?wx_fmt=png

我們寫兩個迴圈，外層迴圈是漢字列表，內層迴圈是字型列表，對於每個漢字會得到一個image_list列表，裡面儲存著這個漢字的所有影象。

640?wx_fmt=png

我們將image_list中影象按照比例分為訓練集和測試集儲存。

640?wx_fmt=png

寫好程式碼後，我們執行如下指令，開始生成印刷體文字漢字集。

640?wx_fmt=png

解析一下上述指令的附屬引數：

--out_dir 表示生成的漢字影象的儲存目錄
--font_dir 表示放置漢字字型檔案的路徑
--width --height 表示生成影象的高度和寬度
--margin 表示字型與邊緣的間隔
--rotate 表示字型旋轉的範圍，[-rotate,rotate]
--rotate_step 表示每次旋轉的間隔

生成這麼一個3755個漢字的資料集的所需的時間還是很久的，估計接近一個小時。其實這個生成過程可以用多執行緒、多程式並行加速，但是考慮到這種文字資料集只需生成一次就好，所以就沒做這方面的優化了。資料集生成完我們可以發現，在dataset資料夾下得到train和test兩個資料夾，train和test資料夾下都有3755個子資料夾，分別儲存著生成的3755個漢字對應的影象，每個子檔案的名字就是該漢字對應的id。隨便選擇一個train資料夾下的一個子資料夾開啟，可以看到所獲得的漢字影象，一共634個。

dataset下自動生成測試集和訓練集

640?wx_fmt=png