File size: 3,127 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 | #include <opencv2/dnn.hpp>
#include <opencv2/imgproc.hpp>
#include <opencv2/imgcodecs.hpp>
#include <array>
#include <cstdint>
#include <iostream>
#include <string>
#include <vector>
using namespace cv;
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", "ssd_inception_v2_coco_2017_11_17_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"));
Mat img = imread(image);
if (img.empty())
{
std::cerr << "could not read image: " << image << std::endl;
return 1;
}
Mat rgb;
cvtColor(img, rgb, COLOR_BGR2RGB);
resize(rgb, rgb, Size(300, 300));
if (!rgb.isContinuous()) rgb = rgb.clone();
int blobShape[] = {1, 300, 300, 3};
Mat blob(4, blobShape, CV_8U, rgb.data);
dnn::Net net = dnn::readNetFromONNX(model, dnn::ENGINE_ORT);
net.setInput(blob);
std::vector<String> out_str = {"detection_boxes:0", "detection_scores:0", "detection_classes:0", "num_detections:0"};
std::vector<Mat> outs;
net.forward(outs, out_str);
const float *boxes = 0, *scores = 0, *classes = 0, *num = 0;
for (size_t i = 0; i < out_str.size(); ++i)
{
const std::string& n = out_str[i];
if (n.find("detection_boxes") != std::string::npos) boxes = (const float*)outs[i].data;
else if (n.find("detection_scores") != std::string::npos) scores = (const float*)outs[i].data;
else if (n.find("detection_classes") != std::string::npos) classes = (const float*)outs[i].data;
else if (n.find("num_detections") != std::string::npos) num = (const float*)outs[i].data;
}
if (!boxes || !scores || !classes || !num)
{
std::cerr << "missing expected output tensors" << std::endl;
return 1;
}
int nd = (int)num[0];
int h = img.rows, w = img.cols;
Mat out = img.clone();
std::vector<std::string> lines;
for (int k = 0; k < nd; ++k)
{
if (scores[k] < conf) continue;
float ymin = boxes[k * 4 + 0], xmin = boxes[k * 4 + 1];
float ymax = boxes[k * 4 + 2], xmax = boxes[k * 4 + 3];
int cls = (int)classes[k];
rectangle(out, Point(int(xmin * w), int(ymin * h)), Point(int(xmax * w), int(ymax * h)), Scalar(0, 255, 0), 2);
putText(out, format("%d:%.2f", cls, scores[k]), Point(int(xmin * w), int(ymin * h) - 5),
FONT_HERSHEY_SIMPLEX, 0.5, Scalar(0, 255, 0), 1);
lines.push_back(format("%d %.3f %.3f %.3f %.3f %.3f", cls, scores[k], xmin, ymin, xmax, ymax));
}
imwrite(output, out);
std::cout << "ssd_inception_v2_coco_2017_11_17 " << lines.size() << " detections" << std::endl;
for (const auto& l : lines) std::cout << l << std::endl;
return 0;
}
|