首次尝试RT-DETR:一种实时印刷电路板(PCB)缺陷检测方法

目录

RT-DETR简介

RT-DETR在PCB缺陷检测中的应用

二、环境准备

2.1 数据准备

2.2 安装PaddleDetection以及依赖

三、模型选型

3.1 找到配置文件

3.2 修改配置文件

四、模型训练

训练100轮模型评估

训练过程可视化

五、预测推理

六、部署模型

6.1导出部署模型

6.2转换模型至ONNX

七、总结


在电子制造业中,对印刷电路板(PCB)的质量检测尤为重要,任何微小的缺陷都可能导致电路板的性能下降,甚至无法使用。为了提高这一过程的效率和准确度,我们决定尝试使用最新的对象检测模型RT-DETR来进行PCB缺陷检测。接下来,让我们详细介绍这一过程。

RT-DETR简介

RT-DETR,即实时对象检测与追踪系统(Real-Time Detection and Tracking System),是一款基于视觉变换器(ViT)模型的全新深度学习算法。该模型的目标是实时检测并追踪各种环境中的潜在风险,并通过预警提高工作人员的安全性。

RT-DETR在PCB缺陷检测中的应用

我们决定使用RT-DETR来对PCB进行缺陷检测,以下是我们的操作步骤:

  1. 数据集构建:首先,我们收集了大量的PCB图片,包括正常的和存在各种缺陷的。然后对这些图片进行标注,构建了一个PCB缺陷检测数据集。

  2. 模型训练:使用RT-DETR模型在我们构建的数据集上进行训练。通过对模型进行反复的训练和调整,使得它能够准确地识别出PCB图片中的各种缺陷。

  3. 实时检测:将训练好的模型部署在PCB生产线上,通过摄像头实时获取PCB图片,然后通过RT-DETR模型对图片进行分析,识别出任何可能的缺陷,并立即发出警告。

最近PaddleDetection的develop分支中新出了RT-DETR的实时端到端目标检测框架,RT-DETR是第一个实时端到端目标检测器。通过高效的混合编码器,解耦尺度内交互和跨尺度融合来高效处理多尺度特征。此外,RT-DETR支持通过使用不同的解码器层来灵活调整推理速度,而不需要重新训练,这有助于实时目标检测器的实际应用。

RT-DETR-L在COCO val2017上实现了53.0%的AP,在T4 GPU上实现了114FPS,RT-DETR-X实现了54.8%的AP和74FPS,在速度和精度方面都优于相同规模的所有YOLO检测器。RT-DETR-R50实现了53.1%的AP和108FPS,RT-DETR-R101实现了54.3%的AP和74FPS,在精度上超过了全部使用相同骨干网络的DETR检测器。可以说RT-DETR是目前目标检测领域又快又好的检测器,并且基于Vit,有巨大发展前景。

工业缺陷检测场景对模型的快速性和准确性要求较高,基于此本文采用RT-DETR-R101在PCB数据集上进行缺陷检测。

根据PaddleDetection给出的模型性能对比,各模型结构和骨干网络的代表模型在COCO数据集上精度mAP和T4 TensorRT FP16上预测速度(FPS)对比图如下:

Model Epoch backbone input shape �����APval ��50���AP50val​ Params(M) FLOPs(G) T4 TensorRT
FP16(FPS)
RT-DETR-R50 6x ResNet-50 640 53.1 71.3 42 136 108
RT-DETR-R101 6x ResNet-101 640 54.3 72.7 76 259 74
RT-DETR-L 6x HGNetv2 640 53.0 71.6 32 110 114
RT-DETR-X 6x HGNetv2 640 54.8 73.1 67 234 74

二、环境准备

2.1 数据准备

本次数据集采用PKU-Market-PCB,该数据集用于印刷电路板(PCB)的瑕疵检测,提供了6种常见的PCB缺陷(漏孔、鼠咬、开路、短路、杂散、杂铜)。 同时,该数据集提供了train_shots10,train_shots30用于小样本缺陷检测,感兴趣的可以通过更换数据集路径进行尝试。

数据集来源:https://robotics.pkusz.edu.cn/resources/dataset/

In [ ]

# 解压数据集(仅运行一次)
!tar -xvf data/data185667/pcb.tar -C ./data/

In [2]

# 组织数据目录
!mkdir MyDataset
!mkdir MyDataset/JPEGImages

In [3]

!mv data/pcb/images/* MyDataset/JPEGImages/
!mv data/pcb/pcb_cocoanno/* MyDataset/

2.2 安装PaddleDetection以及依赖

In [4]

!git clone https://gitee.com/paddlepaddle/PaddleDetection.git
正克隆到 'PaddleDetection'...
remote: Enumerating objects: 257323, done.
remote: Counting objects: 100% (2320/2320), done.
remote: Compressing objects: 100% (1349/1349), done.
remote: Total 257323 (delta 1420), reused 1762 (delta 952), pack-reused 255003
接收对象中: 100% (257323/257323), 413.67 MiB | 5.12 MiB/s, 完成.
处理 delta 中: 100% (210448/210448), 完成.
检查连接... 完成。

In [1]

%cd PaddleDetection
/home/aistudio/PaddleDetection

In [ ]

# 安装依赖
!pip install -r requirements.txt --user
!python setup.py install --user

In [3]

# 切换到develop分支
!git checkout develop
M	configs/datasets/coco_detection.yml
M	configs/rtdetr/_base_/optimizer_6x.yml
M	configs/rtdetr/rtdetr_r101vd_6x_coco.yml
M	configs/runtime.yml
已经位于 'develop'
您的分支与上游分支 'origin/develop' 一致。

三、模型选型

RT-DETR提供了四种模型进行选择,具体的模型选型可以根据自己数据集的情况选取精度与速度相适应的模型,本次采用RT-DETR-R101

3.1 找到配置文件

配置文件在PaddleDetection/configs/rtdetr/rtdetr_r101vd_6x_coco.yml,主要包含五个部分,为了方便更改,将其copy到根目录下 

3.2 修改配置文件

找到配置文件后,需要对其进行修改,参考链接https://github.com/PaddlePaddle/PaddleDetection/blob/release/2.6/docs/advanced_tutorials/customization/detection.md

在这里主要修改:

  • coco_detection 文件修改数据集路径,将num_classes=6,因为有6类缺陷,将训练和验证集路径修改到解压好的个人数据集路径。

  • 学习率训练轮数等超参数,因为官方是在4卡进行训练,因此在单卡情况下修改base_lr=0.0001/4 = 0.000025

In [4]

# 训练配置文件覆盖
!cp ../rtdetr_r101vd_6x_coco.yml configs/rtdetr/rtdetr_r101vd_6x_coco.yml
!cp ../coco_detection.yml configs/datasets/coco_detection.yml
!cp ../runtime.yml configs/runtime.yml
!cp ../optimizer_6x.yml configs/rtdetr/_base_/optimizer_6x.yml
!cp ../rtdetr_r50vd.yml configs/rtdetr/_base_/rtdetr_r50vd.yml
!cp ../rtdetr_reader.yml configs/rtdetr/_base_/rtdetr_reader.yml

四、模型训练

训练100轮map为0.46,average FPS: 18.4

In [ ]

# 开始训练
# 恢复训练 -r output/rtdetr_r101vd_6x_coco/best_model
!python tools/train.py \
-c configs/rtdetr/rtdetr_r101vd_6x_coco.yml \
--use_vdl=True \
--vdl_log_dir=../work/ \
--eval 

训练100轮模型评估

训练过程可视化

In [7]

# 模型评估
!python -u tools/eval.py \
-c configs/rtdetr/rtdetr_r101vd_6x_coco.yml \
-o weights=output/rtdetr_r101vd_6x_coco/best_model
/opt/conda/envs/python35-paddle120-env/lib/python3.7/site-packages/matplotlib/__init__.py:107: DeprecationWarning: Using or importing the ABCs from 'collections' instead of from 'collections.abc' is deprecated, and in 3.8 it will stop working
  from collections import MutableMapping
/opt/conda/envs/python35-paddle120-env/lib/python3.7/site-packages/matplotlib/rcsetup.py:20: DeprecationWarning: Using or importing the ABCs from 'collections' instead of from 'collections.abc' is deprecated, and in 3.8 it will stop working
  from collections import Iterable, Mapping
/opt/conda/envs/python35-paddle120-env/lib/python3.7/site-packages/matplotlib/colors.py:53: DeprecationWarning: Using or importing the ABCs from 'collections' instead of from 'collections.abc' is deprecated, and in 3.8 it will stop working
  from collections import Sized
W0421 18:43:59.137140 12824 gpu_resources.cc:61] Please NOTE: device: 0, GPU Compute Capability: 7.0, Driver API Version: 11.2, Runtime API Version: 11.2
W0421 18:43:59.142907 12824 gpu_resources.cc:91] device: 0, cuDNN Version: 8.2.
loading annotations into memory...
Done (t=0.00s)
creating index...
index created!
[04/21 18:44:01] ppdet.data.source.coco INFO: Load [138 samples valid, 0 samples invalid] in file /home/aistudio/MyDataset/val.json.
[04/21 18:44:04] ppdet.utils.checkpoint INFO: Finish loading model weights: output/rtdetr_r101vd_6x_coco/best_model.pdparams
[04/21 18:44:06] ppdet.engine INFO: Eval iter: 0
[04/21 18:44:13] ppdet.metrics.metrics INFO: The bbox result is saved to bbox.json.
loading annotations into memory...
Done (t=0.11s)
creating index...
index created!
[04/21 18:44:13] ppdet.metrics.coco_utils INFO: Start evaluate...
Loading and preparing results...
DONE (t=0.56s)
creating index...
index created!
Running per image evaluation...
Evaluate annotation type *bbox*
DONE (t=2.29s).
Accumulating evaluation results...
DONE (t=0.28s).
 Average Precision  (AP) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.464
 Average Precision  (AP) @[ IoU=0.50      | area=   all | maxDets=100 ] = 0.921
 Average Precision  (AP) @[ IoU=0.75      | area=   all | maxDets=100 ] = 0.400
 Average Precision  (AP) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.233
 Average Precision  (AP) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.467
 Average Precision  (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.465
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=  1 ] = 0.134
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets= 10 ] = 0.550
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.596
 Average Recall     (AR) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.233
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.600
 Average Recall     (AR) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.580
[04/21 18:44:16] ppdet.engine INFO: Total sample number: 138, average FPS: 18.449997011669787

五、预测推理

在paddledetection套件中可以预测单张图片也可以批量预测,只需修改一些参数,十分方便 

In [ ]

# 挑一张验证集的图片展示预测效果
!python tools/infer.py \
-c configs/rtdetr/rtdetr_r101vd_6x_coco.yml \
-o weights=output/rtdetr_r101vd_6x_coco/best_model \
--infer_img=../MyDataset/JPEGImages/01_missing_hole_15.jpg \
--save_results=True

In [ ]

# 批量预测
!python tools/infer.py \
-c configs/rtdetr/rtdetr_r101vd_6x_coco.yml \
-o weights=output/rtdetr_r101vd_6x_coco/best_model \
--infer_dir=../MyDataset/JPEGImages/ \
--output_dir=../output/infer_images/ \
--draw_threshold=0.5 \
--save_results=True

六、部署模型

6.1导出部署模型

In [ ]

#导出部署模型
!python tools/export_model.py \
-c configs/rtdetr/rtdetr_r101vd_6x_coco.yml \
-o weights=output/rtdetr_r101vd_6x_coco/best_model

6.2转换模型至ONNX

In [ ]

#安装Paddle2ONNX 和 ONNX
!pip install onnx==1.13.0
!pip install paddle2onnx==1.0.5

In [ ]

!paddle2onnx --model_dir=./output_inference/rtdetr_r101vd_6x_coco/ \
            --model_filename model.pdmodel  \
            --params_filename model.pdiparams \
            --opset_version 16 \
            --save_file rtdetr_r101vd_6x_coco.onnx

七、总结

  • 本文利用RT-DETR-R101在工业缺陷检测数据集上进行模型的训练验证与推理部署,通过100轮的训练模型的map为0.46,V100环境下average FPS: 18.4,实现了又快又好的效果。

  • 由于是PCB缺陷检测,属于小目标,可以尝试通过SIHI切图方式进行进一步的AP提升

  • DETR实现了真正端到端的目标检测,真正的难题在于推理速度慢,RT-DETR较好的解决这一问题,希望FastDeploy可以尽快支持RT-DETR进行模型部署。

猜你喜欢

转载自blog.csdn.net/m0_68036862/article/details/131348586