forked from shouxieai/tensorRT_Pro
-
Notifications
You must be signed in to change notification settings - Fork 1
Expand file tree
/
Copy pathapp_centernet.cpp
More file actions
113 lines (86 loc) · 4.04 KB
/
Copy pathapp_centernet.cpp
File metadata and controls
113 lines (86 loc) · 4.04 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
#include <builder/trt_builder.hpp>
#include <infer/trt_infer.hpp>
#include <common/ilogger.hpp>
#include "app_centernet/centernet.hpp"
using namespace std;
static const char* cocolabels[] = {
"person", "bicycle", "car", "motorcycle", "airplane",
"bus", "train", "truck", "boat", "traffic light", "fire hydrant",
"stop sign", "parking meter", "bench", "bird", "cat", "dog", "horse",
"sheep", "cow", "elephant", "bear", "zebra", "giraffe", "backpack",
"umbrella", "handbag", "tie", "suitcase", "frisbee", "skis",
"snowboard", "sports ball", "kite", "baseball bat", "baseball glove",
"skateboard", "surfboard", "tennis racket", "bottle", "wine glass",
"cup", "fork", "knife", "spoon", "bowl", "banana", "apple", "sandwich",
"orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake",
"chair", "couch", "potted plant", "bed", "dining table", "toilet", "tv",
"laptop", "mouse", "remote", "keyboard", "cell phone", "microwave",
"oven", "toaster", "sink", "refrigerator", "book", "clock", "vase",
"scissors", "teddy bear", "hair drier", "toothbrush"
};
bool requires(const char* name);
static void forward_engine(const string& engine_file){
auto engine = CenterNet::create_infer(engine_file, 0, 0.35f);
if(engine == nullptr){
INFOE("Engine is nullptr");
return;
}
string root = "CenterNet_result";
iLogger::rmtree(root);
iLogger::mkdir(root);
auto files = iLogger::find_files("inference", "*.jpg;*.jpeg;*.png;*.gif;*.tif");
vector<cv::Mat> images;
for(int i = 0; i < files.size(); ++i){
auto image = cv::imread(files[i]);
images.emplace_back(image);
}
// warmup
engine->commit(images[0]).get();
auto t0 = iLogger::timestamp_now_float();
auto boxes_array = engine->commits(images);
// wait batch result
boxes_array.back().get();
float inference_average_time = (iLogger::timestamp_now_float() - t0) / boxes_array.size();
for(int i = 0; i < boxes_array.size(); ++i){
auto& image = images[i];
auto boxes = boxes_array[i].get();
for(auto& obj : boxes){
uint8_t b, g, r;
tie(b, g, r) = iLogger::random_color(obj.class_label);
cv::rectangle(image, cv::Point(obj.left, obj.top), cv::Point(obj.right, obj.bottom), cv::Scalar(b, g, r), 5);
auto name = cocolabels[obj.class_label];
auto caption = iLogger::format("%s %.2f", name, obj.confidence);
int width = cv::getTextSize(caption, 0, 1, 1, nullptr).width;
cv::rectangle(image, cv::Point(obj.left-3, obj.top-33), cv::Point(obj.left + width-30, obj.top), cv::Scalar(b, g, r), -1);
cv::putText(image, caption, cv::Point(obj.left, obj.top-5), 0, 0.8, cv::Scalar::all(0), 2, 16);
}
string file_name = iLogger::file_name(files[i], false);
string save_path = iLogger::format("%s/%s.jpg", root.c_str(), file_name.c_str());
INFO("Save to %s, %d object, average time %.2f ms", save_path.c_str(), boxes.size(), inference_average_time);
cv::imwrite(save_path, image);
}
}
static void test(TRT::Mode mode, const string& model){
auto mode_name = TRT::mode_string(mode);
TRT::set_device(0);
const char* name = model.c_str();
INFO("===================== test %s %s ==================================", mode_name, name);
if(not requires(name))
return;
string onnx_file = iLogger::format("%s.onnx", name);
string model_file = iLogger::format("%s.%s.trtmodel", name, mode_name);
int test_batch_size = 16;
if(not iLogger::exists(model_file)){
TRT::compile(
mode, // FP32
test_batch_size, // max batch size
onnx_file, // source
model_file // save to
);
}
forward_engine(model_file);
}
int app_centernet(){
test(TRT::Mode::FP32, "centernet_r18_dcn"); // Type indicates the task. name string indicates the model size and structure.
return 0;
}