Skip to content

Latest commit

 

History

History
51 lines (40 loc) · 2.07 KB

README.md

File metadata and controls

51 lines (40 loc) · 2.07 KB
Raw

Keypoint Detection In Tensorflow and TensorRT C++

1.Modified hourglass (Hourglass-104) and ResNet-101

Introduction

此项目为关键点检测训练以及推理加速代码。训练部分用python3 + tensorflow-1.14完成,推理部分用C++ + tensorRT-6完成。
训练数据集主要为COCO，模型为Hourglass。

Quick Start

  • 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

Checkpoints

https://drive.google.com/drive/folders/1pjOH1XUQOuMXlfGddQPvVEjXaXXPU7u1?usp=sharing

Data Format

如果需要使用自己的数据集进行训练，首先需要将数据转换成如下的格式
(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 ...)
...

Inference

在core/infer/infer_utils.py中的一些api可以用来构建一个简单的inference模型。通过Flask包装一下就可以实现简单的线上推理了。操作示例在infer_hourglass.py中，其中bbx需要通过其他模型获取。

注意事项

TensorRT部分已经转移到新的仓库下
https://github.com/Syencil/tensorRT

2.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 point
  • 4.tensorRT C++数据预处理和python有点不同，并不影响太多，懒得改了。
  • 5.Int 8量化矫正，有空再更新