main.cc 1.1 KB

1234567891011121314151617181920212223242526272829303132333435363738394041
  1. #include "inference.h"
  2. #include <iostream>
  3. #include <opencv2/highgui.hpp>
  4. int main(int argc, char **argv) {
  5. // Check if the correct number of arguments is provided
  6. if (argc != 3) {
  7. std::cerr << "usage: " << argv[0] << " <model_path> <image_path>" << std::endl;
  8. return 1;
  9. }
  10. // Get the model and image paths from the command-line arguments
  11. const std::string model_path = argv[1];
  12. const std::string image_path = argv[2];
  13. // Read the input image
  14. cv::Mat image = cv::imread(image_path);
  15. // Check if the image was successfully loaded
  16. if (image.empty()) {
  17. std::cerr << "ERROR: image is empty" << std::endl;
  18. return 1;
  19. }
  20. // Define the confidence and NMS thresholds
  21. const float confidence_threshold = 0.5;
  22. const float NMS_threshold = 0.5;
  23. // Initialize the YOLO inference with the specified model and parameters
  24. yolo::Inference inference(model_path, cv::Size(640, 640), confidence_threshold, NMS_threshold);
  25. // Run inference on the input image
  26. inference.RunInference(image);
  27. // Display the image with the detections
  28. cv::imshow("image", image);
  29. cv::waitKey(0);
  30. return 0;
  31. }