68 lines
1.7 KiB
C++
68 lines
1.7 KiB
C++
|
#include <opencv2/imgproc.hpp>
|
||
|
|
#include <opencv2/highgui.hpp>
|
||
|
|
#include <opencv2/dnn.hpp>
|
||
|
|
|
||
|
|
#include <iostream>
|
||
|
|
#include <cstdlib>
|
||
|
|
|
||
|
|
int main(int argc, char **argv)
|
||
|
|
{
|
||
|
|
|
||
|
|
if (argc < 4)
|
||
|
|
{
|
||
|
|
std::cerr << "Usage " << argv[0] << ": "
|
||
|
|
<< "<model-definition-file> " << " "
|
||
|
|
<< "<model-weights-file> " << " "
|
||
|
|
<< "<test-image>\n";
|
||
|
|
return -1;
|
||
|
|
|
||
|
|
}
|
||
|
|
cv::String model_prototxt = argv[1];
|
||
|
|
cv::String model_binary = argv[2];
|
||
|
|
cv::String test_image = argv[3];
|
||
|
|
cv::dnn::Net net = cv::dnn::readNetFromCaffe(model_prototxt, model_binary);
|
||
|
|
|
||
|
|
if (net.empty())
|
||
|
|
{
|
||
|
|
std::cerr << "Couldn't load the model !\n";
|
||
|
|
return -2;
|
||
|
|
}
|
||
|
|
cv::Mat img = cv::imread(test_image);
|
||
|
|
if (img.empty())
|
||
|
|
{
|
||
|
|
std::cerr << "Couldn't load image: " << test_image << "\n";
|
||
|
|
return -3;
|
||
|
|
}
|
||
|
|
|
||
|
|
cv::Mat input_blob = cv::dnn::blobFromImage(
|
||
|
|
img, 1.0, cv::Size(416, 416), cv::Scalar(104, 117, 123), false);
|
||
|
|
|
||
|
|
cv::Mat prob;
|
||
|
|
cv::TickMeter t;
|
||
|
|
|
||
|
|
net.setInput(input_blob);
|
||
|
|
t.start();
|
||
|
|
prob = net.forward("predictions");
|
||
|
|
t.stop();
|
||
|
|
|
||
|
|
int prob_size[3] = {1000, 1, 1};
|
||
|
|
cv::Mat prob_data(3, prob_size, CV_32F, prob.ptr<float>(0));
|
||
|
|
|
||
|
|
double max_prob = -1.0;
|
||
|
|
int class_idx = -1;
|
||
|
|
for (int idx = 0; idx < prob.size[1]; ++idx)
|
||
|
|
{
|
||
|
|
double current_prob = prob_data.at<float>(idx, 0, 0);
|
||
|
|
if (current_prob > max_prob)
|
||
|
|
{
|
||
|
|
max_prob = current_prob;
|
||
|
|
class_idx = idx;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
std::cout << "Best class Index: " << class_idx << "\n";
|
||
|
|
std::cout << "Time taken: " << t.getTimeSec() << "\n";
|
||
|
|
std::cout << "Probability: " << max_prob * 100.0<< "\n";
|
||
|
|
|
||
|
|
return 0;
|
||
|
|
}
|