forked from shouxieai/tensorRT_Pro
-
Notifications
You must be signed in to change notification settings - Fork 1
Expand file tree
/
Copy pathapp_alphapose.cpp
More file actions
71 lines (59 loc) · 2.22 KB
/
app_alphapose.cpp
File metadata and controls
71 lines (59 loc) · 2.22 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
/**
* @file _main.cpp
* @author 手写AI (zifuture.com:8090)
* @date 2021-07-26
*
* 实现了基于TensorRT对yolox的推理工作
* 1. 基于FP32的模型编译、和推理执行
* 2. 基于INT8的模型编译、和推理执行
* 3. 自定义插件的实现，从pytorch导出到推理编译，并支持FP16
*
* 预处理、后处理采用CPU实现（若想GPU可以自行实现）
* 一次推理5张图获取结果
*
* 我们是一群热血的个人组织者，力图发布免费高质量内容
* 我们的博客地址：http://zifuture.com:8090
* 我们的B站地址：https://space.bilibili.com/1413433465
*
* 如果想要深入学习关于tensorRT的技术栈，请通过博客中的二维码联系我们（免费崔更即可）
* 请关注B站，我们根据情况发布相关教程视频（免费）
*/
#include <builder/trt_builder.hpp>
#include <infer/trt_infer.hpp>
#include <common/ilogger.hpp>
#include "app_alphapose/alpha_pose.hpp"
using namespace std;
using namespace cv;
bool requires(const char* name);
int app_alphapose(){
TRT::set_device(0);
INFO("===================== test alphapose fp32 ==================================");
const char* name = "sppe";
if(not requires(name))
return 0;
string onnx_file = iLogger::format("%s.onnx", name);
string model_file = iLogger::format("%s.FP32.trtmodel", name);
int test_batch_size = 16;
if(!iLogger::exists(model_file)){
TRT::compile(
TRT::Mode::FP32, // FP32、FP16、INT8
test_batch_size, // max_batch_size
onnx_file, // source
model_file // save to
);
}
Mat image = imread("inference/gril.jpg");
auto engine = AlphaPose::create_infer(model_file, 0);
auto box = Rect(158, 104, 176, 693);
auto keys = engine->commit(make_tuple(image, box)).get();
for(int i = 0; i < keys.size(); ++i){
float x = keys[i].x;
float y = keys[i].y;
cv::circle(image, Point(x, y), 5, Scalar(0, 255, 0), -1, 16);
}
auto save_file = "pose.show.jpg";
INFO("Save to %s", save_file);
imwrite(save_file, image);
INFO("Done");
return 0;
}