RT-DETR初体验：实现PCB缺陷检测

Paddledection-RT-DETR实现PCB缺陷检测

AI Studio

1169人浏览 · 2023-05-05 11:12:52

AI Studio · 2023-05-05 11:12:52 发布

★★★ 本文源自AlStudio社区精品项目，【点击此处】查看更多精品内容 >>>

一、项目背景

论文链接：DETRs Beat YOLOs on Real-time Object Detection

最近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	$AP^{val}$	$AP^{val}_{50}$	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/

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

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

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

2.2 安装PaddleDetection以及依赖

!git clone https://gitee.com/paddlepaddle/PaddleDetection.git

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

%cd PaddleDetection

/home/aistudio/PaddleDetection

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

# 切换到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

# 训练配置文件覆盖
!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

# 开始训练
# 恢复训练 -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轮模型评估

训练过程可视化

# 模型评估
!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

Total sample number: 138, average FPS: 18.449997011669787

五、预测推理

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

# 挑一张验证集的图片展示预测效果
!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

# 批量预测
!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导出部署模型

#导出部署模型
!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

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

!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进行模型部署。

此文章为搬运
原项目链接