-
Notifications
You must be signed in to change notification settings - Fork 111
/
yolov3_darknet_main.cpp
105 lines (96 loc) · 3.39 KB
/
yolov3_darknet_main.cpp
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
// Created by luozhiwang ([email protected])
// Date: 2021/1/26
#include "yolov3_darknet.h"
void initInputParams(common::InputParams &inputParams){
inputParams.ImgH = 416;
inputParams.ImgW = 416;
inputParams.ImgC = 3;
inputParams.BatchSize = 1;
inputParams.HWC = false;
inputParams.IsPadding = false;
inputParams.InputTensorNames = std::vector<std::string>{"input"};
inputParams.OutputTensorNames = std::vector<std::string>{"boxes", "confs"};
inputParams.pFunction = [](const unsigned char &x){return static_cast<float>(x) /255;};
}
void initTrtParams(common::TrtParams &trtParams){
trtParams.ExtraWorkSpace = 0;
trtParams.FP32 = true;
trtParams.FP16 = false;
trtParams.Int32 = false;
trtParams.Int8 = false;
trtParams.worker = 4;
trtParams.MaxBatch = 100;
trtParams.MinTimingIteration = 1;
trtParams.AvgTimingIteration = 2;
trtParams.CalibrationTablePath = "/work/tensorRT-7/data/darknetInt8.calibration";
trtParams.CalibrationImageDir = "/data/dataset/coco/images/train2017";
trtParams.OnnxPath = "/work/tensorRT-7/data/onnx/darknet.onnx";
trtParams.SerializedPath = "/work/tensorRT-7/data/onnx/darknet.serialize";
}
std::vector<common::Anchor> initAnchors(){
std::vector<common::Anchor> anchors;
common::Anchor anchor;
// 10,13, 16,30, 33,23, 30,61, 62,45, 59,119, 116,90, 156,198, 373,326
anchor.width = 10;
anchor.height = 13;
anchors.emplace_back(anchor);
anchor.width = 16;
anchor.height = 30;
anchors.emplace_back(anchor);
anchor.width = 32;
anchor.height = 23;
anchors.emplace_back(anchor);
anchor.width = 30;
anchor.height = 61;
anchors.emplace_back(anchor);
anchor.width = 62;
anchor.height = 45;
anchors.emplace_back(anchor);
anchor.width = 59;
anchor.height = 119;
anchors.emplace_back(anchor);
anchor.width = 116;
anchor.height = 90;
anchors.emplace_back(anchor);
anchor.width = 156;
anchor.height = 198;
anchors.emplace_back(anchor);
anchor.width = 373;
anchor.height = 326;
anchors.emplace_back(anchor);
return anchors;
}
void initDetectParams(common::DetectParams &yoloParams){
yoloParams.Strides = std::vector<int> {8, 16, 32};
yoloParams.Anchors = initAnchors();
yoloParams.AnchorPerScale = 3;
yoloParams.NumClass = 80;
yoloParams.NMSThreshold = 0.5;
yoloParams.PostThreshold = 0.6;
}
int main(int args, char **argv){
// 46ms ===> 35ms worker=4
// 21.7fps ===> 28.5fps
common::InputParams inputParams;
common::TrtParams trtParams;
common::DetectParams yoloParams;
initInputParams(inputParams);
initTrtParams(trtParams);
initDetectParams(yoloParams);
Darknet yolo(inputParams, trtParams, yoloParams);
yolo.initSession(0);
cv::Mat image = cv::imread("/work/tensorRT-7/data/image/coco_1.jpg");
std::vector<common::Bbox> bboxes;
for(int i=0; i<20; ++i){
const auto start_t = std::chrono::high_resolution_clock::now();
bboxes = yolo.predOneImage(image);
const auto end_t = std::chrono::high_resolution_clock::now();
std::cout
<< "Wall clock time passed: "
<< std::chrono::duration<double, std::milli>(end_t-start_t).count()<<"ms"
<< std::endl;
}
image = renderBoundingBox(image, bboxes);
cv::imwrite("/work/tensorRT-7/data/image/render.jpg", image);
return 0;
}