maskrcnn onnxruntime c++

#include "opencv.hpp"
#include 
#include 
#include 

#include 
#include 
#include 
#include 
#include 

#include 
#include 
#include 
#include 
#include 
#include 
#include 

using namespace cv;
using namespace std;

const float THRESHOLD = 0.5;
const bool& isGPU = true;
std::array<float, 3> mean_vec = {102.9801, 115.9465, 122.7717};

const int IMAGE_SHAPE = 512;

inline cv::Mat visualizeOneImageWithMask(const cv::Mat& img, const std::vector<std::array<float, 4>>& bboxes,
                                         const std::vector<uint64_t>& classIndices, const std::vector<cv::Mat>& masks,
                                
                                         const float maskThreshold = 0.5)
{
    assert(bboxes.size() == classIndices.size());

    cv::Mat result = img.clone();

    for (size_t i = 0; i < bboxes.size(); ++i) {
        const auto& curBbox = bboxes[i];
        const uint64_t classIdx = classIndices[i];
        cv::Mat curMask = masks[i].clone();
        const cv::Scalar& curColor = (0,0,255);
        const std::string curLabel = "box" ;

        cv::rectangle(result, cv::Point(curBbox[0], curBbox[1]), cv::Point(curBbox[2], curBbox[3]), curColor, 2);

        int baseLine = 0;
        cv::Size labelSize = cv::getTextSize(curLabel, cv::FONT_HERSHEY_COMPLEX, 0.35, 1, &baseLine);
        cv::rectangle(result, cv::Point(curBbox[0], curBbox[1]),
                      cv::Point(curBbox[0] + labelSize.width, curBbox[1] + static_cast<int>(1.3 * labelSize.height)),
                      curColor, -1);
        cv::putText(result, curLabel, cv::Point(curBbox[0], curBbox[1] + labelSize.height), cv::FONT_HERSHEY_COMPLEX,
                    0.35, cv::Scalar(255, 255, 255));

        // ---------------------------------------------------------------------//
        // Visualize masks

        const cv::Rect curBoxRect(cv::Point(curBbox[0], curBbox[1]), cv::Point(curBbox[2], curBbox[3]));

        cv::resize(curMask, curMask, curBoxRect.size());

        cv::Mat finalMask = (curMask > maskThreshold);

        cv::Mat coloredRoi = (0.3 * curColor + 0.7 * result(curBoxRect));

        coloredRoi.convertTo(coloredRoi, CV_8UC3);

        std::vector<cv::Mat> contours;
        cv::Mat hierarchy;
        finalMask.convertTo(finalMask, CV_8U);

        cv::findContours(finalMask, contours, hierarchy, cv::RETR_TREE, cv::CHAIN_APPROX_SIMPLE);
        cv::drawContours(coloredRoi, contours, -1, curColor, 5, cv::LINE_8, hierarchy, 100);
        coloredRoi.copyTo(result(curBoxRect), finalMask);
    }

    return result;
}


void main()
{
    Ort::Env env(ORT_LOGGING_LEVEL_WARNING, "test"); /** 初始化环境，每个进程一个环境.环境保留了线程池和其他状态信息 **/
    Ort::SessionOptions session_options;

    session_options.SetIntraOpNumThreads(1);
    std::vector<std::string> availableProviders = Ort::GetAvailableProviders();
    auto cudaAvailable = std::find(availableProviders.begin(), availableProviders.end(), "CUDAExecutionProvider");
    OrtCUDAProviderOptions cudaOption{0};

    session_options.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_EXTENDED); /** 设置图像优化级别 **/
    if (isGPU && (cudaAvailable == availableProviders.end())) {
        std::cout << "GPU is not supported by your ONNXRuntime build. Fallback to CPU." << std::endl;
        std::cout << "Inference device: CPU" << std::endl;
    } else if (isGPU && (cudaAvailable != availableProviders.end())) {
        std::cout << "Inference device: GPU" << std::endl;
        session_options.AppendExecutionProvider_CUDA(cudaOption);
    } else {
        std::cout << "Inference device: CPU" << std::endl;
    }

    //*************************************************************************
    // 创建Session并把模型载入内存
#ifdef _WIN64
    const wchar_t* model_path = L"E:/2022/work/Maskrcnn_onnx/x64/Debug/MaskRCNN-10.onnx";
#else
    const char* model_path = "E:/2022/work/Maskrcnn_onnx/x64/Debug/model.onnx";
#endif
    printf("Using Onnxruntime C++ API\n");
    Ort::Session session(env, model_path, session_options);

    ///1.Prepare for the "input_image"//

    Mat img = imread("E:/2022/work/Maskrcnn_onnx/x64/Debug/demo.jpg");
    Mat dst;
    cout << img.cols << '\n' << img.rows << endl;
    // Resize
    //float ratio = 800.0 / min(img.cols, img.rows);
    //const int re_width = int(ratio * img.cols);
    //const int re_height = int(ratio * img.rows);
    //cout << re_width << '\n' << re_height << endl;
    cv::resize(img, dst, cv::Size(IMAGE_SHAPE, IMAGE_SHAPE));
    // Convert to BGR
    cv::Mat bgrimg, bgrimg_f;

    cvtColor(dst, bgrimg, cv::COLOR_RGB2BGR);

    bgrimg.convertTo(bgrimg_f, CV_32FC3);

    // hwc -> chw
    float* blob = new float[bgrimg_f.cols * bgrimg_f.rows * bgrimg_f.channels()];
    cv::Size floatImageSize{bgrimg_f.cols, bgrimg_f.rows};
    std::vector<cv::Mat> chw(bgrimg_f.channels());
    for (int i = 0; i < bgrimg_f.channels(); ++i) {
        chw[i] = cv::Mat(floatImageSize, CV_32FC1, blob + i * floatImageSize.width * floatImageSize.height);
    }
    cv::split(bgrimg_f, chw);

    for (int i = 0; i < bgrimg_f.rows; i++) {
        for (int j = 0; j < bgrimg_f.cols; j++) {
            bgrimg_f.at<Vec3f>(i, j)[0] -= mean_vec[0];
            bgrimg_f.at<Vec3f>(i, j)[1] -= mean_vec[1];
            bgrimg_f.at<Vec3f>(i, j)[2] -= mean_vec[2];
        }
    }
    const int padd_h = int(ceil(bgrimg_f.rows / 32) * 32);
    const int padd_w = int(ceil(bgrimg_f.cols / 32) * 32);
    cv::Mat paddimg;
    cv::resize(bgrimg_f, paddimg, cv::Size(padd_w, padd_h));
    cout << paddimg.cols << '\n' << paddimg.rows << endl;

    std::vector<float> output;

    for (int c = 0; c < 3; c++) {
        for (int i = 0; i < padd_h; i++) {
            for (int j = 0; j < padd_w; j++) {
                output.push_back(paddimg.ptr<float>(i)[j * 3 + c]);
            }
        }
    }
    std::vector<int64_t> input_shape_ = {3, padd_h, padd_w};
    /Prepare for the "input"/

    auto allocator_info = Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeCPU);
    Ort::Value input_tensor_ = Ort::Value::CreateTensor<float>(allocator_info, output.data(), output.size(),
                                                               input_shape_.data(), input_shape_.size());

    const char* input_names[] = {"image"};
    const char* output_names[] = {"6568", "6570", "6572", "6887"};

    std::vector<Ort::Value> ort_inputs;
    ort_inputs.push_back(std::move(input_tensor_));


    for (int i = 0; i < 10; i++) {
        double timeStart = (double)getTickCount();
        auto outputTensors = session.Run(Ort::RunOptions{nullptr}, input_names, ort_inputs.data(), 1, output_names, 4);
        double nTime = ((double)getTickCount() - timeStart) / getTickFrequency();
        cout << "running time ：" << nTime << "sec\n" << endl;

        using DataOutputType = std::pair<float*, std::vector<int64_t>>;
        std::vector<DataOutputType> outputData;
        outputData.reserve(4);

        int count = 1;
        for (auto& elem : outputTensors) {
            outputData.emplace_back(std::make_pair(std::move(elem.GetTensorMutableData<float>()),
                                                   elem.GetTensorTypeAndShapeInfo().GetShape()));
        }


        /.Get Outputs and check them/

        size_t nBoxes = outputData[1].second[0];

        std::vector<std::array<float, 4>> bboxes;
        std::vector<uint64_t> classIndices;
        std::vector<cv::Mat> masks;

        bboxes.reserve(nBoxes);
        classIndices.reserve(nBoxes);
        masks.reserve(nBoxes);

        for (size_t i = 0; i < nBoxes; ++i) {
            if (outputData[2].first[i] > 0.7) {
                float xmin = outputData[0].first[i * 4 + 0];
                float ymin = outputData[0].first[i * 4 + 1];
                float xmax = outputData[0].first[i * 4 + 2];
                float ymax = outputData[0].first[i * 4 + 3];

                xmin = std::max<float>(xmin, 0);
                ymin = std::max<float>(ymin, 0);
                xmax = std::min<float>(xmax, 512 - 1);
                ymax = std::min<float>(ymax, 512 - 1);

                bboxes.emplace_back(std::array<float, 4>{xmin, ymin, xmax, ymax});
                classIndices.emplace_back(reinterpret_cast<int64_t*>(outputData[1].first)[i]);

                cv::Mat curMask(28, 28, CV_32FC1);
                memcpy(curMask.data, outputData[3].first + i * 28 * 28, 28 * 28 * sizeof(float));
                masks.emplace_back(curMask);
            }
        }
        cv::Mat result = visualizeOneImageWithMask(dst, bboxes, classIndices, masks);

        cv::imshow("result", result);
        cv::waitKey(0);
    }

}