此项目为关键点检测训练以及推理加速代码。训练部分用python3 + tensorflow-1.14完成,推理部分用C++ + tensorRT-6完成。
训练数据集主要为COCO,模型为Hourglass。
- python3 train_hourglass_coco.py
- python3 core/infer/freeze_graph.py -CUDA 0 -c checkpoints/coco/Hourglass_coco.ckpt -o Hourglass.pb
- python3 demon.py
https://drive.google.com/drive/folders/1pjOH1XUQOuMXlfGddQPvVEjXaXXPU7u1?usp=sharing
如果需要使用自己的数据集进行训练,首先需要将数据转换成如下的格式
(filename1 bxmin,bymin,bxmax,bymax px,py px,py ...)
If multi points have same label
(filename1 bxmin,bymin,bxmax,bymax px,py|px,py px,py ...)
(filename2 bxmin,bymin,bxmax,bymax px,py|px,py px,py|px,py ...)
...
在core/infer/infer_utils.py中的一些api可以用来构建一个简单的inference模型。通过Flask包装一下就可以实现简单的线上推理了。操作示例在infer_hourglass.py中,其中bbx需要通过其他模型获取。
TensorRT部分已经转移到新的仓库下
https://github.com/Syencil/tensorRT
此处项目采用CUDA 10 + tensorRT-6完成推理阶段,可实现模型推理加速,支持FP32,FP16
- 1.pb转uff
- cd tensorRT/python
- python3 pb2uff.py
- 2.编译C++文件
- cd tensorRT/c++
- cmake .
- make
1.数据增强 主要是图像旋转增强这一块有问题,会尽快将包括其他的增强方式加入项目2.TensorRT C++中对upsample plugin的实现,框架现已搭好,会尽快更新3.通过Hourglass-101构建今年大火的Anchor-free检测器CenterNet:Object as point4.tensorRT C++数据预处理和python有点不同,并不影响太多,懒得改了。5.Int 8量化矫正,有空再更新