forked from shouxieai/tensorRT_Pro
-
Notifications
You must be signed in to change notification settings - Fork 1
Expand file tree
/
Copy pathapp_scrfd.cpp
More file actions
131 lines (99 loc) · 4.7 KB
/
app_scrfd.cpp
File metadata and controls
131 lines (99 loc) · 4.7 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
118
119
120
121
122
123
124
125
126
127
128
129
130
131
#include <builder/trt_builder.hpp>
#include <infer/trt_infer.hpp>
#include <common/ilogger.hpp>
#include "app_scrfd/scrfd.hpp"
using namespace std;
using namespace cv;
bool requires(const char* name);
bool compile_scrfd(int input_width, int input_height, string& out_model_file, TRT::Mode mode = TRT::Mode::FP32){
const char* name = "scrfd_2.5g_bnkps";
if(not requires(name))
return false;
string onnx_file = iLogger::format("%s.onnx", name);
string model_file = iLogger::format("%s.%dx%d.%s.trtmodel", name, input_width, input_height, TRT::mode_string(mode));
int test_batch_size = 6;
out_model_file = model_file;
if(iLogger::exists(model_file))
return true;
input_width = iLogger::upbound(input_width);
input_height = iLogger::upbound(input_height);
TRT::set_layer_hook_reshape([&](const string& name, const std::vector<int64_t>& shape){
INFOV("%s, %s", name.c_str(), iLogger::join_dims(shape).c_str());
vector<string> layerset{
"Reshape_108", "Reshape_110", "Reshape_112",
"Reshape_126", "Reshape_128", "Reshape_130",
"Reshape_144", "Reshape_146", "Reshape_148"
};
vector<int> strides{8, 8, 8, 16, 16, 16, 32, 32, 32};
auto layer_iter = std::find_if(layerset.begin(),layerset.end(), [&](const string& item){return item==name;});
if(layer_iter != layerset.end()){
int pos = layer_iter - layerset.begin();
int stride = strides[pos];
return vector<int64_t>{-1, input_height * input_width / stride / stride * 2, shape[2]};
}
return shape;
});
return TRT::compile(
TRT::Mode::FP32, // FP32、FP16、INT8
test_batch_size, // max batch size
onnx_file, // source
model_file, // save to
{TRT::InputDims({1, 3, input_height, input_width})}
);
}
static void scrfd_performance(shared_ptr<Scrfd::Infer> infer){
auto files = iLogger::find_files("inference", "*.jpg;*.jpeg;*.png;*.gif;*.tif");
vector<Mat> images;
for(int i = 0; i < files.size(); ++i)
images.emplace_back(imread(files[i]));
int ntest = 1000;
INFO("Do scrfd performance test..., %d testing, %d images, take a moment...", ntest, images.size());
// warmup
infer->commits(images).back().get();
auto begin = iLogger::timestamp_now_float();
for(int i = 0; i < ntest; ++i){
infer->commits(images).back().get();
}
auto time_fee_per_image = (iLogger::timestamp_now_float() - begin) / ntest / images.size();
INFO("***********************************************************************************************");
INFO("Retinaface limit case performance: %.3f ms / image, fps = %.2f", time_fee_per_image, 1000 / time_fee_per_image);
INFO("***********************************************************************************************");
}
int app_scrfd(){
TRT::set_device(0);
INFO("===================== test scrfd FP32 ==================================");
string model_file;
if(!compile_scrfd(640, 640, model_file))
return 0;
auto engine = Scrfd::create_infer(model_file, 0, 0.7);
scrfd_performance(engine);
auto save_root = "scrfd_result";
iLogger::rmtree(save_root);
iLogger::mkdirs(save_root);
auto files = iLogger::find_files("inference", "*.jpg;*.jpeg;*.png;*.gif;*.tif");
vector<Mat> images;
for(int i = 0; i < files.size(); ++i)
images.emplace_back(imread(files[i]));
// warmup
engine->commits(images).back().get();
auto time_begin = iLogger::timestamp_now_float();
auto images_faces = engine->commits(images);
images_faces.back().get();
float average_time = (iLogger::timestamp_now_float() - time_begin) / files.size();
for(int i = 0; i < images.size(); ++i){
auto& image = images[i];
auto faces = images_faces[i].get();
for(int i = 0; i < faces.size(); ++i){
auto& face = faces[i];
rectangle(image, cv::Point(face.left, face.top), cv::Point(face.right, face.bottom), Scalar(0, 255, 0), 2);
for(int j = 0; j < 5; ++j)
circle(image, Point(face.landmark[j*2+0], face.landmark[j*2+1]), 3, Scalar(0, 255, 0), -1, 16);
putText(image, iLogger::format("%.3f", face.confidence), cv::Point(face.left, face.top), 0, 1, Scalar(0, 255, 0), 1, 16);
}
auto save_file = iLogger::format("%s/%s", save_root, iLogger::file_name(files[i]).c_str());
INFO("Save to %s, %d faces, average time: %.2f ms", save_file.c_str(), faces.size(), average_time);
imwrite(save_file, image);
}
INFO("Done");
return 0;
}