forked from shouxieai/tensorRT_Pro
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathapp_high_performance.cpp
More file actions
117 lines (92 loc) · 3.3 KB
/
Copy pathapp_high_performance.cpp
File metadata and controls
117 lines (92 loc) · 3.3 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
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
#include "app_high_performance/high_performance.hpp"
#include "app_high_performance/yolo_high_perf.hpp"
#include "app_high_performance/alpha_pose_high_perf.hpp"
using namespace std;
using namespace HighPerformance;
class ImageData : public Data, public cv::Mat{
public:
SetupData(ImageData);
ImageData() = default;
ImageData(const cv::Mat& image):cv::Mat(image){}
};
class YoloNode : public Node{
public:
virtual bool startup(const string& file, YoloHighPerf::Type type, int gpuid = 0){
infer_ = YoloHighPerf::create_infer(file, type, gpuid);
if(infer_ == nullptr)
return false;
Node::startup();
return true;
}
protected:
virtual void forward(vector<shared_ptr<Data>>& inputs_data) override{
auto image = dynamic_obj_cast(inputs_data[0], ImageData);
auto output_future = infer_->commit(*image);
outputs_[0]->commit(output_future);
}
private:
shared_ptr<YoloHighPerf::Infer> infer_;
};
class PoseNode : public Node{
public:
virtual bool startup(const string& file, int gpuid = 0){
infer_ = AlphaPoseHighPerf::create_infer(file, gpuid);
if(infer_ == nullptr)
return false;
Node::startup();
return true;
}
protected:
virtual void forward(vector<shared_ptr<Data>>& inputs_data) override{
auto image = dynamic_obj_cast(inputs_data[0], cv::Mat);
auto boxarray = dynamic_obj_cast(inputs_data[1], YoloHighPerf::ObjectBoxArray);
vector<shared_future<DataPtr>> keypoints(boxarray->size());
for(int i = 0; i < boxarray->size(); ++i){
auto& box = boxarray->at(i);
keypoints[i] = infer_->commit(make_tuple(*image, cv::Rect(box.left, box.top, box.right-box.left, box.bottom-box.top)));
}
outputs_[0]->commit(
make_data_future(
make_shared<Container<tuple<shared_ptr<cv::Mat>, vector<shared_future<DataPtr>>>>>(image, keypoints)
)
);
}
private:
shared_ptr<AlphaPoseHighPerf::Infer> infer_;
};
int app_high_performance(){
InputNode camera;
OutputNode output;
YoloNode yolo;
PoseNode pose;
connect(camera, yolo);
connect(camera, pose);
connect(yolo, pose);
connect(pose, output);
camera.startup([](vector<shared_ptr<Pipeline>>& output_pipe){
cv::VideoCapture cap("exp/fall_video.mp4");
cv::Mat image;
int num_frame = 0;
while(cap.read(image)){
num_frame++;
for(int i = 0; i < output_pipe.size(); ++i)
output_pipe[i]->commit(make_data_future(make_shared<ImageData>(image.clone())));
}
INFO("num frame %d", num_frame);
});
output.startup([](vector<shared_ptr<Data>>& datas){
using ThisData = Container<tuple<shared_ptr<cv::Mat>, vector<shared_future<DataPtr>>>>;
auto data = dynamic_obj_cast(datas[0], ThisData);
auto& image = *get<0>(data->value);
auto& points = get<1>(data->value);
for(auto& keys_fut : points){
auto keys = keys_fut.get();
auto keys_kpt = dynamic_obj_cast(keys, AlphaPoseHighPerf::PointArray);
}
});
yolo.startup("yolox_s.FP32.trtmodel", YoloHighPerf::Type::X);
pose.startup("sppe.fp32.trtmodel");
getchar();
camera.stop();
return 0;
}