main.cpp 2.2 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970
  1. #include <iostream>
  2. #include <vector>
  3. #include <getopt.h>
  4. #include <opencv2/opencv.hpp>
  5. #include "inference.h"
  6. using namespace std;
  7. using namespace cv;
  8. int main(int argc, char **argv)
  9. {
  10. std::string projectBasePath = "/home/user/ultralytics"; // Set your ultralytics base path
  11. bool runOnGPU = true;
  12. //
  13. // Pass in either:
  14. //
  15. // "yolov8s.onnx" or "yolov5s.onnx"
  16. //
  17. // To run Inference with yolov8/yolov5 (ONNX)
  18. //
  19. // Note that in this example the classes are hard-coded and 'classes.txt' is a place holder.
  20. Inference inf(projectBasePath + "/yolov8s.onnx", cv::Size(640, 480), "classes.txt", runOnGPU);
  21. std::vector<std::string> imageNames;
  22. imageNames.push_back(projectBasePath + "/ultralytics/assets/bus.jpg");
  23. imageNames.push_back(projectBasePath + "/ultralytics/assets/zidane.jpg");
  24. for (int i = 0; i < imageNames.size(); ++i)
  25. {
  26. cv::Mat frame = cv::imread(imageNames[i]);
  27. // Inference starts here...
  28. std::vector<Detection> output = inf.runInference(frame);
  29. int detections = output.size();
  30. std::cout << "Number of detections:" << detections << std::endl;
  31. for (int i = 0; i < detections; ++i)
  32. {
  33. Detection detection = output[i];
  34. cv::Rect box = detection.box;
  35. cv::Scalar color = detection.color;
  36. // Detection box
  37. cv::rectangle(frame, box, color, 2);
  38. // Detection box text
  39. std::string classString = detection.className + ' ' + std::to_string(detection.confidence).substr(0, 4);
  40. cv::Size textSize = cv::getTextSize(classString, cv::FONT_HERSHEY_DUPLEX, 1, 2, 0);
  41. cv::Rect textBox(box.x, box.y - 40, textSize.width + 10, textSize.height + 20);
  42. cv::rectangle(frame, textBox, color, cv::FILLED);
  43. cv::putText(frame, classString, cv::Point(box.x + 5, box.y - 10), cv::FONT_HERSHEY_DUPLEX, 1, cv::Scalar(0, 0, 0), 2, 0);
  44. }
  45. // Inference ends here...
  46. // This is only for preview purposes
  47. float scale = 0.8;
  48. cv::resize(frame, frame, cv::Size(frame.cols*scale, frame.rows*scale));
  49. cv::imshow("Inference", frame);
  50. cv::waitKey(-1);
  51. }
  52. }