diff --git a/src/yolo_v2_class.cpp b/src/yolo_v2_class.cpp index 8643a22a..d79021a1 100644 --- a/src/yolo_v2_class.cpp +++ b/src/yolo_v2_class.cpp @@ -201,4 +201,54 @@ YOLODLL_API std::vector Detector::detect(image_t img, float thresh) cudaSetDevice(old_gpu_index); return bbox_vec; +} + +YOLODLL_API std::vector Detector::tracking(std::vector cur_bbox_vec, int const frames_story) +{ + bool prev_track_id_present = false; + for (auto &i : prev_bbox_vec_deque) + if (i.size() > 0) prev_track_id_present = true; + + static unsigned int track_id = 1; + + if (!prev_track_id_present) { + //track_id = 1; + for (size_t i = 0; i < cur_bbox_vec.size(); ++i) + cur_bbox_vec[i].track_id = track_id++; + prev_bbox_vec_deque.push_front(cur_bbox_vec); + if (prev_bbox_vec_deque.size() > frames_story) prev_bbox_vec_deque.pop_back(); + return cur_bbox_vec; + } + + std::vector dist_vec(cur_bbox_vec.size(), std::numeric_limits::max()); + + for (auto &prev_bbox_vec : prev_bbox_vec_deque) { + for (auto &i : prev_bbox_vec) { + int cur_index = -1; + for (size_t m = 0; m < cur_bbox_vec.size(); ++m) { + bbox_t const& k = cur_bbox_vec[m]; + if (i.obj_id == k.obj_id) { + unsigned int cur_dist = sqrt(((float)i.x - k.x)*((float)i.x - k.x) + ((float)i.y - k.y)*((float)i.y - k.y)); + if (cur_dist < 100 && (k.track_id == 0 || dist_vec[m] > cur_dist)) { + dist_vec[m] = cur_dist; + cur_index = m; + } + } + } + + bool track_id_absent = !std::any_of(cur_bbox_vec.begin(), cur_bbox_vec.end(), [&](bbox_t const& b) { return b.track_id == i.track_id; }); + + if (cur_index >= 0 && track_id_absent) + cur_bbox_vec[cur_index].track_id = i.track_id; + } + } + + for (size_t i = 0; i < cur_bbox_vec.size(); ++i) + if (cur_bbox_vec[i].track_id == 0) + cur_bbox_vec[i].track_id = track_id++; + + prev_bbox_vec_deque.push_front(cur_bbox_vec); + if (prev_bbox_vec_deque.size() > frames_story) prev_bbox_vec_deque.pop_back(); + + return cur_bbox_vec; } \ No newline at end of file diff --git a/src/yolo_v2_class.hpp b/src/yolo_v2_class.hpp index 37fcd61a..34e220a6 100644 --- a/src/yolo_v2_class.hpp +++ b/src/yolo_v2_class.hpp @@ -44,6 +44,8 @@ public: static YOLODLL_API image_t load_image(std::string image_filename); static YOLODLL_API void free_image(image_t m); + YOLODLL_API std::vector tracking(std::vector cur_bbox_vec, int const frames_story = 4); + #ifdef OPENCV std::vector detect(cv::Mat mat, float thresh = 0.2) { @@ -113,57 +115,6 @@ private: #endif // OPENCV std::deque> prev_bbox_vec_deque; - -public: - std::vector tracking(std::vector cur_bbox_vec, int const frames_story = 4) - { - bool prev_track_id_present = false; - for (auto &i : prev_bbox_vec_deque) - if (i.size() > 0) prev_track_id_present = true; - - static unsigned int track_id = 1; - - if(!prev_track_id_present) { - //track_id = 1; - for (size_t i = 0; i < cur_bbox_vec.size(); ++i) - cur_bbox_vec[i].track_id = track_id++; - prev_bbox_vec_deque.push_front(cur_bbox_vec); - if (prev_bbox_vec_deque.size() > frames_story) prev_bbox_vec_deque.pop_back(); - return cur_bbox_vec; - } - - std::vector dist_vec(cur_bbox_vec.size(), std::numeric_limits::max()); - - for (auto &prev_bbox_vec : prev_bbox_vec_deque) { - for (auto &i : prev_bbox_vec) { - int cur_index = -1; - for (size_t m = 0; m < cur_bbox_vec.size(); ++m) { - bbox_t const& k = cur_bbox_vec[m]; - if (i.obj_id == k.obj_id) { - unsigned int cur_dist = sqrt(((float)i.x - k.x)*((float)i.x - k.x) + ((float)i.y - k.y)*((float)i.y - k.y)); - if (cur_dist < 100 && (k.track_id == 0 || dist_vec[m] > cur_dist)) { - dist_vec[m] = cur_dist; - cur_index = m; - } - } - } - - bool track_id_absent = !std::any_of(cur_bbox_vec.begin(), cur_bbox_vec.end(), [&](bbox_t const& b) { return b.track_id == i.track_id; }); - - if (cur_index >= 0 && track_id_absent) - cur_bbox_vec[cur_index].track_id = i.track_id; - } - } - - for (size_t i = 0; i < cur_bbox_vec.size(); ++i) - if (cur_bbox_vec[i].track_id == 0) - cur_bbox_vec[i].track_id = track_id++; - - prev_bbox_vec_deque.push_front(cur_bbox_vec); - if (prev_bbox_vec_deque.size() > frames_story) prev_bbox_vec_deque.pop_back(); - - return cur_bbox_vec; - } };