[AI-Challenge 2021] Nhận dạng chữ tiếng Việt trong ảnh ngoại cảnh

Bài viết này mình sẽ giới thiệu về cách triển khai model detect (phát hiện) và recogize (nhận diện) chữ trong ảnh ngoại cảnh mà mình sử dụng trong cuộc thi AI-Challenge 2021. Trong cuộc thi này team mình sử dụng framework PaddleOCR, mình sử dụng model SAST để predict ra những bounding box chứa text và đưa qua model SRN để predict text ở trong bouding box, ở framework này documents của nó hướng dẫn rất đầy đủ.

Mục lục

1. Giới thiệu về cuộc thi AIC
2. Bài toán nhận diện chữ tiếng Việt trong ảnh ngoại cảnh và sinh hoạt hằng ngày
3. Các bước triển khai mô hình
4. Model Detection (SAST)
5. Model Recognition (SRN)
6. Evaluation
7. Inference với model đã train xong
8. Tạo docker image
9. Kết luận
References

1. Giới thiệu về cuộc thi AIC

AI-Challenge là một cuộc thi về trí tuệ nhân tạo được tổ chức hàng năm. Trong đó chú trọng về học tập và ứng dụng trí tuệ nhân tạo vào chương trình xây dựng thành phố Hồ Chí Minh trở thành đô thị thông minh giải, quyết các bài toán quan trọng phục vụ cuộc sống.

2. Bài toán nhận diện chữ tiếng Việt trong ảnh ngoại cảnh và sinh hoạt hằng ngày

Bài toán trong cuộc thi AI năm 2021 là “Nhận diện chữ tiếng Việt trong ảnh ngoại cảnh và sinh hoạt hằng ngày”. Mục tiêu là phát hiện (detect) và nhận diện (recognize) chữ trong ảnh, cụ thể ở đây sẽ tập trung vào chữ trong khung cảnh (scene text) được thu lại từ nhiều nguồn camera khác nhau ở Việt Nam.

Vấn đề đặt ra là một phần quan trọng cho nhiều hệ thống thông minh hiện nay như robot, xe tự lái… Các hệ thống này yêu cầu hiểu được cảnh vật xung quanh và chữ trong cảnh vật nắm giữ rất nhiều thông tin quan trọng, có thể phục vụ phát triển du lịch thông minh, bảo tàng thông minh, xe tự lái, robot tự hành… Để tăng tính ứng dụng thực tiễn của giải pháp, mô hình cần đáp ứng tốt được cả về độ chính xác cũng như về thời gian xử lý.

3. Các bước triển khai mô hình

Phần này mình sẽ hướng dẫn tổng quan train model trên colab.
Bật GPU: Runtime → Change runtime type và chọn GPU ở Hardware accelerator.

Ảnh có chứa văn bản

Mô tả được tạo tự động

Tiếp theo là tải source code mà mình đã chuẩn bị sẵn cho dễ sử dụng (link repo gốc https://github.com/PaddlePaddle/PaddleOCR):

!git clone https://github.com/hungcao0402/PaddleOCR-Vietnamese.git

Mình có tạo sẵn notebook lưu trong file Notebook.ipynb ở repo trên, bạn đọc có thể tham khảo.

Cài đặt thư viện cần thiết:

!pip install -r requirements.txt
!pip install paddlepaddle-gpu

Cả hai phần detect và recognize đều gồm các bước sau:

Bước 1: Tiền xử lý dữ liệu
Bước 2: Cấu hình file config
Bước 3: Train model
Bước 4: Evaluation
Bước 5: Convert sang model inference
Bước 6: Predict

4. Model Detection (SAST)

4.1. SAST

SAST (A Single-Shot Arbitrarily-Shaped Text Detector based on Context Attended Multi-Task Learning) là một mô hình hiệu quả trong việc dự đoán văn bản có hình dạng tùy ý tuy nhiên các vùng chữ nhỏ đang là vấn đề với nó.

4.2. Chuẩn bị dữ liệu

Bạn có thể tập dữ liệu huấn luyện từ repo https://github.com/VinAIResearch/dict-guided và download dataset Original với format x1,y1,x2,y2,x3,y3,x4,y4,TRANSCRIPT.

Sau khi tải về và giải nén nó ra chúng ta sẽ có:

Folder labels – chứa các file annotation của từng image
Folder train_images – chứa 1200 ảnh từ im0001 đến im1200
Folder test_image – chứa 300 ảnh từ im1201 đến 1500
Folder unseen_test_images – chứa 500 ảnh từ im1501 đến im2000
File general_dict.txt
File vn_dictionary.txt

Tiền xử lý dữ liệu, convert về format của PaddleOCR để train:

Để train model detection PaddleOCR, file annotation cần có format như sau:

Image file name Image annotation information encoded by json.dumps”
img_001.jpg [{“transcription”: “text”, “points”: [[310, 104], [416, 141], [418, 216], [312, 179]]}, {…}]

Ở đây

points sẽ là các cặp (x, y) 4 góc của text box theo chiều ngược kim đồng hồ, bắt đầu từ góc dưới bên trái.
transcription là text ở trong text box hiện tại. Khi chứa “###” thì có nghĩa là text box này invalid và sẽ skip đi khi train model.

Quan sát tập dữ liệu ta thấy:

Folder labels chứa từng file annotation cho mỗi ảnh.
Mỗi dòng trong file annotation chứa một text box.
Các cặp points được xếp theo chiều ngược kim đồng hồ và bắt đầu từ góc dưới bên trên và tách nhau bởi dấu phẩy, cuối cùng là text của box đó.

Ý tưởng để convert về format PaddleOCR:

Dùng một vòng lặp for đọc từng file annotation của mỗi ảnh.
Chạy một vòng lặp for đọc các dòng trong file annotation cho đến hết để lấy các text box.
Dùng split(‘,’, 8) để tách từng point và text ra sau đó lưu vào dictionary (trong lệnh split thêm tham số 8 để tránh trường hợp text chứa dấu phẩy sẽ bị thiếu dấu phẩy trong text).

Dưới đây là code convert của mình:

import glob 
import os
import numpy as np 
from tqdm import tqdm 
import pandas as pd  
import json

root_path = glob.glob('/content/vietnamese/labels/*')
train_label = open("train_label.txt","w")
test_label = open("test_label.txt","w")
useen_label = open("useen_label.txt","w")
for file in root_path:
    with open(file) as f:
      content = f.readlines()
      f.close()
    content = [x.strip() for x in content]
    text = []
    for i in content:
      label = {}
      i = i.split(',',8)
      label['transcription'] = i[-1]
      label['points'] = [[i[0],i[1]],[i[2],i[3]], [i[4],i[5]],[i[6],i[7]]]
      text.append(label)
    content = []
    text = json.dumps(text, ensure_ascii=False)
    img_name = os.path.basename(file).split('.')[0].split('_')[1]
    int_img = int(img_name)
    img_name = 'im' + "{:04n}".format(int(img_name)) + '.jpg'
    if int_img > 1500:
      useen_label.write( img_name+ '\t'+f'{text}' + '\n')
    elif int_img > 1200:
      test_label.write( img_name+ '\t'+f'{text}' + '\n')
    else:
      train_label.write( img_name+ '\t'+f'{text}' + '\n')

Vậy là ta đã chuẩn bị xong data train model detection, tiếp theo là bước chuẩn bị file config để train.

4.3. Chuẩn bị file config

Trong bước này bạn chọn model detection mà mình sử dụng. Sau khi thử train các model detection mà PaddleOCR hỗ trợ thì mình thấy SAST đạt hiệu quả tốt nhất cho nên mình sử dụng SAST với Backbone là ResNet50_vd trong cuộc thi này.

Sau khi quyết định được mô hình mà mình sẽ sử dụng, mình download pretrained model đã được train sẵn với dataset tiếng anh đạt kết quả cao với ICDAR2015 dataset.

Download Pretrained Model

Bạn đọc vào link sau đây và download pretrained SAST + ResNet50_vd:
https://github.com/PaddlePaddle/PaddleOCR/blob/release/2.4/doc/doc_en/algorithm_overview_en.md

Global:
       debug: false
       use_gpu: true
       epoch_num: 300
       log_smooth_window: 20
       print_batch_step: 2
       save_model_dir: ./output/SAST/
       save_epoch_step: 1
       eval_batch_step:
       - 40000
       - 50000
       cal_metric_during_train: false
       pretrained_model: ./pretrain_models/SAST/latest
       checkpoints:
       save_inference_dir: ./inference/SAST
       use_visualdl: True
       infer_img: null
       save_res_path: ./output/sast.txt
 …
 Train:
       dataset:
         name: SimpleDataSet
         data_dir: ./data #path_img_dir
         label_file_list: [ train_label.txt ]
         ratio_list: [1]
 …
 Eval:
       dataset:
         name: SimpleDataSet
         data_dir: #path_img_dir
         label_file_list:
         - unseen_label.txt

Chú thích

Global:
- use_gpu: sử dụng gpu hoặc không
- epoch_num: số epoch tối đa khi train .
- save_model_dir: là đường dẫn thư mục lưu checkpoint cho mỗi epoch khi train.
- pretrained_model: là đường dẫn thư mục pretrained model.
- checkpoints: là đường dẫn thư mục load checkpoint khi train.
Train:
- data_dir: là đường dẫn thư mục chứa ảnh train
- label_file_list: list các file annotation
- ratio_list: tỉ lệ train data set

4.4. Training

python3 tools/train.py -c ./configs/det/SAST.yml

Trong dòng code trên, sử dụng -c để chọn file cấu hình. Ngoài ra bạn còn có thể sử dụng -o để thay đổi các parameter mà không cần phải sửa file yml.

Sau mỗi save_epoch_step, checkpoint sẽ được lưu ở đường dẫn được đặt tại dòng Global.save_model_dir trong file config gồm 3 file:

iter_epoch_1.pdopt
iter_epoch_1.pdparams
iter_epoch_1.states

Ví dụ ta có thể train tiếp từ checkpoint trên mà không cần sửa file config bằng cách sau:

python3 tools/train.py -c ./configs/det/SAST.yml    \
    -o Global.checkpoints=./output/SAST/iter_epoch_1

Chú thích: Global.checkpoints sẽ được ưu tiên hơn Global.pretrained_model do đó nếu cả 2 đều được set thì Global.checkpoints sẽ được load để train. Nếu đường dẫn Global.checkpoints sai, Global.pretrained_model sẽ được load.

5. Model Recognition (SRN)

5.1. SRN

SRN (Semantic Reasoning Network) kết hợp cả phần visual information và global semantic information, hai phần này chạy song song trong quá trình huấn luyện và đưa ra kết quả.
SRN là một end-to-end trainable pipeline cho bài toán scene text recognition, bao gồm 4 phần:

Backbone network
Parallel visual attention module (PVAM)
Global semantic reasoning module (GSRM)
Visual semantic fusion decoder (VSFD)

Pipeline: Khi đưa ảnh đầu vào, đầu tiên, backbone network sẽ được dùng để trích xuất các thuộc tính 2D (V). Sau đó, PVAM được dùng để tạo ra N thuộc tính 1D sau khi được căn chỉnh, tại đó, mỗi thuộc tính tương ứng với một ký tự trong văn bản. N thuộc tính 1D này sau đó sẽ được đưa vào GSRM để thu thập thông tin ngữ nghĩa (S). Cuối cùng, các thuộc tính ảnh sau khi được căn chỉnh và ngữ nghĩa của nó sẽ được VSFD hợp nhất để dự đoán các ký tự.

5.2. Chuẩn bị dữ liệu

Data train model recognition thì nó sẽ khác so với model detection đó là data sẽ là các ảnh nhỏ và một file txt lưu đường dẫn ảnh cùng với chữ trong ảnh đó.

Format data

Có thể thấy rằng data để train model recognition là mỗi ảnh chứa một chữ trong khi data đã download lại là ảnh chứa nhiều chữ. Do vậy ta phải cắt nhỏ ảnh đã có từ các box chứa text.

import json
import os
import cv2
import copy
import numpy as np
from tools.infer.utility import draw_ocr_box_txt, get_rotate_crop_image
def print_draw_crop_rec_res( img_crop_list, img_name):
        bbox_num = len(img_crop_list)
        for bno in range(bbox_num):
          crop_name=img_name+'_'+str(bno)+'.jpg'
          crop_name_w = "./train/img_crop/{}".format(crop_name)
          cv2.imwrite(crop_name_w, img_crop_list[bno])
          crop_label.write("{0}\t{1}\n".format(crop_name, text[bno]))
crop_label = open('./train/crop_label.txt','w')
with open('./train/train_label.txt','r') as file_text:
  img_files=file_text.readlines()
count=0
for img_file in img_files:
  content = json.loads(img_file.split('\t')[1].strip())
  dt_boxes=[]
  text=[]
  for i in content:
    content = i['points']
    if i['transcription'] == "###":
      count+=1
      continue
    bb = np.array(i['points'],dtype=np.float32)
    dt_boxes.append(bb)
    text.append(i['transcription'])
  image_file = './train/vietnamese/train_images/' +img_file.split('\t')[0]
  img = cv2.imread(image_file)
  ori_im=img.copy()
  img_crop_list=[]
  for bno in range(len(dt_boxes)):
    tmp_box = copy.deepcopy(dt_boxes[bno])
    img_crop = get_rotate_crop_image(ori_im, tmp_box)
    img_crop_list.append(img_crop)
  img_name = img_file.split('\t')[0].split('.')[0]
  print_draw_crop_rec_res(img_crop_list,img_name)

5.3. Dictionary

Để dự đoán được tiếng việt thì ta cũng cần một file dictionary dành cho tiếng Việt chứa tất cả các kí tự. Vì PaddleOCR chưa hỗ trợ tiếng Việt nên mình đã tự làm một file dictionary riêng để train. Dictionary mình tạo lưu ở ppocr/utils/dict/vi_vietnam.txt.

5.4. Chuẩn bị file config

Global:

      debug: false

      use_gpu: true

      epoch_num: 200

      log_smooth_window: 20

      print_batch_step: 5

      save_model_dir: ./output/SRN

      save_epoch_step: 1

      eval_batch_step:

      - 30000

      - 40000

      cal_metric_during_train: false

      pretrained_model: ./pretrain_models/SRN/latest

      checkpoints: 

      save_inference_dir: ./inference/SRN

      use_visualdl: false

      infer_img: doc/imgs_words/ch/word_1.jpg

      character_dict_path: /content/drive/MyDrive/PaddleOCR/PaddleOCR/ppocr/utils/dict/vi_vietnam.txt

      character_type: ch

      max_text_length: 25

      num_heads: 8

      infer_mode: false

      use_space_char: true

…

Cách chỉnh sửa file config model recognition cũng tương tự như model detection. Chỉ có một số điểm khác là character_dict_path – đường dẫn file dictionary, use_space_char – dự đoán khoảng trắng hay không, max_text_length – độ dài tối đa kí tự trong một box.

5.5. Training

python3 tools/train.py -c ./configs/rec/SRN.yml \
           -o Global.use_gpu=True \
           Global.pretrained_model=pretrained_model \ 
           Global.character_type=ch

6. Evaluation

Một bước không thể thiếu đó là đánh giá mô hình để biết mô hình có đang hoạt động tốt không, chuyện gì xảy ra với mô hình. Từ đó giúp ta xác định được nên làm gì tiếp theo. Dataset dùng để evaluation được set ở mục Eval.dataset trong file config.

#evaluation model detection
python3 tools/eval.py -c SAST.yml \
                       -o Global.checkpoints _model=path_checkpoints

#evaluation model recognition
python3 tools/eval.py -c ./configs/rec/SRN.yml  \
                   -o Global.checkpoints=.path_checkpoints \
                      Global.character_type=ch  \
                       Global.character_dict_path=./ppocr/utils/dict/vi_vietnam.txt

7. Inference với model đã train xong

Cả 2 model detection và recognition được convert như sau:

python3 tools/export_model.py -c ./configs/det/SAST.yml  \
                              -o Global.pretrained_model= #path_pretrained \
                                 Global.save_inference_dir=./inference/SAST

python3 tools/export_model.py -c ./configs/rec/SRN.yml  \
                              -o Global.pretrained_model= #path_pretrained \
                                 Global.save_inference_dir=./inference/SRN

Convert xong ta sẽ có file và bây giờ thử inference. Bạn có thể thử predict với model mình đã train, link model inference trained:
https://drive.google.com/file/d/1-M5EMuLyFlOfp2bNMX6tibNmY3WReIxo/view?usp=sharing
https://drive.google.com/file/d/1slORHgS8awv7BeuSDahJvV1POPQtP_I_/view?usp=sharing

python3 tools/infer/predict_det.py --det_algorithm="SAST" --use_gpu=True \
                                   --det_model_dir="./inference/SAST"  \
                                   --image_dir=./test_data/TestA/img_3.jpg

Ảnh predict lưu ở inference_results:

python3 tools/infer/predict_rec.py --image_dir="/content/im0005_31.jpg" \
                             --use_gpu=True \
                                   --rec_algorithm="SRN" \
                                   --rec_model_dir="./inference/SRN/"  \
                                   --rec_image_shape="1, 64, 256"  \
                                   --rec_char_type="ch"   \
                                   --rec_char_dict_path="./ppocr/utils/dict/vi_vietnam.txt"

Kết hợp hai model detection và recognition với nhau:

Cuối cùng kết hợp cả hai model detection và recognition với nhau để predict ra kết quả:

python3 /content/drive/MyDrive/PaddleOCR/PaddleOCR/tools/infer/predict_system.py     
           --use_gpu=True  \
           --det_algorithm="SAST"  \
           --det_model_dir="./inference/SAST"  \
           --rec_algorithm="SRN" \
           --rec_model_dir="./inference/SRN/"  \
           --rec_image_shape="1, 64, 256"  \
           --image_dir=#path_img \ 
           --rec_char_type="ch"  \
           --drop_score=0.7  \
           --rec_char_dict_path="./ppocr/utils/dict/vi_vietnam.txt"

Sau khi chạy lệnh trên ta được kết quả sau đây:

8. Tạo docker image

Vòng chung kết sẽ phải nộp bài làm của mình dưới dạng docker image nên mình build docker image từ docker file (ở đây mình đẽ tạo sẵn trong repo có tên Dockerfile).

Khi đã có docker file, ta chạy dòng lệnh sau để tạo một docker image:

docker build -t sast_srn .

Cuối cùng chạy docker đã build trên với lệnh sau:

docker run -v test_data:/data/test_data:ro
submission_output:/data/submission_output sast_srn /bin/bash run.sh