YOLOv7-Pose
实现YOLOv7:可训练的免费套件为实时目标检测设置了最新技术标准
YOLOv7-Pose的姿态估计是基于YOLO-Pose的。关键点标签采用MS COCO 2017数据集。
训练
使用预训练模型yolov7-w6-person.pt
进行训练。训练命令如下:
python -m torch.distributed.launch --nproc_per_node 8 --master_port 9527 train.py --data data/coco_kpts.yaml --cfg cfg/yolov7-w6-pose.yaml --weights weights/yolov7-w6-person.pt --batch-size 128 --img 960 --kpt-label --sync-bn --device 0,1,2,3,4,5,6,7 --name yolov7-w6-pose --hyp data/hyp.pose.yaml
部署
- 导出ONNX模型
运行以下命令生成onnx模型和引擎模型:
python models/export_onnx.py \
--weights weights/yolov7-w6-pose.pt \
--img-size 832 \
--device 0 \
--batch-size 1 \
--simplify
- 导出TensorRT模型
使用脚本:
python models/export_TRT.py \
--onnx weights/yolov7-w6-pose.onnx \
--batch-size 1 \
--device 1 \
--fp16
或者使用trtexec:
trtexec \
--onnx=weights/yolov7-w6-pose.onnx \
--workspace=4096 \
--saveEngine=weights/yolov7-w6-pose-FP16.engine \
--fp16
推理
- PyTorch模型推理
python detect_multi_backend.py \
--weights weights/yolov7-w6-pose.pt \
--source data/images \
--device 0 \
--img-size 832 \
--kpt-label
- ONNX模型推理
python detect_multi_backend.py \
--weights weights/yolov7-w6-pose.onnx \
--source data/images \
--device 0 \
--img-size 832 \
--kpt-label
- TensorRT模型推理
python detect_multi_backend.py \
--weights weights/yolov7-w6-pose.engine \
--source data/images \
--device 0 \
--img-size 832 \
--kpt-label
测试
使用yolov7-w6-pose.pt
进行测试:
官方YOLOv7-pose和YOLO-Pose代码只在test.py
中计算检测mAP。若要计算关键点mAP,需使用COCO API。在此仓库中实现的oks_iou
矩阵计算加速了关键点mAP的计算。测试关键点mAP时,oks
区域设置为0.6
乘以ground truth box的区域。
- 测试PyTorch模型
python test_multi_backend.py \
--weights weights/yolov7-w6-pose.pt \
--data data/coco_kpts.yaml \
--img-size 832 \
--conf-thres 0.001 \
--iou-thres 0.6 \
--task val \
--device 0 \
--kpt-label
- 测试ONNX模型
python test_multi_backend.py \
--weights weights/yolov7-w6-pose.onnx \
--data data/custom_kpts.yaml \
--img-size 832 \
--conf-thres 0.001 \
--iou-thres 0.6 \
--task val \
--device 0 \
--kpt-label
- 测试TensorRT模型
python test_multi_backend.py \
--weights weights/yolov7-w6-pose-FP16.engine \
--data data/coco_kpts.yaml \
--img-size 832 \
--conf-thres 0.001 \
--iou-thres 0.6 \
--task val \
--device 0 \
--kpt-label
INT8校准
python models/export_TRT.py \
--onnx weights/yolov7-w6-pose.onnx \