From f4a30ef60a3c9288967386d3b023fa3cf7ffb465 Mon Sep 17 00:00:00 2001 From: AlexeyAB Date: Tue, 8 Aug 2017 22:16:28 +0300 Subject: [PATCH] Track id fixed --- src/yolo_console_dll.cpp | 13 ++++++++----- src/yolo_v2_class.cpp | 10 ++++++---- 2 files changed, 14 insertions(+), 9 deletions(-) diff --git a/src/yolo_console_dll.cpp b/src/yolo_console_dll.cpp index cada95ec..e1f2c7e8 100644 --- a/src/yolo_console_dll.cpp +++ b/src/yolo_console_dll.cpp @@ -63,17 +63,19 @@ std::vector objects_names_from_file(std::string const filename) { } -int main() +int main(int argc, char *argv[]) { + std::string filename; + if (argc > 1) filename = argv[1]; + Detector detector("yolo-voc.cfg", "yolo-voc.weights"); auto obj_names = objects_names_from_file("data/voc.names"); while (true) - { - std::string filename; + { std::cout << "input image or video filename: "; - std::cin >> filename; + if(filename.size() == 0) std::cin >> filename; if (filename.size() == 0) break; try { @@ -94,7 +96,7 @@ int main() td.join(); result_vec = thread_result_vec; det_frame = frame; - td = std::thread([&]() { thread_result_vec = detector.detect(det_frame, 0.2, true); ready_flag = true; }); + td = std::thread([&]() { thread_result_vec = detector.detect(det_frame, 0.24, true); ready_flag = true; }); if (!prev_frame.empty()) { result_vec = detector.tracking(result_vec); // comment it - if track_id is not required @@ -138,6 +140,7 @@ int main() } catch (std::exception &e) { std::cerr << "exception: " << e.what() << "\n"; getchar(); } catch (...) { std::cerr << "unknown exception \n"; getchar(); } + filename.clear(); } return 0; diff --git a/src/yolo_v2_class.cpp b/src/yolo_v2_class.cpp index f036b2a5..4418ce91 100644 --- a/src/yolo_v2_class.cpp +++ b/src/yolo_v2_class.cpp @@ -257,7 +257,6 @@ YOLODLL_API std::vector Detector::tracking(std::vector cur_bbox_ if (i.size() > 0) prev_track_id_present = true; 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 = det_gpu.track_id[cur_bbox_vec[i].obj_id]++; prev_bbox_vec_deque.push_front(cur_bbox_vec); @@ -273,7 +272,9 @@ YOLODLL_API std::vector Detector::tracking(std::vector cur_bbox_ 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)); + float center_x_diff = (float)(i.x + i.w/2) - (float)(k.x + k.w/2); + float center_y_diff = (float)(i.y + i.h/2) - (float)(k.y + k.h/2); + unsigned int cur_dist = sqrt(center_x_diff*center_x_diff + center_y_diff*center_y_diff); if (cur_dist < 100 && (k.track_id == 0 || dist_vec[m] > cur_dist)) { dist_vec[m] = cur_dist; cur_index = m; @@ -281,9 +282,10 @@ YOLODLL_API std::vector Detector::tracking(std::vector cur_bbox_ } } - 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; }); + bool track_id_absent = !std::any_of(cur_bbox_vec.begin(), cur_bbox_vec.end(), + [&i](bbox_t const& b) { return b.track_id == i.track_id && b.obj_id == i.obj_id; }); - if (cur_index >= 0 && track_id_absent) { + if (cur_index >= 0 && track_id_absent){ cur_bbox_vec[cur_index].track_id = i.track_id; cur_bbox_vec[cur_index].w = (cur_bbox_vec[cur_index].w + i.w) / 2; cur_bbox_vec[cur_index].h = (cur_bbox_vec[cur_index].h + i.h) / 2;