Skip to content

Keypoints-detection in tensorflow and tensorRT C++

Notifications You must be signed in to change notification settings

Syencil/Keypoints

Repository files navigation

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量化矫正,有空再更新

About

Keypoints-detection in tensorflow and tensorRT C++

Topics

Resources

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published