├── .gitignore
├── CMakeLists.txt
├── README.md
├── include
    ├── centerTrack.h
    ├── tracker.h
    └── util.h
└── src
    ├── centerTrack.cpp
    ├── main.cpp
    ├── tracker.cpp
    └── util.cpp


/.gitignore:
--------------------------------------------------------------------------------
1 | /build/
2 | /output_video/
3 | /videos/
4 | /model/
5 | .vscode
6 | 


--------------------------------------------------------------------------------
/CMakeLists.txt:
--------------------------------------------------------------------------------
 1 | cmake_minimum_required(VERSION 3.0 FATAL_ERROR)
 2 | project(example-app)
 3 | 
 4 | # set(CMAKE_PREFIX_PATH /home/lx/programs/libtorch/share/cmake/Torch)
 5 | # find_package(Torch REQUIRED)
 6 | find_package(OpenCV REQUIRED)
 7 | set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${TORCH_CXX_FLAGS}")
 8 | include_directories("./include" 
 9 | "/home/lx/work/hi3519/vehicle_detect/caffe_ssd_head_tail/caffe_ssd/include"
10 | "/usr/local/cuda-10.0/include")
11 | 
12 | aux_source_directory(./src SRC)
13 | add_executable(example-app ${SRC})
14 | target_link_libraries(example-app "${TORCH_LIBRARIES}" "${OpenCV_LIBS}"
15 | "/home/lx/work/hi3519/vehicle_detect/caffe_ssd_head_tail/caffe_ssd/build/lib/libcaffe.so"
16 | "/usr/lib/x86_64-linux-gnu/libglog.so"
17 | "/usr/lib/x86_64-linux-gnu/libboost_system.so"
18 | "/usr/local/cuda/lib64/libcudart.so"
19 | "/usr/local/cuda-10.0/lib64/libcublas.so")
20 | set_property(TARGET example-app PROPERTY CXX_STANDARD 14)


--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
 1 | # CenterTrack 2D tracking C++ caffe 版本
 2 | 
 3 | ## 根据[**CenterTrack官方pytorch版本**](https://github.com/xingyizhou/CenterTrack)改写.
 4 | > 话说CenteTrack是真的丝滑,检测加跟踪一步到位,检测效果比yolov3还好,跟踪效果跟deepsort有的一拼,屏内跟踪的不二之选.
 5 | 
 6 | ## 使用[**此版本的CenterTrack**](https://github.com/lrjbdss/CenterTrack_2D_train)进行训练,然后使用[**模型转换工具**](https://github.com/xxradon/PytorchToCaffe)转成caffemodel
 7 | > 有坑,caffe的pooling只有ceil_mode模式,但pytorch默认ceil_model为False,转模型会报错.屏蔽掉ceil_model关键字,打开ceil_mode模式即可,因为模型是基于floor_model模式训练的,输出尺寸跟训练时不同,后处理的坐标转换有点麻烦.  
 8 | >另外,考虑到简化处理,mns阶段的sigmoid和maxpooling被放到了模型中,但caffe模型不支持CenterTrack的nms操作,而且量化后的海思wk模型也不能正常的maxpooling(pooling后最大像素值也会改变),所以模型输出的是sigmoid层,后续的nms由cpu完成.
 9 | 
10 | ## 主干网络使用resnet18 
11 | > resnest18也试了,没感觉有明显提升


--------------------------------------------------------------------------------
/include/centerTrack.h:
--------------------------------------------------------------------------------
 1 | #pragma once
 2 | #include "caffe/caffe.hpp"
 3 | // #include "util.h"
 4 | #include "tracker.h"
 5 | #include <string>
 6 | #include <vector>
 7 | 
 8 | class CenterTrack
 9 | {
10 | public:
11 |     CenterTrack(const std::string prototxtPath,
12 |                 const std::string caffeModelPath,
13 |                 float score_threshold = 0.3,
14 |                 bool input_pre_hm = false,
15 |                 int num_class = 1,
16 |                 int topK = 25,
17 |                 // std::array<float, 3> mean = {105.9422484375, 105.21612951388889, 105.52638559027778},
18 |                 // float std = 78.21761860831944);
19 |                 std::array<float, 3> mean = {0.40789655, 0.44719303, 0.47026116},
20 |                 std::array<float, 3> std = {0.279, 0.279, 0.279});
21 | 
22 |     std::vector<Track> tracking(cv::Mat &img, cv::Mat pre_img);
23 | 
24 | private:
25 |     void WrapInputLayer(std::vector<cv::Mat> *input_channels,
26 |                         caffe::Blob<float> *input_layer);
27 | 
28 |     void Preprocess(const cv::Mat &img,
29 |                     std::vector<cv::Mat> &input_channels);
30 |     cv::Mat img_preprocess(cv::Mat img);
31 |     std::vector<Track> post_process();
32 | 
33 | private:
34 |     std::array<float, 3> img_mean;
35 |     std::array<float, 3> img_std;
36 |     cv::Mat pre_img;
37 |     cv::Mat pre_hm;
38 |     bool input_pre_hm;
39 |     std::shared_ptr<caffe::Net<float>> net_;
40 |     float score_threshold;
41 |     int inp_w, inp_h, ori_w, ori_h;
42 |     int num_class;
43 |     int topK;
44 |     Tracker tracker;
45 | };
46 | 


--------------------------------------------------------------------------------
/include/tracker.h:
--------------------------------------------------------------------------------
 1 | #pragma once
 2 | #include <opencv2/opencv.hpp>
 3 | #include <array>
 4 | #include <vector>
 5 | #include <map>
 6 | #include <math.h>
 7 | #include "util.h"
 8 | 
 9 | class Tracker
10 | {
11 | public:
12 |     Tracker();
13 | 
14 |     void reset();
15 | 
16 |     std::vector<Track> step(std::vector<Track> results);
17 | 
18 |     cv::Mat gen_pre_hm(int w, int h);
19 | 
20 | private:
21 |     int id_count;
22 |     std::vector<Track> tracks;
23 | 
24 |     // void add_track(Track det);
25 |     float cal_dist(cv::Point2f p1, cv::Point2f p2);
26 | };


--------------------------------------------------------------------------------
/include/util.h:
--------------------------------------------------------------------------------
 1 | #pragma once
 2 | #include <opencv2/opencv.hpp>
 3 | #include <vector>
 4 | #include <numeric>
 5 | 
 6 | struct Track
 7 | {
 8 |     float score;
 9 |     float class_id;
10 |     cv::Point2f center;
11 |     // std::array<float, 2> tracking;
12 |     cv::Vec2f tracking;
13 |     cv::Rect2f bbox;
14 |     int track_id;
15 |     int active;
16 |     // void update(Track det)
17 |     // {
18 |     //     score = det.score;
19 |     //     center = det.center;
20 |     //     tracking = det.tracking;
21 |     //     bbox = det.bbox;
22 |     //     active += 1;
23 |     // }
24 | };
25 | 
26 | cv::Scalar color_map(int64_t n);
27 | // {
28 | //     auto bit_get = [](int64_t x, int64_t i) { return x & (1 << i); };
29 | 
30 | //     int64_t r = 0, g = 0, b = 0;
31 | //     int64_t i = n;
32 | //     for (int64_t j = 7; j >= 0; --j)
33 | //     {
34 | //         r |= bit_get(i, 0) << j;
35 | //         g |= bit_get(i, 1) << j;
36 | //         b |= bit_get(i, 2) << j;
37 | //         i >>= 3;
38 | //     }
39 | //     return cv::Scalar(b, g, r);
40 | // }
41 | 
42 | void draw_text(cv::Mat &img, const std::string &str,
43 |                const cv::Scalar &color, cv::Point pos);
44 | // {
45 | //     auto t_size = cv::getTextSize(str, cv::FONT_HERSHEY_PLAIN, 1, 1, nullptr);
46 | //     cv::Point bottom_left, upper_right;
47 | 
48 | //     bottom_left = pos;
49 | //     upper_right = cv::Point(bottom_left.x + t_size.width, bottom_left.y - t_size.height);
50 | 
51 | //     cv::rectangle(img, bottom_left, upper_right, color, -1);
52 | //     cv::putText(img, str, bottom_left, cv::FONT_HERSHEY_PLAIN, 1, cv::Scalar(255, 255, 255) - color);
53 | // }
54 | 
55 | void draw_result(cv::Mat &img, Track track);
56 | // {
57 | //     auto img_box = cv::Rect2f(bbox.x * img.cols,
58 | //                               bbox.y * img.rows,
59 | //                               bbox.width * img.cols,
60 | //                               bbox.height * img.rows);
61 | 
62 | //     cv::rectangle(img, img_box, color);
63 | 
64 | //     if (!label.empty())
65 | //     {
66 | //         draw_text(img, label, color, img_box.tl());
67 | //     }
68 | // }
69 | 
70 | template <typename T>
71 | std::vector<int> sort_index(const std::vector<T> &v)
72 | {
73 |     std::vector<int> idx(v.size());
74 |     std::iota(idx.begin(), idx.end(), 0);
75 |     std::sort(idx.begin(), idx.end(), [&v](int i1, int i2) { return v[i1] > v[i2]; });
76 |     return idx;
77 | }
78 | 
79 | void draw_gaussian(cv::Mat &input_hm, cv::Rect2f bbox);


--------------------------------------------------------------------------------
/src/centerTrack.cpp:
--------------------------------------------------------------------------------
  1 | #include "centerTrack.h"
  2 | #include "util.h"
  3 | #include <iostream>
  4 | #include <array>
  5 | // #include "caffe/caffe.hpp"
  6 | 
  7 | using boost::shared_ptr;
  8 | using caffe::Blob;
  9 | using caffe::Caffe;
 10 | using caffe::Layer;
 11 | using caffe::Net;
 12 | 
 13 | CenterTrack::CenterTrack(const std::string prototxtPath,
 14 |                          const std::string caffeModelPath,
 15 |                          float score,
 16 |                          bool add_pre_hm,
 17 |                          int classes,
 18 |                          int K,
 19 |                          std::array<float, 3> mean,
 20 |                          std::array<float, 3> std)
 21 |     : net_(new Net<float>(prototxtPath, caffe::Phase::TEST)),
 22 |       score_threshold(score),
 23 |       num_class(classes),
 24 |       topK(K),
 25 |       input_pre_hm(add_pre_hm),
 26 |       img_mean(mean),
 27 |       img_std(std),
 28 |       tracker()
 29 | {
 30 |     Caffe::set_mode(Caffe::GPU);
 31 |     net_->CopyTrainedLayersFrom(caffeModelPath);
 32 | 
 33 |     int num_inputs = 2;
 34 |     if (input_pre_hm)
 35 |         num_inputs = 3;
 36 |     CHECK_EQ(net_->num_inputs(), num_inputs) << "Network should have exactly %d input.", num_inputs;
 37 |     CHECK_EQ(net_->num_outputs(), 4) << "Network should have exactly 4 output.";
 38 |     Blob<float> *input_layer = net_->input_blobs()[0];
 39 |     inp_w = input_layer->width();
 40 |     inp_h = input_layer->height();
 41 |     // pre_hm = tracker.gen_pre_hm(inp_w, inp_h);
 42 |     // input_geometry_ = cv::Size(input_layer->width(), input_layer->height());
 43 | }
 44 | 
 45 | cv::Mat CenterTrack::img_preprocess(cv::Mat img)
 46 | {
 47 |     if (img.size() != cv::Size(inp_w, inp_h))
 48 |     {
 49 |         ori_h = img.rows;
 50 |         ori_w = img.cols;
 51 | 
 52 |         cv::Point2f ori[3], inp[3];
 53 |         ori[0] = cv::Point2f(ori_w / 2., ori_h / 2.);
 54 |         inp[0] = cv::Point2f(inp_w / 2., inp_h / 2.);
 55 |         if (ori_w / float(ori_h) >= inp_w / float(inp_h))
 56 |         {
 57 |             float delta_h = (inp_h / float(inp_w) * ori_w - ori_h) / 2.;
 58 |             ori[1] = cv::Point2f(0., ori_h / 2.);
 59 |             inp[1] = cv::Point2f(0., inp_h / 2.);
 60 |             ori[2] = cv::Point2f(0., -delta_h);
 61 |             inp[2] = cv::Point2f(0., 0.);
 62 |         }
 63 |         else
 64 |         {
 65 |             float delta_w = (inp_w / float(inp_h) * ori_h - ori_w) / 2.;
 66 |             ori[1] = cv::Point2f(ori_w / 2., 0);
 67 |             inp[1] = cv::Point2f(inp_w / 2., 0);
 68 |             ori[2] = cv::Point2f(-delta_w, 0);
 69 |             inp[2] = cv::Point2f(0., 0.);
 70 |         }
 71 |         cv::Mat trans = cv::getAffineTransform(ori, inp);
 72 |         cv::Mat resized_img;
 73 |         cv::warpAffine(img, resized_img, trans, cv::Size2d(inp_w, inp_h));
 74 |         return resized_img;
 75 |     }
 76 |     return img;
 77 | }
 78 | 
 79 | std::vector<Track> CenterTrack::tracking(cv::Mat &img, cv::Mat input_pre_img)
 80 | {
 81 | 
 82 |     cv::Mat resized_img = img_preprocess(img);
 83 | 
 84 |     // cv::imshow("resized_img", resized_img);
 85 |     // cv::waitKey(0);
 86 | 
 87 |     std::vector<cv::Mat> img_channels;
 88 |     WrapInputLayer(&img_channels, net_->input_blobs()[0]);
 89 |     Preprocess(resized_img, img_channels);
 90 |     printf("***img_channels[0].at<float>(0, 0): %f\n", img_channels[0].at<float>(0, 0));
 91 |     printf("***resized_img.data[0, 0, 0]: %d, resized: %f\n",
 92 |            resized_img.data[0, 0, 0], (resized_img.data[0, 0, 0] / 255. - img_mean[0]) / img_std[0]);
 93 |     for (int i = 0; i < img_channels.size(); ++i)
 94 |         img_channels[i] = (img_channels[i] / 255. - img_mean[i]) / img_std[i];
 95 |     // img_channels[i] = (img_channels[i] - img_mean[i]) / img_std;
 96 |     printf("***img_channels[0].at<float>(0, 0): %f\n", img_channels[0].at<float>(0, 0));
 97 |     if (input_pre_img.data)
 98 |         pre_img = img_preprocess(input_pre_img);
 99 |     if (!pre_img.data)
100 |         pre_img = resized_img;
101 |     std::vector<cv::Mat> pre_img_channels;
102 | 
103 |     WrapInputLayer(&pre_img_channels, net_->input_blobs()[1]);
104 |     Preprocess(pre_img, pre_img_channels);
105 |     for (int i = 0; i < pre_img_channels.size(); ++i)
106 |         pre_img_channels[i] = (pre_img_channels[i] / 255. - img_mean[i]) / img_std[i];
107 |     // pre_img_channels[i] = (pre_img_channels[i] - img_mean[i]) / img_std;
108 |     pre_img = resized_img;
109 | 
110 |     if (input_pre_hm)
111 |     {
112 |         std::vector<cv::Mat> pre_hm_channels;
113 |         WrapInputLayer(&pre_hm_channels, net_->input_blobs()[2]);
114 |         pre_hm = img_preprocess(pre_hm);
115 |         Preprocess(pre_hm, pre_hm_channels);
116 |     }
117 | 
118 |     net_->Forward();
119 | 
120 |     // *** test
121 |     // Blob<float> *add_blob10 = net_->blob_by_name("add_blob10").get();
122 |     // const float *add_blob10_data = add_blob10->cpu_data();
123 | 
124 |     std::vector<Track> dets = post_process();
125 |     auto results = tracker.step(dets);
126 |     // pre_hm = tracker.gen_pre_hm(ori_w, ori_h);
127 |     return results;
128 | }
129 | 
130 | std::vector<Track> CenterTrack::post_process()
131 | {
132 |     Blob<float> *hm = net_->blob_by_name("sigmoid_blob1").get();
133 |     float *hm_dat = hm->mutable_cpu_data();
134 |     // Blob<float> *hm_pool = net_->blob_by_name("max_pool_blob2").get();
135 |     // const float *hm_pool_data = hm_pool->cpu_data();
136 |     // int hm_w = int(inp_w / 4);
137 |     // int hm_h = int(inp_h / 4);
138 |     int hm_w = 88; // 本来应该320下采样4倍变成80,因为caffe的pooling只有ceil_mode,
139 |     int hm_h = 48; // 结果变成了88,与训练时有差异
140 |     CHECK_EQ(hm->channels(), num_class) << "output hm channel should equal to num_class.";
141 |     CHECK_EQ(hm->width(), hm_w) << "output hm width should have equal to input_width/4.";
142 |     CHECK_EQ(hm->height(), hm_h) << "output hm height should have equal to input_height/4.";
143 | 
144 |     // nms
145 |     int row, col;
146 |     float hm_data[num_class * hm_w * hm_h] = {0};
147 |     for (int i = 0; i < num_class * hm_w * hm_h; ++i)
148 |     {
149 |         // if (hm_data[i] != hm_pool_data[i])
150 |         //     hm_data[i] = 0.;
151 |         row = i / hm_w % hm_h;
152 |         col = i % hm_w;
153 |         if (col > 0)
154 |         {
155 |             // left
156 |             if (hm_dat[i] < hm_dat[i - 1])
157 |                 continue;
158 |         }
159 |         if (col < hm_w - 1)
160 |         {
161 |             // right
162 |             if (hm_dat[i] <= hm_dat[i + 1])
163 |                 continue;
164 |         }
165 | 
166 |         if (row > 0)
167 |         {
168 |             // up
169 |             if (hm_dat[i] < hm_dat[i - hm_w])
170 |                 continue;
171 | 
172 |             // up-left
173 |             if (col > 0)
174 |             {
175 |                 if (hm_dat[i] < hm_dat[i - hm_w - 1])
176 |                     continue;
177 |             }
178 |             // up-right
179 |             if (col < hm_h - 1)
180 |             {
181 |                 if (hm_dat[i] <= hm_dat[i - hm_w + 1])
182 |                     continue;
183 |             }
184 |         }
185 | 
186 |         if (row < hm_h - 1)
187 |         {
188 |             // down
189 |             if (hm_dat[i] <= hm_dat[i + hm_w])
190 |                 continue;
191 | 
192 |             // down-left
193 |             if (col > 0)
194 |             {
195 |                 if (hm_dat[i] < hm_dat[i + hm_w - 1])
196 |                     continue;
197 |             }
198 |             // down-right
199 |             if (col < hm_w - 1)
200 |             {
201 |                 if (hm_dat[i] <= hm_dat[i + hm_w + 1])
202 |                     continue;
203 |             }
204 |         }
205 |         hm_data[i] = hm_dat[i];
206 |     }
207 | 
208 |     const std::vector<float> hm_vec(hm_data, hm_data + num_class * hm_w * hm_h);
209 |     std::vector<int> max_index = sort_index(hm_vec);
210 | 
211 |     int channel_size = hm_w * hm_h;
212 | 
213 |     Blob<float> *reg = net_->blob_by_name("conv_blob25").get();
214 |     const float *reg_data = reg->cpu_data();
215 | 
216 |     Blob<float> *wh = net_->blob_by_name("conv_blob27").get();
217 |     const float *wh_data = wh->cpu_data();
218 | 
219 |     Blob<float> *track = net_->blob_by_name("conv_blob29").get();
220 |     const float *track_data = track->cpu_data();
221 | 
222 |     // 下面三个变量用来坐标转换
223 |     float scale = std::max(ori_w / 80., ori_h / 45.);
224 |     // float delta_w = 0.;
225 |     // float delta_h = 95.;
226 |     float delta_w = (scale * 80 - ori_w) / 2.;
227 |     float delta_h = (scale * 45 - ori_h) / 2.;
228 | 
229 |     std::vector<float> topK_score;
230 |     std::vector<float> topK_x, topK_y;
231 |     std::vector<int> topK_class;
232 |     std::vector<float> reg_x, reg_y;
233 |     std::vector<float> wh_x, wh_y;
234 |     std::vector<float> track_x, track_y;
235 |     std::vector<cv::Rect2f> topK_bbox;
236 |     int obj_count = 0;
237 |     for (; obj_count < topK; ++obj_count)
238 |     {
239 |         int idx = max_index[obj_count];
240 |         if (hm_vec[idx] < score_threshold)
241 |             break;
242 |         topK_score.push_back(hm_vec[idx]);
243 |         topK_class.push_back(int(idx / channel_size));
244 |         idx %= channel_size;
245 |         topK_y.push_back(idx / hm_w * scale - delta_h);
246 |         topK_x.push_back(idx % hm_w * scale - delta_w);
247 |         reg_x.push_back(reg_data[idx] * scale);
248 |         reg_y.push_back(reg_data[idx + channel_size] * scale);
249 |         wh_x.push_back(wh_data[idx] * scale);
250 |         wh_y.push_back(wh_data[idx + channel_size] * scale);
251 |         track_x.push_back(track_data[idx] * scale);
252 |         track_y.push_back(track_data[idx + channel_size] * scale);
253 |         topK_bbox.push_back(cv::Rect2f(topK_x[obj_count] + reg_x[obj_count] - wh_x[obj_count] / 2.,
254 |                                        topK_y[obj_count] + reg_y[obj_count] - wh_y[obj_count] / 2.,
255 |                                        wh_x[obj_count], wh_y[obj_count]));
256 |     }
257 |     // printf("obj_count: %d\n", obj_count);
258 |     std::vector<Track> dets;
259 |     for (int i = 0; i < obj_count; ++i)
260 |     {
261 |         Track item;
262 |         item.score = topK_score[i];
263 |         item.class_id = topK_class[i];
264 |         item.center = cv::Point2i(topK_x[i], topK_y[i]);
265 |         item.tracking = cv::Vec2f(track_x[i], track_y[i]);
266 |         item.bbox = topK_bbox[i];
267 |         dets.push_back(item);
268 |     }
269 |     return dets;
270 | }
271 | 
272 | void CenterTrack::WrapInputLayer(std::vector<cv::Mat> *input_channels,
273 |                                  Blob<float> *input_layer)
274 | {
275 |     // Blob<float> *input_layer = net_->input_blobs()[0];
276 | 
277 |     float *input_data = input_layer->mutable_cpu_data();
278 |     for (int i = 0; i < input_layer->channels(); ++i)
279 |     {
280 |         cv::Mat channel(inp_h, inp_w, CV_32FC1, input_data);
281 |         input_channels->push_back(channel);
282 |         input_data += inp_w * inp_h;
283 |     }
284 | }
285 | 
286 | void CenterTrack::Preprocess(const cv::Mat &img,
287 |                              std::vector<cv::Mat> &input_channels)
288 | {
289 |     cv::Mat sample_resized;
290 |     if (img.size() != cv::Size(inp_w, inp_h))
291 |     {
292 |         printf("img size: (%d, %d)\n", img.cols, img.rows);
293 |         cv::resize(img, sample_resized, cv::Size(inp_w, inp_h));
294 |     }
295 |     else
296 |         sample_resized = img;
297 | 
298 |     cv::Mat sample_float;
299 |     if (img.channels() == 3)
300 |     {
301 |         sample_resized.convertTo(sample_float, CV_32FC3);
302 |     }
303 |     else
304 |         sample_resized.convertTo(sample_float, CV_32FC1);
305 | 
306 |     /* This operation will write the separate BGR planes directly to the
307 |    * input layer of the network because it is wrapped by the cv::Mat
308 |    * objects in input_channels. */
309 |     cv::split(sample_float, input_channels);
310 | }
311 | 


--------------------------------------------------------------------------------
/src/main.cpp:
--------------------------------------------------------------------------------
 1 | #include "centerTrack.h"
 2 | // #include "util.h"
 3 | #include <array>
 4 | 
 5 | using namespace std;
 6 | 
 7 | int main(int argc, const char *argv[])
 8 | {
 9 | 
10 |     auto input_path = string("/home/lx/work/hi3519/nfs100.69/centerTrack/data/500.avi");
11 | 
12 |     cv::VideoCapture cap(input_path);
13 |     if (!cap.isOpened())
14 |     {
15 |         throw runtime_error("Cannot open cv::VideoCapture");
16 |     }
17 | 
18 |     std::string prototxtPath = "/home/lx/work/centerTrack_SDK/model/mosaic1_1533.prototxt";
19 |     std::string caffeModelPath = "/home/lx/work/centerTrack_SDK/model/mosaic1_1533.caffemodel";
20 |     CenterTrack centerTrack(prototxtPath, caffeModelPath);
21 |     // Tracker tracker = Tracker();
22 | 
23 |     cv::Size orig_dim(int(cap.get(cv::CAP_PROP_FRAME_WIDTH)),
24 |                       int(cap.get(cv::CAP_PROP_FRAME_HEIGHT)));
25 |     cv::VideoWriter outputCap("../output_video/res18_centerTrack.avi",
26 |                               CV_FOURCC('M', 'P', '4', '2'), 25.0, orig_dim);
27 |     auto image = cv::Mat();
28 |     int frame_count = 0;
29 |     while (cap.read(image))
30 |     // while (true)
31 |     {
32 |         // image = cv::imread("/media/lx/5f5da3ff-2264-4479-82cc-7053cf09f640/datasets/vehicle_type_data/convert_to_coco/image/0-2017-10-23-11-01-13-323-01-300972.jpg");
33 |         auto results = centerTrack.tracking(image, cv::Mat());
34 | 
35 |         for (auto track : results)
36 |         {
37 |             draw_result(image, track);
38 |         }
39 |         cv::putText(image, std::to_string(frame_count++), cv::Point(0, 20), CV_FONT_HERSHEY_PLAIN, 2, cv::Scalar(0, 0, 0));
40 |         outputCap << image;
41 |         // cv::imshow("image", image);
42 |         // if (cv::waitKey(0) == 27)
43 |         //     break;
44 |     }
45 |     outputCap.release();
46 | 
47 |     return 0;
48 | }
49 | 


--------------------------------------------------------------------------------
/src/tracker.cpp:
--------------------------------------------------------------------------------
 1 | #include "tracker.h"
 2 | 
 3 | Tracker::Tracker()
 4 |     : id_count(1)
 5 | {
 6 | }
 7 | 
 8 | void Tracker::reset()
 9 | {
10 |     id_count = 1;
11 |     tracks.clear();
12 | }
13 | 
14 | // void Tracker::add_track(Track result)
15 | // {
16 | //     result.track_id = id_count;
17 | //     tracks[id_count++] = result;
18 | // }
19 | 
20 | float Tracker::cal_dist(cv::Point2f p1, cv::Point2f p2)
21 | {
22 |     float distance = powf((p1.x - p2.x), 2) + powf((p1.y - p2.y), 2);
23 |     // return sqrtf(distance);
24 |     return distance;
25 | }
26 | 
27 | std::vector<Track> Tracker::step(std::vector<Track> dets)
28 | {
29 |     float first_thresh = 0.3; // 对首次出现的目标,设置更高的阈值,来增加跟踪的稳定性
30 |     // for (auto &det : dets)
31 |     for (auto iter = dets.begin(); iter != dets.end();)
32 |     {
33 |         cv::Point2f last_ct(iter->center);
34 |         last_ct.x += iter->tracking[0];
35 |         last_ct.y += iter->tracking[1];
36 |         float det_size = iter->bbox.area();
37 |         float min_dist = 1e18;
38 |         int match_idx = -1;
39 | 
40 |         for (int i = 0; i < tracks.size(); ++i)
41 |         {
42 |             auto track = tracks[i];
43 |             if (track.class_id != iter->class_id)
44 |                 continue;
45 |             float track_size = track.bbox.area();
46 | 
47 |             float dist = cal_dist(track.center, last_ct);
48 |             if (dist < track_size && dist < det_size && dist < min_dist)
49 |             {
50 |                 min_dist = dist;
51 |                 match_idx = i;
52 |             }
53 |         }
54 |         if (match_idx >= 0)
55 |         {
56 |             iter->track_id = tracks[match_idx].track_id;
57 |             iter->active = tracks[match_idx].active + 1;
58 |             ++iter;
59 |             tracks.erase(tracks.begin() + match_idx);
60 |         }
61 |         else
62 |         {
63 |             if (iter->score > first_thresh)
64 |             {
65 |                 iter->track_id = id_count++;
66 |                 iter->active = 1;
67 |                 ++iter;
68 |             }
69 |             else
70 |                 iter = dets.erase(iter);
71 |         }
72 |     }
73 |     tracks = dets;
74 |     return tracks;
75 | }
76 | 
77 | cv::Mat Tracker::gen_pre_hm(int w, int h)
78 | {
79 |     cv::Mat pre_hm = cv::Mat(h, w, CV_32FC1) * 0;
80 | 
81 |     for (auto item : tracks)
82 |     {
83 |         cv::Rect2f bbox = item.bbox;
84 |         draw_gaussian(pre_hm, bbox);
85 |     }
86 |     return pre_hm;
87 | }


--------------------------------------------------------------------------------
/src/util.cpp:
--------------------------------------------------------------------------------
 1 | #include "util.h"
 2 | #include <opencv2/opencv.hpp>
 3 | 
 4 | void draw_gaussian(cv::Mat &input_hm, cv::Rect2f bbox)
 5 | {
 6 |     float height = bbox.height;
 7 |     float width = bbox.width;
 8 | 
 9 |     float min_overlap = 0.7;
10 |     float a1 = 1;
11 |     float b1 = height + width;
12 |     float c1 = width * height * (1 - min_overlap) / (1 + min_overlap);
13 |     float sq1 = sqrtf(powf(b1, 2) - 4 * a1 * c1);
14 |     float r1 = (b1 + sq1) / 2;
15 | 
16 |     float a2 = 4;
17 |     float b2 = 2 * (height + width);
18 |     float c2 = (1 - min_overlap) * width * height;
19 |     float sq2 = sqrtf(powf(b2, 2) - 4 * a2 * c2);
20 |     float r2 = (b2 + sq2) / 2;
21 | 
22 |     float a3 = 4 * min_overlap;
23 |     float b3 = -2 * min_overlap * (height + width);
24 |     float c3 = (min_overlap - 1) * width * height;
25 |     float sq3 = sqrtf(powf(b3, 2) - 4 * a3 * c3);
26 |     float r3 = (b3 + sq3) / 2;
27 |     float radius = std::min(r1, r2);
28 |     radius = std::min(radius, r3);
29 | 
30 |     float diameter = radius * 2 + 1;
31 |     float sigma = diameter * 6;
32 |     cv::Mat gaussian_kernel = cv::getGaussianKernel(diameter, sigma);
33 | 
34 |     cv::Mat roi = input_hm(bbox);
35 |     cv::MatExpr expr = cv::max(roi, gaussian_kernel);
36 |     expr.op->assign(expr, roi);
37 | }
38 | 
39 | cv::Scalar color_map(int64_t n)
40 | {
41 |     auto bit_get = [](int64_t x, int64_t i) { return x & (1 << i); };
42 | 
43 |     int64_t r = 0, g = 0, b = 0;
44 |     int64_t i = n;
45 |     for (int64_t j = 7; j >= 0; --j)
46 |     {
47 |         r |= bit_get(i, 0) << j;
48 |         g |= bit_get(i, 1) << j;
49 |         b |= bit_get(i, 2) << j;
50 |         i >>= 3;
51 |     }
52 |     return cv::Scalar(b, g, r);
53 | }
54 | 
55 | void draw_text(cv::Mat &img, const std::string &str,
56 |                const cv::Scalar &color, cv::Point pos)
57 | {
58 |     auto t_size = cv::getTextSize(str, cv::FONT_HERSHEY_PLAIN, 1, 1, nullptr);
59 |     cv::Point bottom_left, upper_right;
60 | 
61 |     bottom_left = pos;
62 |     upper_right = cv::Point(bottom_left.x + t_size.width, bottom_left.y - t_size.height);
63 | 
64 |     cv::rectangle(img, bottom_left, upper_right, color, -1);
65 |     cv::putText(img, str, bottom_left, cv::FONT_HERSHEY_PLAIN, 1, cv::Scalar(255, 255, 255) - color);
66 | }
67 | 
68 | void draw_result(cv::Mat &img, Track track)
69 | {
70 |     auto color = color_map(track.track_id);
71 |     cv::rectangle(img, track.bbox, color);
72 |     cv::Point2f pred_pre(track.center);
73 |     pred_pre.x += track.tracking[0];
74 |     pred_pre.y += track.tracking[1];
75 |     cv::arrowedLine(img, track.center, pred_pre, cv::Scalar(255, 0, 255), 3, 8, 0, 0.2);
76 |     draw_text(img, std::to_string(track.track_id), color, track.bbox.tl());
77 | }


--------------------------------------------------------------------------------