nanotrack_tracker.cpp 5.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183
  1. // NanoTrack
  2. // Link to original inference code: https://github.com/HonglinChu/NanoTrack
  3. // Link to original training repo: https://github.com/HonglinChu/SiamTrackers/tree/master/NanoTrack
  4. // backBone model: https://github.com/HonglinChu/SiamTrackers/blob/master/NanoTrack/models/onnx/nanotrack_backbone_sim.onnx
  5. // headNeck model: https://github.com/HonglinChu/SiamTrackers/blob/master/NanoTrack/models/onnx/nanotrack_head_sim.onnx
  6. #include <iostream>
  7. #include <cmath>
  8. #include <opencv2/dnn.hpp>
  9. #include <opencv2/imgproc.hpp>
  10. #include <opencv2/highgui.hpp>
  11. #include <opencv2/video.hpp>
  12. using namespace cv;
  13. using namespace cv::dnn;
  14. const char *keys =
  15. "{ help h | | Print help message }"
  16. "{ input i | | Full path to input video folder, the specific camera index. (empty for camera 0) }"
  17. "{ backbone | backbone.onnx | Path to onnx model of backbone.onnx}"
  18. "{ headneck | headneck.onnx | Path to onnx model of headneck.onnx }"
  19. "{ backend | 0 | Choose one of computation backends: "
  20. "0: automatically (by default), "
  21. "1: Halide language (http://halide-lang.org/), "
  22. "2: Intel's Deep Learning Inference Engine (https://software.intel.com/openvino-toolkit), "
  23. "3: OpenCV implementation, "
  24. "4: VKCOM, "
  25. "5: CUDA },"
  26. "{ target | 0 | Choose one of target computation devices: "
  27. "0: CPU target (by default), "
  28. "1: OpenCL, "
  29. "2: OpenCL fp16 (half-float precision), "
  30. "3: VPU, "
  31. "4: Vulkan, "
  32. "6: CUDA, "
  33. "7: CUDA fp16 (half-float preprocess) }"
  34. ;
  35. static
  36. int run(int argc, char** argv)
  37. {
  38. // Parse command line arguments.
  39. CommandLineParser parser(argc, argv, keys);
  40. if (parser.has("help"))
  41. {
  42. parser.printMessage();
  43. return 0;
  44. }
  45. std::string inputName = parser.get<String>("input");
  46. std::string backbone = parser.get<String>("backbone");
  47. std::string headneck = parser.get<String>("headneck");
  48. int backend = parser.get<int>("backend");
  49. int target = parser.get<int>("target");
  50. Ptr<TrackerNano> tracker;
  51. try
  52. {
  53. TrackerNano::Params params;
  54. params.backbone = samples::findFile(backbone);
  55. params.neckhead = samples::findFile(headneck);
  56. params.backend = backend;
  57. params.target = target;
  58. tracker = TrackerNano::create(params);
  59. }
  60. catch (const cv::Exception& ee)
  61. {
  62. std::cerr << "Exception: " << ee.what() << std::endl;
  63. std::cout << "Can't load the network by using the following files:" << std::endl;
  64. std::cout << "backbone : " << backbone << std::endl;
  65. std::cout << "headneck : " << headneck << std::endl;
  66. return 2;
  67. }
  68. const std::string winName = "NanoTrack";
  69. namedWindow(winName, WINDOW_AUTOSIZE);
  70. // Open a video file or an image file or a camera stream.
  71. VideoCapture cap;
  72. if (inputName.empty() || (isdigit(inputName[0]) && inputName.size() == 1))
  73. {
  74. int c = inputName.empty() ? 0 : inputName[0] - '0';
  75. std::cout << "Trying to open camera #" << c << " ..." << std::endl;
  76. if (!cap.open(c))
  77. {
  78. std::cout << "Capture from camera #" << c << " didn't work. Specify -i=<video> parameter to read from video file" << std::endl;
  79. return 2;
  80. }
  81. }
  82. else if (inputName.size())
  83. {
  84. inputName = samples::findFileOrKeep(inputName);
  85. if (!cap.open(inputName))
  86. {
  87. std::cout << "Could not open: " << inputName << std::endl;
  88. return 2;
  89. }
  90. }
  91. // Read the first image.
  92. Mat image;
  93. cap >> image;
  94. if (image.empty())
  95. {
  96. std::cerr << "Can't capture frame!" << std::endl;
  97. return 2;
  98. }
  99. Mat image_select = image.clone();
  100. putText(image_select, "Select initial bounding box you want to track.", Point(0, 15), FONT_HERSHEY_SIMPLEX, 0.5, Scalar(0, 255, 0));
  101. putText(image_select, "And Press the ENTER key.", Point(0, 35), FONT_HERSHEY_SIMPLEX, 0.5, Scalar(0, 255, 0));
  102. Rect selectRect = selectROI(winName, image_select);
  103. std::cout << "ROI=" << selectRect << std::endl;
  104. tracker->init(image, selectRect);
  105. TickMeter tickMeter;
  106. for (int count = 0; ; ++count)
  107. {
  108. cap >> image;
  109. if (image.empty())
  110. {
  111. std::cerr << "Can't capture frame " << count << ". End of video stream?" << std::endl;
  112. break;
  113. }
  114. Rect rect;
  115. tickMeter.start();
  116. bool ok = tracker->update(image, rect);
  117. tickMeter.stop();
  118. float score = tracker->getTrackingScore();
  119. std::cout << "frame " << count <<
  120. ": predicted score=" << score <<
  121. " rect=" << rect <<
  122. " time=" << tickMeter.getTimeMilli() << "ms" <<
  123. std::endl;
  124. Mat render_image = image.clone();
  125. if (ok)
  126. {
  127. rectangle(render_image, rect, Scalar(0, 255, 0), 2);
  128. std::string timeLabel = format("Inference time: %.2f ms", tickMeter.getTimeMilli());
  129. std::string scoreLabel = format("Score: %f", score);
  130. putText(render_image, timeLabel, Point(0, 15), FONT_HERSHEY_SIMPLEX, 0.5, Scalar(0, 255, 0));
  131. putText(render_image, scoreLabel, Point(0, 35), FONT_HERSHEY_SIMPLEX, 0.5, Scalar(0, 255, 0));
  132. }
  133. imshow(winName, render_image);
  134. tickMeter.reset();
  135. int c = waitKey(1);
  136. if (c == 27 /*ESC*/)
  137. break;
  138. }
  139. std::cout << "Exit" << std::endl;
  140. return 0;
  141. }
  142. int main(int argc, char **argv)
  143. {
  144. try
  145. {
  146. return run(argc, argv);
  147. }
  148. catch (const std::exception& e)
  149. {
  150. std::cerr << "FATAL: C++ exception: " << e.what() << std::endl;
  151. return 1;
  152. }
  153. }