#include <opencv2/core.hpp>
#include <opencv2/dnn.hpp>

#include <iostream>

using namespace std;

void switch_to_cuda(cv::dnn::Net &net)
{
    for(auto const &vpair : cv::dnn::getAvailableBackends()){
        if(vpair.first == cv::dnn::DNN_BACKEND_CUDA && vpair.second == cv::dnn::DNN_TARGET_CUDA){
            net.setPreferableBackend(cv::dnn::DNN_BACKEND_CUDA);
            net.setPreferableTarget(cv::dnn::DNN_TARGET_CUDA);
            cout<<__func__<<":can switch to cuda"<<endl; //can switch to cuda
            break;
        }
    }
}

int main()
{
    cv::Mat input(128, 256, CV_32FC3);
    auto blob = cv::dnn::blobFromImage(input);
    cout<<__func__<<":blob get"<<endl;

    cout<<__func__<<":set input"<<endl;
    auto net = cv::dnn::readNetFromONNX("osnet_ain_x1_0.onnx");

    //switch_to_cuda(net); //crash if switch to cuda

    net.setInput(blob);
    cout<<__func__<<":forward"<<endl;
    net.forward();
    cout<<__func__<<":forward done"<<endl;

    std::cout<<__func__<<":forward done"<<std::endl;

    return 0;
}