1+ #include < fstream>
2+ #include < map>
3+ #include < sstream>
4+ #include < vector>
5+ #include < opencv2/opencv.hpp>
6+
7+ #include " NvInfer.h"
8+ #include " yololayer.h"
9+ #include " mish.h"
10+
11+
12+ using namespace nvinfer1 ;
13+
14+ cv::Mat preprocess_img (cv::Mat& img) {
15+ int w, h, x, y;
16+ float r_w = Yolo::INPUT_W / (img.cols *1.0 );
17+ float r_h = Yolo::INPUT_H / (img.rows *1.0 );
18+ if (r_h > r_w) {
19+ w = Yolo::INPUT_W;
20+ h = r_w * img.rows ;
21+ x = 0 ;
22+ y = (Yolo::INPUT_H - h) / 2 ;
23+ } else {
24+ w = r_h* img.cols ;
25+ h = Yolo::INPUT_H;
26+ x = (Yolo::INPUT_W - w) / 2 ;
27+ y = 0 ;
28+ }
29+ cv::Mat re (h, w, CV_8UC3);
30+ cv::resize (img, re, re.size ());
31+ cv::Mat out (Yolo::INPUT_H, Yolo::INPUT_W, CV_8UC3, cv::Scalar (128 , 128 , 128 ));
32+ re.copyTo (out (cv::Rect (x, y, re.cols , re.rows )));
33+ return out;
34+ }
35+
36+ cv::Rect get_rect (cv::Mat& img, float bbox[4 ]) {
37+ int l, r, t, b;
38+ float r_w = Yolo::INPUT_W / (img.cols * 1.0 );
39+ float r_h = Yolo::INPUT_H / (img.rows * 1.0 );
40+ if (r_h > r_w) {
41+ l = bbox[0 ] - bbox[2 ]/2 .f ;
42+ r = bbox[0 ] + bbox[2 ]/2 .f ;
43+ t = bbox[1 ] - bbox[3 ]/2 .f - (Yolo::INPUT_H - r_w * img.rows ) / 2 ;
44+ b = bbox[1 ] + bbox[3 ]/2 .f - (Yolo::INPUT_H - r_w * img.rows ) / 2 ;
45+ l = l / r_w;
46+ r = r / r_w;
47+ t = t / r_w;
48+ b = b / r_w;
49+ } else {
50+ l = bbox[0 ] - bbox[2 ]/2 .f - (Yolo::INPUT_W - r_h * img.cols ) / 2 ;
51+ r = bbox[0 ] + bbox[2 ]/2 .f - (Yolo::INPUT_W - r_h * img.cols ) / 2 ;
52+ t = bbox[1 ] - bbox[3 ]/2 .f ;
53+ b = bbox[1 ] + bbox[3 ]/2 .f ;
54+ l = l / r_h;
55+ r = r / r_h;
56+ t = t / r_h;
57+ b = b / r_h;
58+ }
59+ return cv::Rect (l, t, r-l, b-t);
60+ }
61+
62+ float iou (float lbox[4 ], float rbox[4 ]) {
63+ float interBox[] = {
64+ std::max (lbox[0 ] - lbox[2 ]/2 .f , rbox[0 ] - rbox[2 ]/2 .f ), // left
65+ std::min (lbox[0 ] + lbox[2 ]/2 .f , rbox[0 ] + rbox[2 ]/2 .f ), // right
66+ std::max (lbox[1 ] - lbox[3 ]/2 .f , rbox[1 ] - rbox[3 ]/2 .f ), // top
67+ std::min (lbox[1 ] + lbox[3 ]/2 .f , rbox[1 ] + rbox[3 ]/2 .f ), // bottom
68+ };
69+
70+ if (interBox[2 ] > interBox[3 ] || interBox[0 ] > interBox[1 ])
71+ return 0 .0f ;
72+
73+ float interBoxS =(interBox[1 ]-interBox[0 ])*(interBox[3 ]-interBox[2 ]);
74+ return interBoxS/(lbox[2 ]*lbox[3 ] + rbox[2 ]*rbox[3 ] -interBoxS);
75+ }
76+
77+ bool cmp (const Yolo::Detection& a, const Yolo::Detection& b) {
78+ return a.det_confidence > b.det_confidence ;
79+ }
80+
81+ void nms (std::vector<Yolo::Detection>& res, float *output, float conf_thresh, float nms_thresh = 0.5 ) {
82+ int det_size = sizeof (Yolo::Detection) / sizeof (float );
83+ std::map<float , std::vector<Yolo::Detection>> m;
84+ for (int i = 0 ; i < output[0 ] && i < Yolo::MAX_OUTPUT_BBOX_COUNT; i++) {
85+ if (output[1 + det_size * i + 4 ] <= conf_thresh) continue ;
86+ Yolo::Detection det;
87+ memcpy (&det, &output[1 + det_size * i], det_size * sizeof (float ));
88+ if (m.count (det.class_id ) == 0 ) m.emplace (det.class_id , std::vector<Yolo::Detection>());
89+ m[det.class_id ].push_back (det);
90+ }
91+ for (auto it = m.begin (); it != m.end (); it++) {
92+ // std::cout << it->second[0].class_id << " --- " << std::endl;
93+ auto & dets = it->second ;
94+ std::sort (dets.begin (), dets.end (), cmp);
95+ for (size_t m = 0 ; m < dets.size (); ++m) {
96+ auto & item = dets[m];
97+ res.push_back (item);
98+ for (size_t n = m + 1 ; n < dets.size (); ++n) {
99+ if (iou (item.bbox , dets[n].bbox ) > nms_thresh) {
100+ dets.erase (dets.begin ()+n);
101+ --n;
102+ }
103+ }
104+ }
105+ }
106+ }
107+
108+ // TensorRT weight files have a simple space delimited format:
109+ // [type] [size] <data x size in hex>
110+ std::map<std::string, Weights> loadWeights (const std::string file) {
111+ std::cout << " Loading weights: "
112+ std::map<std::string, Weights> weightMap;
113+
114+ // Open weights file
115+ std::ifstream input (file);
116+ assert (input.is_open () && " Unable to load weight file."
117+
118+ // Read number of weight blobs
119+ int32_t count;
120+ input >> count;
121+ assert (count > 0 && " Invalid weight map file."
122+
123+ while (count--)
124+ {
125+ Weights wt{DataType::kFLOAT , nullptr , 0 };
126+ uint32_t size;
127+
128+ // Read name and type of blob
129+ std::string name;
130+ input >> name >> std::dec >> size;
131+ wt.type = DataType::kFLOAT ;
132+
133+ // Load blob
134+ uint32_t * val = reinterpret_cast <uint32_t *>(malloc (sizeof (val) * size));
135+ for (uint32_t x = 0 , y = size; x < y; ++x)
136+ {
137+ input >> std::hex >> val[x];
138+ }
139+ wt.values = val;
140+
141+ wt.count = size;
142+ weightMap[name] = wt;
143+ }
144+
145+ return weightMap;
146+ }
147+
148+ IScaleLayer* addBatchNorm2d (INetworkDefinition *network, std::map<std::string, Weights>& weightMap, ITensor& input, std::string lname, float eps) {
149+ float *gamma = (float *)weightMap[lname + " .weight" values ;
150+ float *beta = (float *)weightMap[lname + " .bias" values ;
151+ float *mean = (float *)weightMap[lname + " .running_mean" values ;
152+ float *var = (float *)weightMap[lname + " .running_var" values ;
153+ int len = weightMap[lname + " .running_var" count ;
154+
155+ float *scval = reinterpret_cast <float *>(malloc (sizeof (float ) * len));
156+ for (int i = 0 ; i < len; i++) {
157+ scval[i] = gamma[i] / sqrt (var[i] + eps);
158+ }
159+ Weights scale{DataType::kFLOAT , scval, len};
160+
161+ float *shval = reinterpret_cast <float *>(malloc (sizeof (float ) * len));
162+ for (int i = 0 ; i < len; i++) {
163+ shval[i] = beta[i] - mean[i] * gamma[i] / sqrt (var[i] + eps);
164+ }
165+ Weights shift{DataType::kFLOAT , shval, len};
166+
167+ float *pval = reinterpret_cast <float *>(malloc (sizeof (float ) * len));
168+ for (int i = 0 ; i < len; i++) {
169+ pval[i] = 1.0 ;
170+ }
171+ Weights power{DataType::kFLOAT , pval, len};
172+
173+ weightMap[lname + " .scale"
174+ weightMap[lname + " .shift"
175+ weightMap[lname + " .power"
176+ IScaleLayer* scale_1 = network->addScale (input, ScaleMode::kCHANNEL , shift, scale, power);
177+ assert (scale_1);
178+ return scale_1;
179+ }
180+
181+ ILayer* convBnMish (INetworkDefinition *network, std::map<std::string, Weights>& weightMap, ITensor& input, int outch, int ksize, int s, int p, int linx) {
182+ Weights emptywts{DataType::kFLOAT , nullptr , 0 };
183+ IConvolutionLayer* conv1 = network->addConvolutionNd (input, outch, DimsHW{ksize, ksize}, weightMap[" module_list." std::to_string (linx) + " .Conv2d.weight"
184+ assert (conv1);
185+ conv1->setStrideNd (DimsHW{s, s});
186+ conv1->setPaddingNd (DimsHW{p, p});
187+
188+ IScaleLayer* bn1 = addBatchNorm2d (network, weightMap, *conv1->getOutput (0 ), " module_list." std::to_string (linx) + " .BatchNorm2d" 1e-4 );
189+
190+ auto creator = getPluginRegistry ()->getPluginCreator (" Mish_TRT" " 1"
191+ const PluginFieldCollection* pluginData = creator->getFieldNames ();
192+ IPluginV2 *pluginObj = creator->createPlugin ((" mish" std::to_string (linx)).c_str (), pluginData);
193+ ITensor* inputTensors[] = {bn1->getOutput (0 )};
194+ auto mish = network->addPluginV2 (&inputTensors[0 ], 1 , *pluginObj);
195+ return mish;
196+ }
0 commit comments