image_classification.cpp 1.7 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667
  1. #include <opencv2/imgproc.hpp>
  2. #include <opencv2/highgui.hpp>
  3. #include <opencv2/dnn.hpp>
  4. #include <iostream>
  5. #include <cstdlib>
  6. int main(int argc, char **argv)
  7. {
  8. if (argc < 4)
  9. {
  10. std::cerr << "Usage " << argv[0] << ": "
  11. << "<model-definition-file> " << " "
  12. << "<model-weights-file> " << " "
  13. << "<test-image>\n";
  14. return -1;
  15. }
  16. cv::String model_prototxt = argv[1];
  17. cv::String model_binary = argv[2];
  18. cv::String test_image = argv[3];
  19. cv::dnn::Net net = cv::dnn::readNetFromCaffe(model_prototxt, model_binary);
  20. if (net.empty())
  21. {
  22. std::cerr << "Couldn't load the model !\n";
  23. return -2;
  24. }
  25. cv::Mat img = cv::imread(test_image);
  26. if (img.empty())
  27. {
  28. std::cerr << "Couldn't load image: " << test_image << "\n";
  29. return -3;
  30. }
  31. cv::Mat input_blob = cv::dnn::blobFromImage(
  32. img, 1.0, cv::Size(416, 416), cv::Scalar(104, 117, 123), false);
  33. cv::Mat prob;
  34. cv::TickMeter t;
  35. net.setInput(input_blob);
  36. t.start();
  37. prob = net.forward("predictions");
  38. t.stop();
  39. int prob_size[3] = {1000, 1, 1};
  40. cv::Mat prob_data(3, prob_size, CV_32F, prob.ptr<float>(0));
  41. double max_prob = -1.0;
  42. int class_idx = -1;
  43. for (int idx = 0; idx < prob.size[1]; ++idx)
  44. {
  45. double current_prob = prob_data.at<float>(idx, 0, 0);
  46. if (current_prob > max_prob)
  47. {
  48. max_prob = current_prob;
  49. class_idx = idx;
  50. }
  51. }
  52. std::cout << "Best class Index: " << class_idx << "\n";
  53. std::cout << "Time taken: " << t.getTimeSec() << "\n";
  54. std::cout << "Probability: " << max_prob * 100.0<< "\n";
  55. return 0;
  56. }