基於PaddleOCR的多視角集裝箱箱號檢測識別

2023-03-21 21:00:56

基於PaddleOCR的多視角集裝箱箱號檢測識別

一、專案介紹

集裝箱號是指裝運出口貨物集裝箱的箱號,填寫託運單時必填此項。標準箱號構成基本概念:採用ISO6346(1995)標準

標準集裝箱箱號由11位編碼組成,如:CBHU 123456 7,包括三個部分:

  1. 第一部分由4位元英文字母組成。前三位程式碼主要說明箱主、經營人,第四位程式碼說明集裝箱的型別。列如CBHU 開頭的標準集裝箱是表明箱主和經營人為中遠集運
  2. 第二部分由6位數位組成。是箱體註冊碼,用於一個集裝箱箱體持有的唯一標識
  3. 第三部分為校驗碼由前4位元字母和6位數位經過校驗規則運算得到,用於識別在校驗時是否發生錯誤。即第11位編號

本教學基於PaddleOCR進行集裝箱箱號檢測識別任務,使用少量資料分別訓練檢測、識別模型,最後將他們串聯在一起實現集裝箱箱號檢測識別的任務

效果展示:

二、環境準備

首先點選左側套件選擇PaddleOCR 進行下載。

三、資料集介紹

本教學所使用的集裝箱箱號資料集,該資料包含3003張解析度為1920×1080的集裝箱影象

1、PaddleOCR檢測模型訓練標註規則如下,中間用"\t"分隔:

" 影象檔名                    json.dumps編碼的影象標註資訊"
ch4_test_images/img_61.jpg    [{"transcription": "MASA", "points": [[310, 104], [416, 141], [418, 216], [312, 179]]}, {...}]

其中json.dumps編碼前的影象標註資訊是包含多個字典的list,字典中的 points 表示文字方塊的四個點的座標(x, y),從左上角的點開始順時針排列。 transcription 表示當前文字方塊的文字,當其內容為「###」時,表示該文字方塊無效,在訓練時會跳過。

2、PaddleOCR識別模型訓練標註規則如下,中間用"\t"分隔:

" 影象檔名                 影象標註資訊 "

train_data/rec/train/word_001.jpg   簡單可依賴
train_data/rec/train/word_002.jpg   用科技讓複雜的世界更簡單

四、資料整理

4.1 檢測模型所需資料準備

將資料集3000張圖片按2:1劃分成訓練集和驗證集,執行以下程式碼

from tqdm import tqdm
finename = "all_label.txt"
f = open(finename)
lines = f.readlines() 
t = open('det_train_label.txt','w')
v = open('det_eval_label.txt','w')
count = 0
for line in tqdm(lines):
    if count < 2000:
        t.writelines(line)
        count += 1
    else:
        v.writelines(line)
f.close()
t.close()
v.close()

4.2 識別模型所需資料準備

我們根據檢測部分的註釋,裁剪資料集儘可能只包含文字部分圖片作為識別的資料,執行以下程式碼

from PIL import Image
import json
from tqdm import tqdm
import os
import numpy as np
import cv2
import math

from PIL import Image, ImageDraw

class Rotate(object):

    def __init__(self, image: Image.Image, coordinate):
        self.image = image.convert('RGB')
        self.coordinate = coordinate
        self.xy = [tuple(self.coordinate[k]) for k in ['left_top', 'right_top', 'right_bottom', 'left_bottom']]
        self._mask = None
        self.image.putalpha(self.mask)

    @property
    def mask(self):
        if not self._mask:
            mask = Image.new('L', self.image.size, 0)
            draw = ImageDraw.Draw(mask, 'L')
            draw.polygon(self.xy, fill=255)
            self._mask = mask
        return self._mask

    def run(self):
        image = self.rotation_angle()
        box = image.getbbox()
        return image.crop(box)

    def rotation_angle(self):
        x1, y1 = self.xy[0]
        x2, y2 = self.xy[1]
        angle = self.angle([x1, y1, x2, y2], [0, 0, 10, 0]) * -1
        return self.image.rotate(angle, expand=True)

    def angle(self, v1, v2):
        dx1 = v1[2] - v1[0]
        dy1 = v1[3] - v1[1]
        dx2 = v2[2] - v2[0]
        dy2 = v2[3] - v2[1]
        angle1 = math.atan2(dy1, dx1)
        angle1 = int(angle1 * 180 / math.pi)
        angle2 = math.atan2(dy2, dx2)
        angle2 = int(angle2 * 180 / math.pi)
        if angle1 * angle2 >= 0:
            included_angle = abs(angle1 - angle2)
        else:
            included_angle = abs(angle1) + abs(angle2)
            if included_angle > 180:
                included_angle = 360 - included_angle
        return included_angle



def image_cut_save(path, bbox, save_path):
    """
    :param path: 圖片路徑
    :param left: 區塊左上角位置的畫素點離圖片左邊界的距離
    :param upper:區塊左上角位置的畫素點離圖片上邊界的距離
    :param right:區塊右下角位置的畫素點離圖片左邊界的距離
    :param lower:區塊右下角位置的畫素點離圖片上邊界的距離
    """
    img_width  = 1920
    img_height = 1080
    img = Image.open(path)
    coordinate = {'left_top': bbox[0], 'right_top': bbox[1], 'right_bottom': bbox[2], 'left_bottom': bbox[3]}
    rotate = Rotate(img, coordinate)
    
    left, upper = bbox[0]
    right, lower = bbox[2]
    if lower-upper > right-left:
        rotate.run().convert('RGB').transpose(Image.ROTATE_90).save(save_path)
    else:
        rotate.run().convert('RGB').save(save_path)
    return True

#讀取檢測標註製作識別資料集
files = ["det_train_label.txt","det_eval_label.txt"]
filetypes =["train","eval"]
for index,filename in enumerate(files):
    f = open(filename)
    l = open('rec_'+filetypes[index]+'_label.txt','w')
    if index == 0:
        data_dir = "RecTrainData"
    else:
        data_dir = "RecEvalData"
    if not os.path.exists(data_dir):
        os.mkdir(data_dir)
    lines = f.readlines() 
    for line in tqdm(lines):
        image_name = line.split("\t")[0].split("/")[-1]
        annos = json.loads(line.split("\t")[-1])
        img_path = os.path.join("/home/aistudio/input0/images",image_name)
        for i,anno in enumerate(annos):
            data_path = os.path.join(data_dir,str(i)+"_"+image_name)
            if image_cut_save(img_path,anno["points"],data_path):
                l.writelines(str(i)+"_"+image_name+"\t"+anno["transcription"]+"\n")
    l.close()
    f.close()

五、實驗

由於資料集比較少,為了模型更好和更快的收斂,這裡選用 PaddleOCR 中的 PP-OCRv3 模型進行檢測和識別。PP-OCRv3在PP-OCRv2的基礎上,中文場景端到端Hmean指標相比於PP-OCRv2提升5%, 英文數位模型端到端效果提升11%。詳細優化細節請參考PP-OCRv3技術報告。

5.1 檢測模型

5.1.1 檢測模型設定

PaddleOCR提供了許多檢測模型,在路徑PaddleOCR-2.6.0/configs/det下可找到模型及其組態檔。如我們選用模型ch_PP-OCRv3_det_student.yml,其組態檔路徑在:PaddleOCR-2.6.0/configs/det/ch_PP-OCRv3/ch_PP-OCRv3_det_student.yml。使用前需對其進行必要的設定,如訓練引數、資料集路徑等。將部分關鍵設定展示如下:

#關鍵訓練引數
use_gpu: true #是否使用顯示卡
epoch_num: 1200 #訓練epoch個數
save_model_dir: ./output/ch_PP-OCR_V3_det/ #模型儲存路徑
save_epoch_step: 200 #每訓練200epoch,儲存一次模型
eval_batch_step: [0, 100] #訓練每迭代100次,進行一次驗證
pretrained_model: ./PaddleOCR-release
2.5/pretrain_models/ch_PP-OCR_V3_det/best_accuracy.pdparams #預訓練模型路徑
#訓練集路徑設定
Train:
  dataset:
    name: SimpleDataSet
    data_dir: /input0/images #圖片資料夾路徑
    label_file_list:
      - ./det_train_label.txt #標籤路徑

檔案直接放在更目錄裡,自行替換即可 /home/aistudio/ch_PP-OCRv3_det_student.yml

5.1.2 模型微調

在notebook中執行如下命令對模型進行微調,其中 -c 傳入的為設定好的模型檔案路徑

%run /home/aistudio/PaddleOCR-2.6.0/tools/train.py \
    -c /home/aistudio/PaddleOCR-2.6.0/configs/det/ch_PP-OCRv3/ch_PP-OCRv3_det_student.yml

使用預設超引數,模型ch_PP-OCRv3_det_student在訓練集上訓練100個epoch後,模型在驗證集上的hmean達到:90.96%,此後再無明顯增長

[2023/03/21 15:57:09] ppocr INFO: best metric, hmean: 0.909551282051282, precision: 0.8977836411609498,
recall: 0.921611681990265, fps: 20.347745459258228, best_epoch: 100

5.2 識別模型

5.2.1 識別模型設定

PaddleOCR也提供了許多識別模型,在路徑PaddleOCR-2.6.0/configs/rec下可找到模型及其組態檔。如我們選用模型ch_PP-OCRv3_rec_distillation,其組態檔路徑在:PaddleOCR-2.6.0/configs/rec/PP-OCRv3/ch_PP-OCRv3_rec_distillation.yml。使用前需對其進行必要的設定,如訓練引數、資料集路徑等。將部分關鍵設定展示如下:

#關鍵訓練引數
use_gpu: true #是否使用顯示卡
epoch_num: 1200 #訓練epoch個數
save_model_dir: ./output/rec_ppocr_v3_distillation #模型儲存路徑
save_epoch_step: 200 #每訓練200epoch,儲存一次模型
eval_batch_step: [0, 100] #訓練每迭代100次,進行一次驗證
pretrained_model: ./PaddleOCR-release-2.5/pretrain_models/PPOCRv3/best_accuracy.pdparams #預訓練模型路徑
#訓練集路徑設定
Train:
  dataset:
    name: SimpleDataSet
    data_dir: ./RecTrainData/ #圖片資料夾路徑
    label_file_list:
      - ./rec_train_label.txt #標籤路徑
      
      Eval:
  dataset:
    name: SimpleDataSet
    data_dir: ./RecEvalData/
    label_file_list:
    - ./rec_eval_label.txt

檔案直接放在更目錄裡,自行替換即可 /home/aistudio/ch_PP-OCRv3_rec_distillation.yml

5.2.2 模型微調

在notebook中執行如下命令對模型進行微調,其中 -c 傳入的為設定好的模型檔案路徑

%run /home/aistudio/PaddleOCR-2.6.0/tools/train.py \
    -c /home/aistudio/PaddleOCR-2.6.0/configs/rec/PP-OCRv3/ch_PP-OCRv3_rec_distillation.yml

使用預設超引數,模型ch_PP-OCRv3_rec_distillation在訓練集上訓練50個epoch後,模型在驗證集上的精度達到:91.11%,此後再無明顯增長

[2023/03/21 20:04:28] ppocr INFO: best metric, acc: 0.9110600272522444, norm_edit_dis: 0.9427426548965615,
Teacher_acc: 0.9040291998159589, Teacher_norm_edit_dis: 0.9405629345025616, fps: 246.029195787707, best_epoch: 50

六、結果展示

6.1 檢測模型推理

在notebook中執行如下命令使用微調過的模型檢測測試圖片中的文字,其中:

  • Global.infer_img 為圖片路徑或圖片資料夾路徑,
  • Global.pretrained_model 為微調過的模型,
  • Global.save_res_path 為推理結果儲存路徑
%run /home/aistudio/PaddleOCR-2.6.0/tools/infer_det.py \
    -c /home/aistudio/PaddleOCR-2.6.0/configs/det/ch_PP-OCRv3/ch_PP-OCRv3_det_student.yml \
    -o Global.infer_img="/home/aistudio/input0/images" Global.pretrained_model="./output/ch_PP-OCR_V3_det/best_accuracy" Global.save_res_path="./output/det_infer_res/predicts.txt"

6.2 識別模型推理

在notebook中執行如下命令使用微調過的模型檢測測試圖片中的文字,其中:

  • Global.infer_img 為圖片路徑或圖片資料夾路徑,
  • Global.pretrained_model 為微調過的模型,
  • Global.save_res_path 為推理結果儲存路徑
%run /home/aistudio/PaddleOCR-2.6.0/tools/infer_rec.py \
    -c /home/aistudio/PaddleOCR-2.6.0/configs/rec/PP-OCRv3/ch_PP-OCRv3_rec_distillation.yml \
    -o Global.infer_img="./RecEvalData/" Global.pretrained_model="./output/rec_ppocr_v3_distillation/best_accuracy" Global.save_res_path="./output/rec_infer_res/predicts.txt"
    

部分結果展示:

./RecEvalData/0_1-122720001-OCR-AS-B01.jpg	{"Student": {"label": "EITU1786393", "score": 0.9737951755523682}, "Teacher": {"label": "EITU1786393", "score": 0.9882291555404663}}
./RecEvalData/0_1-122720001-OCR-LB-C02.jpg	{"Student": {"label": "EITU1786393", "score": 0.9709678888320923}, "Teacher": {"label": "EITU1786393", "score": 0.9925146698951721}}
./RecEvalData/0_1-122720001-OCR-RF-D01.jpg	{"Student": {"label": "EITU1786393", "score": 0.9985160231590271}, "Teacher": {"label": "EITU1786393", "score": 0.9967824816703796}}
./RecEvalData/0_1-122728001-OCR-RF-D01.jpg	{"Student": {"label": "DFSU4119250", "score": 0.9663339257240295}, "Teacher": {"label": "DFSU4119250", "score": 0.9600133299827576}}
./RecEvalData/0_1-122740001-OCR-AH-A01.jpg	{"Student": {"label": "MRKU4306585", "score": 0.9916775226593018}, "Teacher": {"label": "MRKU4306585", "score": 0.9929805994033813}}
./RecEvalData/0_1-122749001-OCR-AH-A01.jpg	{"Student": {"label": "FCGU4996010", "score": 0.9195910096168518}, "Teacher": {"label": "FCGU4996010", "score": 0.9424482583999634}}
./RecEvalData/0_1-122830001-OCR-AS-B01.jpg	{"Student": {"label": "MEDU4024195", "score": 0.9861812591552734}, "Teacher": {"label": "MEDU4024195", "score": 0.9718942642211914}}
./RecEvalData/0_1-122843001-OCR-RF-D01.jpg	{"Student": {"label": "TGU864295", "score": 0.9045045375823975}, "Teacher": {"label": "TGU864395", "score": 0.8963061571121216}}

6.3 檢測識別模型串聯推理

6.3.1 模型轉換

在串聯推理前首先需要將訓練儲存的模型轉換成推理模型,分別執行如下檢測命令即可。其中,-c傳入要轉換模型的組態檔路徑,-o Global.pretrained_model為要被轉換的模型檔案,Global.save_inference_dir為轉換得到推理模型的儲存路徑

# 檢測模型轉換
%run /home/aistudio/PaddleOCR-2.6.0/tools/export_model.py \
-c /home/aistudio/PaddleOCR-2.6.0/configs/det/ch_PP-OCRv3/ch_PP-OCRv3_det_student.yml  \
-o Global.pretrained_model="./output/ch_PP-OCR_V3_det/best_accuracy" Global.save_inference_dir="./output/det_inference/"


# 識別模型轉換
%run /home/aistudio/PaddleOCR-2.6.0/tools/export_model.py \
-c /home/aistudio/PaddleOCR-2.6.0/configs/rec/PP-OCRv3/ch_PP-OCRv3_rec_distillation.yml \
-o Global.pretrained_model="./output/rec_ppocr_v3_distillation/best_accuracy" Global.save_inference_dir="./output/rec_inference/"

6.3.2 模型串聯推理

轉換完畢後,PaddleOCR提供了檢測和識別模型的串聯工具,可以將訓練好的任一檢測模型和任一識別模型串聯成兩階段的文字識別系統。輸入影象經過文字檢測、檢測框矯正、文字識別、得分過濾四個主要階段輸出文字位置和識別結果。執行程式碼如下,其中image_dir為單張影象或者影象集合的路徑,det_model_dir為檢測inference模型的路徑,rec_model_dir為識別inference模型的路徑。視覺化識別結果預設儲存到 ./inference_results 資料夾裡面

%run /home/aistudio/PaddleOCR-2.6.0/tools/infer/predict_system.py \
--image_dir="OCRTest" \
--det_model_dir="./output/det_inference/" \
--rec_model_dir="./output/rec_inference/Student/"

結果展示:

1-122700001-OCR-LF-C01.jpg	[{"transcription": "TTEMU3108252", "points": [[1226, 133], [1322, 133], [1322, 883], [1226, 883]]}, {"transcription": "22G1", "points": [[1417, 214], [1479, 216], [1471, 463], [1409, 461]]}]
1-122720001-OCR-AH-A01.jpg	[{"transcription": "ITU1786393", "points": [[225, 206], [918, 215], [917, 318], [224, 309]]}]
1-122720001-OCR-AS-B01.jpg	[{"transcription": "EITU1786393", "points": [[919, 283], [1389, 296], [1387, 372], [917, 359]]}, {"transcription": "45G1", "points": [[1104, 399], [1288, 399], [1288, 486], [1104, 486]]}]
1-122720001-OCR-LB-C02.jpg	[{"transcription": "TU1", "points": [[226, 10], [515, 6], [516, 97], [227, 102]]}, {"transcription": "45G1", "points": [[489, 114], [784, 104], [787, 204], [492, 213]]}]
1-122720001-OCR-RF-D01.jpg	[{"transcription": "EITU1786393", "points": [[216, 38], [941, 27], [942, 125], [217, 135]]}, {"transcription": "45G1", "points": [[452, 137], [719, 133], [720, 218], [453, 223]]}]

以1-122720001-OCR-AS-B01.jpg測試樣例進行展示:

多視角識別結果為:EITU1786393

七.總結

本專案做了基於PaddleOCR的多視角集裝箱箱號檢測識別,使用少量資料分別訓練檢測、識別模型,最後將他們串聯在一起實現集裝箱箱號檢測識別的任務。其中集裝箱號是指裝運出口貨物集裝箱的箱號,填寫託運單時必填此項。標準箱號構成基本概念:採用ISO6346(1995)標準。

從結果上看,基於PaddleOCR的多視角集裝箱箱號檢測識別取得了不錯的效果,但也存在一些改進地方。

  • 資料集樣本量和豐富度不夠
  • 訓練不充分,因為算力和時間限制,本人就簡單訓練了100 epochs左右停止了。
  • 模型引數調優(目前預設引數)
    等等

源專案連結:https://aistudio.baidu.com/aistudio/projectdetail/5766320?contributionType=1

歡迎關注fork 三連