ssd_inception_v2_coco_2017_11_17/demo.cpp

#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;
}