faster_rcnn_resnet50_coco_2018_01_28/demo.cpp

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