【OpenMMLab AI训练营第二期】Class3: RTMPose

概述

结合MMDetection、MMPose、MMDeploy等算法库来实现三角板关键点检测。
其中采用了两阶段来结合两个不同任务来共同完成三角板关键点检测。
首先使用目标检测，分别训练Faster RCNN和RTMDet-Tiny对图像、视频进行检测。然后使用RTMPose-s来进行关键点检测。最后可以通过将模型转为ONNX格式来部署模型到终端上进行推理。

安装

安装MMCV

pip install -U openmim
mim install mmengine
mim install "mmcv==2.0.0rc3"
mim install "mmdet>=3.0.0rc6"

安装MMPose

git clone https://github.com/open-mmlab/mmpose.git -b tutorial2023
cd MMPose
mim install -e .

然后在MMPose目录下分别创建四个文件夹：checkpoint 用于存放预训练权重模型、outputs 用于存放预测的结果、data 用于存放训练数据、data/test 用于存放测试数据

安装MMDeteciton

git clone https://github.com/open-mmlab/mmdetection.git
cd mmdetection
pip install mmdetection -v -e .

同样建立四个子文件夹

三角板目标检测

分别采用Faster RCNN和RTMDet

下载并解压数据集

下载数据集

wget https://zihao-download.obs.cn-east-3.myhuaweicloud.com/yolov8/datasets/Traingle_215_Dataset/Triangle_215_Keypoint_coco.zip
unzip data/Traingle_215_Keypoint_coco.zip -d data >> /dev//null

下载测试数据

wget https://zihao-download.obs.myhuaweicloud.com/20220610-mmpose/triangle_datasets/test_img/triangle_1.jpg -P data/test
wget https://zihao-download.obs.myhuaweicloud.com/20220610-mmpose/triangle_datasets/test_img/triangle_2.jpg -P data/test
wget https://zihao-download.obs.myhuaweicloud.com/20220610-mmpose/triangle_datasets/test_img/triangle_3.jpg -P data/test
wget https://zihao-download.obs.myhuaweicloud.com/20220610-mmpose/triangle_datasets/test_img/triangle_4.jpg -P data/test
 
wget https://zihao-download.obs.myhuaweicloud.com/20220610-mmpose/triangle_datasets/videos/triangle_6.mp4 -P data/test
wget https://zihao-download.obs.myhuaweicloud.com/20220610-mmpose/triangle_datasets/videos/triangle_7.mp4 -P data/test
wget https://zihao-download.obs.myhuaweicloud.com/20220610-mmpose/triangle_datasets/videos/triangle_9.mp4 -P data/test

下载config

Faster RCNN

wget https://zihao-download.obs.myhuaweicloud.com/20220610-mmpose/triangle_datasets/faster_r_cnn_triangle.py -P data

RTMDet-Tiny

wget https://zihao-download.obs.myhuaweicloud.com/20220610-mmpose/triangle_datasets/rtmdet_tiny_triangle.py -P data

训练

cd MMPose
python tools/train.py data/faster_r_cnn_triangle.py #faster rcnn
python tools/train.py data/rtmdet_tiny_triangle.py #rtmdet-tiny

模型轻量化转换

因为保存的模型中有很多训练时的参数，在推理中是不需要的，所以可以删去。

python tools/model_converters/publish_model.py \
            work_dirs/faster_r_cnn_triangle/epoch_50.pth \
            checkpoint/faster_r_cnn_triangle_epoch_50_202305120846.pth

预测

python tools/model_converters/publish_model.py \
            work_dirs/faster_r_cnn_triangle/epoch_50.pth \
            checkpoint/faster_r_cnn_triangle_epoch_50_202305120846.pth

结果类似如下：

RTMPose关键点检测

训练和推理过程如上面章节