基于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 三连

03-21 19:01