RobustVideoMatting_ONNXRUNTIME_CPP/main.cpp at main · RandomPrototypes/RobustVideoMatting_ONNXRUNTIME_CPP · GitHub

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
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
#include <onnxruntime_cxx_api.h>
#include <stdio.h>
#include <opencv2/opencv.hpp>
#include <opencv2/dnn/dnn.hpp>

int main(int argc, char **argv)
{
	cv::Size size = cv::Size(1280,720);
	const char *onnxModelFilename = argc >= 2 ? argv[1] : "rvm_mobilenetv3_fp32.onnx";
	bool use_CUDA = false;
	if(argc >= 3) {
		if(!strcmp(argv[2], "GPU") || !strcmp(argv[2], "CUDA"))
			use_CUDA = true;
		else if(!strcmp(argv[2], "CPU"))
			use_CUDA = false;
		else {
			printf("invalid argument 2 : %s, should be CPU, GPU or CUDA\n", argv[2]);
		}
	}
	int deviceID = 0;

	//creates the onnx runtime environment
	Ort::Env env(OrtLoggingLevel::ORT_LOGGING_LEVEL_WARNING, "segmentation");
    Ort::SessionOptions sessionOptions;
    sessionOptions.SetIntraOpNumThreads(1);

    //activates the CUDA backend
    if(use_CUDA) {
    	OrtCUDAProviderOptions cuda_options;
    	sessionOptions.AppendExecutionProvider_CUDA(cuda_options);
    }
    sessionOptions.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_EXTENDED);
	Ort::Session session(env, onnxModelFilename, sessionOptions);
	Ort::IoBinding io_binding(session);

	Ort::AllocatorWithDefaultOptions allocator;


    cv::VideoCapture cap;
	cap.open(deviceID);
	if (!cap.isOpened()) {
		printf("can not open device %d\n", deviceID);
		return 0;
	}

	cap.set(cv::CAP_PROP_FRAME_WIDTH, size.width);
	cap.set(cv::CAP_PROP_FRAME_HEIGHT, size.height);

	printf("create tensors\n");


	Ort::MemoryInfo memoryInfo = Ort::MemoryInfo::CreateCpu(OrtAllocatorType::OrtArenaAllocator, OrtMemType::OrtMemTypeDefault);

	//Not sure if this really allocate on the GPU, there is currently no documentation on it...
	Ort::MemoryInfo memoryInfoCuda("Cuda", OrtAllocatorType::OrtDeviceAllocator, 0, OrtMemType::OrtMemTypeDefault);

	std::vector<float> src_data(size.width * size.height * 3);
	std::vector<int64_t> src_dims = {1, 3, size.height, size.width};
	Ort::Value src_tensor = Ort::Value::CreateTensor<float>(memoryInfo, src_data.data(), src_data.size(), src_dims.data(), 4);

	float downsample_ratio = 0.25f;
	int64_t downsample_ratio_dims[] = {1};
	Ort::Value downsample_ratio_tensor = Ort::Value::CreateTensor<float>(memoryInfo, &downsample_ratio, 1, downsample_ratio_dims, 1);

	float rec_data = 0.0f;
	int64_t rec_dims[] = {1, 1, 1, 1};
	Ort::Value r1i = Ort::Value::CreateTensor<float>(memoryInfo, &rec_data, 1, rec_dims, 4);

	io_binding.BindOutput("fgr", memoryInfoCuda);
	io_binding.BindOutput("pha", memoryInfo);
	io_binding.BindOutput("r1o", memoryInfoCuda);
	io_binding.BindOutput("r2o", memoryInfoCuda);
	io_binding.BindOutput("r3o", memoryInfoCuda);
	io_binding.BindOutput("r4o", memoryInfoCuda);

	io_binding.BindInput("r1i", r1i);
    io_binding.BindInput("r2i", r1i);
    io_binding.BindInput("r3i", r1i);
    io_binding.BindInput("r4i", r1i);
    io_binding.BindInput("downsample_ratio", downsample_ratio_tensor);

	printf("start\n");
    cv::Mat frame;
    while(true)
	{
		cap.read(frame);
		if (frame.empty()) {
            printf("error : empty frame grabbed");
            break;
        }

        cv::Mat blobMat;
        cv::dnn::blobFromImage(frame, blobMat);

        src_data.assign(blobMat.begin<float>(), blobMat.end<float>());
        for(size_t i = 0; i < src_data.size(); i++)
        	src_data[i] /= 255;

        io_binding.BindInput("src", src_tensor);
        session.Run(Ort::RunOptions{nullptr}, io_binding);

        std::vector<std::string> outputNames = io_binding.GetOutputNames();
        std::vector<Ort::Value> outputValues = io_binding.GetOutputValues();

        cv::Mat mask(size.height, size.width, CV_8UC1);
        for(int i = 0; i < outputNames.size(); i++) {
        	if(outputNames[i] == "pha") {
        		const cv::Mat outputImg(size.height, size.width, CV_32FC1, const_cast<float*>(outputValues[i].GetTensorData<float>()));
        		outputImg.convertTo(mask, CV_8UC1, 255.0);
        	} else if(outputNames[i] == "r1o") {
        		io_binding.BindInput("r1i", outputValues[i]);
        	} else if(outputNames[i] == "r2o") {
        		io_binding.BindInput("r2i", outputValues[i]);
        	} else if(outputNames[i] == "r3o") {
        		io_binding.BindInput("r3i", outputValues[i]);
        	} else if(outputNames[i] == "r4o") {
        		io_binding.BindInput("r4i", outputValues[i]);
        	}
        }
        cv::Mat img;
        cv::bitwise_and(frame, frame, img, mask);
        cv::imshow("img", img);
        cv::imshow("mask", mask);
        int key = cv::waitKey(10);
        if(key > 0)
        	break;


	}
    return 0;
}