iT邦幫忙

2021 iThome 鐵人賽

DAY 3
1
AI & Data

手寫中文字之影像辨識系列 第 3

【第3天】資料前處理-YOLOv4與自動框選中文字

現況

  1. 觀察主辦單位提供的資料集(約7萬張圖檔),發現圖檔大致分為下列幾種。
    1.1 圖檔內只有1個中文字

    1.2 圖檔內中文字有其他字跡或只有半個字

    1.3 圖檔內含2個以上中文字

    1.4 圖檔內無文字(僅有空白背景)

  2. 為了分類上述圖檔,決定採用one-stage物件偵測演算法-YOLOv4,小量抽樣建模後,自動框選中文字。

  3. 此次抽樣建模的資料集較小(3000張),適合部署到Colab上以免費GPU訓練。若資料集過大,不建議以Colab訓練模型,曾經遇到VMware硬碟空間不足、上傳資料時間過久、模型訓練到一半Colab斷線或提示GPU用量額滿的慘痛經驗。


工具/套件

  1. LabelImg
  2. YOLOv4
  3. Google Colab

內容

  1. LabelImg

    1.1 Windows環境-安裝與執行

    • <方法一> 以Python執行

      ※註:詳細步驟或其他環境的安裝,可參考官方文件

    • <方法二> 以EXE檔執行
      官方I/O下載解壓縮後,即可執行。

    1.2 操作過程

    • 執行LabelImg後,分別點選Open Dir與Change Save Dir,設定訓練集圖檔路徑及標記Annotations(XML檔)的儲存路徑。

    • 勾選選單中View的Auto Save

    • 框選中文字體(Bounding box),並標記為word (快捷鍵:W拉框框、D下一張、A上一張)

    • 每標記完一張圖檔 都會產出一個對應的XML檔案(切記!檔名盡量不要使用中文字)

  2. YOLOv4部署到Colab訓練

    2.1 在Colab上訓練YOLOv4,可以避免Windows+YOlOv4,在呼叫本地端GPU時出現異常。若有Windows+YOlOv4的穩定訓練方法,也歡迎大家留言告訴我。

    2.2 事前準備

    • 按此下載 train.rar,並解壓縮成train資料夾。

    • 將資料集(3000張)放進train下的\VOCdevkit\VOC2021\JPEGImages

    • 將LabelImg標記產生的XML檔(3000個),放進train底下\VOCdevkit\VOC2021\Annotations

    • 執行train資料夾中的gen_train_val.py,將資料集以8:2比例分配train與val。

    import os, random
    from os.path import join, splitext
    
    base = 'VOC2021'
    source = join('VOCdevkit',base, 'JPEGImages/')
    train_txt = join('VOCdevkit',base, 'ImageSets/Main/train.txt')
    val_txt = join('VOCdevkit',base, 'ImageSets/Main/val.txt')
    
    files = os.listdir(source)
    random.shuffle(files)
    
    f_train = open(train_txt,'a')
    f_val = open(val_txt,'a')
    
    for i, file in enumerate(files):
        name = splitext(file)[0]
        if (i >=len(files)*0.2):
            f_train.write(name+'\n')
        else:
            f_val.write(name+'\n')
    
    f_train.close()
    f_val.close()
    
    • 執行gen_train_val.py後,得到train.txt與val.txt

    • train資料夾中voc_label.py,用於YOLOv4訓練模型前的datasets預處理(標記Train / Test/Val資料集),將在部署到Colab之後執行。

      ※註:如果想增加YOLOv4辨識的標籤種類(classes),需修改train資料夾中檔案的參數設定,請參閱下圖。

    • 接下來是將整個train資料夾壓縮成train3.zip

    2.3 部署到Colab訓練:最後一步是將train3.zip上傳到Colab訓練,完成模型訓練後,壓縮成train3_finished.tar.gz,再下載到本地端。考量到篇幅過長,詳細過程請參閱我的Github

  3. 模型成效
    3.1 將train3_finished.tar.gz解壓縮後,train資料夾中predit.py,執行後可自動框選中文字。

    import cv2
    import numpy as np
    import os
    
    #讀取模型與訓練權重
    def initNet():
        CONFIG = 'yolov4-tiny-myobj.cfg'
        WEIGHT = 'yolov4-tiny-myobj_last.weights'
        net = cv2.dnn.readNet(CONFIG,WEIGHT)
        model = cv2.dnn_DetectionModel(net)
        model.setInputParams(size=(416, 416), scale=1/255.0)
        model.setInputSwapRB(True)
        return model
    
    #物件偵測
    def nnProcess(image, model):
        classes, confs, boxes = model.detect(image, 0.4, 0.1)
        return classes, confs, boxes
    
    #框選偵測到的物件
    def drawBox(image, classes, confs, boxes):
        new_image = image.copy()
        for (classid, conf, box) in zip(classes, confs,boxes):
         x, y, w, h = box
            if x - 2 < 0:
                x = 2
            if y - 2 < 0:
                y = 2
            cv2.rectangle(new_image, (x - 2, y - 2), (x + w + 2, y + h + 2), (0, 255, 0), 2)
        return new_image
    
    if __name__ == '__main__':
        # 主辦單位提供的資料集(約7萬張)
        source = './01_origin/'
        files = os.listdir(source)
        # 依照正整數排序
        files.sort(key=lambda x:int(x[:-6]))
        model = initNet()
        for file in files:
            img = cv2.imdecode(np.fromfile(source+file,dtype=np.uint8), -1)
            classes, confs, boxes = nnProcess(img, model)
            try:
                frame = drawBox(img, classes, confs, boxes)
                frame = cv2.resize(frame, (240, 200), interpolation=cv2.INTER_CUBIC)
                # 顯示框選後的圖片
                cv2.imshow('img', frame)
                cv2.waitKey()
            except:
                continue
        print('程式執行完畢')
    

    3.2 框選中文字效果

    • 圖檔內只有1個中文字

    • 圖檔內中文字有其他字跡或只有半個字

    • 圖檔內含2個以上中文字

    • 圖檔內無文字(僅有空白背景)


小結

  1. 從3.2的結果得知,訓練後的YOLOv4模型,可以正確地框選出中文字。唯一美中不足的地方,大概是「體」被模型當成2個字框選。

  2. 承上,推測可能是小量抽樣(3000張)中,「體」這種情況的樣本出現頻率較低。增加抽樣數量、重新訓練模型後,應該可改善框選表現。

  3. 考量到:「正式比賽時,每張圖檔內只會有一個最正確的中文字」。因此,下一章的目標是:「從7萬張圖檔中,取出YOLOv4模型框選為1個中文字的圖檔,再將物件偵測框(綠色邊界框)內的中文字裁切、另存圖檔,做為新的資料集」。

讓我們繼續看下去...


參考資料

  1. LabelImg詳細說明-tzutalin
  2. YOLOv4詳細說明-AlexeyAB

上一篇
【第2天】開賽準備
下一篇
【第4天】資料前處理-圖檔分類與裁切
系列文
手寫中文字之影像辨識31

尚未有邦友留言

立即登入留言