[C++] onnx 결과값 Mat 타입으로 변경

segmentation 모델 결과를 onnx로 받아보려고 시도하다가 잘 안되었는데,

방법을 찾아 공유하고자 글을 쓴다.

 

나의 경우는 segmentation 모델 결과가 128 x 128 x 1 의 float 형태의 결과였다.

따라서 128 x 128 크기의 Mat에 _result 포인터를 넣었더니 해결되었고,

0~1사이 값이므로 255를 곱해 8bit로 변경하였더니 결과가 잘 나왔다.

auto output_tensors = session_.Run(Ort::RunOptions{ nullptr }, input_names.data(), inputTensors.data(), 1, output_names.data(), 1);
float* _result = output_tensors.front().GetTensorMutableData<float>(); // result 부분을 float 포인터로 받음
cv::Mat1f result_1f = cv::Mat1f(128, 128, _result); // pointer 부분을 Float Mat으로 변경
result_1f *= 255;
result_1f.convertTo(result, CV_8U);

 

아래는 웹캠으로부터 segmentation을 진행하는 전체 코드이다.

int main(void) {
	cv::VideoCapture cap(1);
	Mat frame, resized_img, image, input_mat, result;
	//std::unique_ptr<ONNX_Model> model;
	int width_ = 128;
	int height_= 128;
	
	std::cout << "Init Model" << std:: endl;
	Ort::Env env;
	Ort::SessionOptions session_options;
	Ort::Session session_{ env, L"my_model.onnx", session_options};
	std::vector<const char*> input_names = { "input" };
	std::vector<const char*> output_names = { "output" };

	std::cout << "Init Model Success" << std::endl;
	std::vector<int64_t> inputDims = { 1, 3, width_, height_ };
	std::vector<int64_t> outputDims = { 1, width_, height_ };


	size_t input_tensor_size = 1 * 3 * width_ * height_;
	std::vector<float> inputTensorValues(input_tensor_size);

	size_t output_tensor_size = 1 * width_ * height_;
	std::vector<float> outputTensorValues(output_tensor_size);

	
	std::vector<Ort::Value> inputTensors;
	std::vector<Ort::Value> outputTensors;
	Ort::MemoryInfo memoryInfo = Ort::MemoryInfo::CreateCpu(
		OrtAllocatorType::OrtArenaAllocator, OrtMemType::OrtMemTypeDefault);
	
	
	outputTensors.push_back(Ort::Value::CreateTensor<float>(
		memoryInfo, outputTensorValues.data(), output_tensor_size,
		outputDims.data(), outputDims.size()));
	while (cap.isOpened()) {
		cap >> frame;
		frame.copyTo(image);
		cv::resize(image, resized_img, Size(128, 128));
		cv::cvtColor(resized_img, resized_img, COLOR_RGB2BGR);
		resized_img.convertTo(resized_img, CV_32F, 1.0 / 255);
		cv::dnn::blobFromImage(resized_img, input_mat);
		clock_t start_time = clock();
		inputTensorValues.assign(input_mat.begin<float>(),
			input_mat.end<float>());
		inputTensors.push_back(Ort::Value::CreateTensor<float>(
			memoryInfo, inputTensorValues.data(), input_tensor_size, inputDims.data(),
			inputDims.size()));
		
		//session_.Run(Ort::RunOptions{ nullptr }, input_names, &input_tensor_, 1, output_names, &output_tensor_, 1);
		auto output_tensors = session_.Run(Ort::RunOptions{ nullptr }, input_names.data(), inputTensors.data(), 1, output_names.data(), 1);
		float* _result = output_tensors.front().GetTensorMutableData<float>();
		cv::Mat1f result_1f = cv::Mat1f(128, 128, _result);
		result_1f *= 255;
		result_1f.convertTo(result, CV_8U);
		clock_t end_time = clock();
		std::cout << (end_time - start_time) / 1000.0 << std::endl;		
		cv::imshow("original", frame);
		cv::imshow("Result", result);
		cv::waitKey(1);
		
	}
}

 

 

 

참고

 

https://leimao.github.io/blog/ONNX-Runtime-CPP-Inference/