File size: 3,138 Bytes
e9eb33d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 | #include <opencv2/dnn.hpp>
#include <opencv2/imgproc.hpp>
#include <opencv2/imgcodecs.hpp>
#include <array>
#include <cstdint>
#include <iostream>
#include <string>
#include <vector>
static std::string argVal(int argc, char** argv, const std::string& key, const std::string& def)
{
for (int i = 1; i + 1 < argc; ++i)
if (key == argv[i]) return argv[i + 1];
return def;
}
int main(int argc, char** argv)
{
std::string model = argVal(argc, argv, "--model", "faster_rcnn_resnet50_coco_2018_01_28_2026jul.onnx");
std::string image = argVal(argc, argv, "--image", "example_outputs/input_image.png");
std::string output = argVal(argc, argv, "--output", "example_outputs/output_image.png");
float conf = std::stof(argVal(argc, argv, "--conf", "0.3"));
cv::Mat img = cv::imread(image);
if (img.empty())
{
std::cerr << "could not read image: " << image << std::endl;
return 1;
}
const int W = 800, H = 600;
cv::Mat rgb;
cv::cvtColor(img, rgb, cv::COLOR_BGR2RGB);
cv::resize(rgb, rgb, cv::Size(W, H));
if (!rgb.isContinuous()) rgb = rgb.clone();
int blobShape[] = {1, H, W, 3};
cv::Mat blob(4, blobShape, CV_8U, rgb.data);
cv::dnn::Net net = cv::dnn::readNetFromONNX(model, cv::dnn::ENGINE_ORT);
net.setInput(blob);
std::vector<cv::String> out_strs = {"detection_boxes:0", "detection_scores:0", "detection_classes:0", "num_detections:0"};
std::vector<cv::Mat> outs;
net.forward(outs, out_strs);
float* boxes = nullptr;
float* scores = nullptr;
float* classes = nullptr;
float* numd = nullptr;
for (size_t i = 0; i < out_strs.size(); ++i)
{
float* p = (float*)outs[i].data;
const std::string& n = out_strs[i];
if (n.find("detection_boxes") != std::string::npos) boxes = p;
else if (n.find("detection_scores") != std::string::npos) scores = p;
else if (n.find("detection_classes") != std::string::npos) classes = p;
else if (n.find("num_detections") != std::string::npos) numd = p;
}
int nd = (int)numd[0];
int w = img.cols, h = img.rows;
std::vector<int> kept;
for (int i = 0; i < nd; ++i)
if (scores[i] >= conf) kept.push_back(i);
std::cout << "faster_rcnn_resnet50_coco_2018_01_28 " << kept.size() << " detections" << std::endl;
for (int i : kept)
{
int cls = (int)classes[i] - 1;
float score = scores[i];
float ymin = boxes[i * 4 + 0], xmin = boxes[i * 4 + 1];
float ymax = boxes[i * 4 + 2], xmax = boxes[i * 4 + 3];
cv::Point p1((int)(xmin * w), (int)(ymin * h));
cv::Point p2((int)(xmax * w), (int)(ymax * h));
cv::rectangle(img, p1, p2, cv::Scalar(0, 255, 0), 2);
cv::putText(img, cv::format("%d:%.2f", cls, score), cv::Point(p1.x, p1.y - 5),
cv::FONT_HERSHEY_SIMPLEX, 0.5, cv::Scalar(0, 255, 0), 1);
std::cout << cls << " " << cv::format("%.3f %.3f %.3f %.3f %.3f", score, xmin, ymin, xmax, ymax) << std::endl;
}
cv::imwrite(output, img);
std::cout << "wrote " << output << std::endl;
return 0;
}
|