ImageRecognition/CircularArresterCurrentALGO.cpp

#include "CircularArresterCurrentALGO.h"
#define _CRT_SECURE_NO_WARNINGS
#include <iostream>
#include <fstream>
#include <string>
#include <math.h>
#include <cmath>
#include <opencv2/imgproc.hpp>
#include <opencv2/highgui.hpp>
#include <opencv2/highgui/highgui.hpp>
#include "opencv2/imgproc/types_c.h"
#include <onnxruntime_cxx_api.h>
#include <vector>

#define SEG_IMAGE_SIZE 512
#define LINE_HEIGH 120
#define LINE_WIDTH 1600
#define CIRCLE_RADIUS 250
#define METER_RANGE 3.0

using namespace cv;
using namespace std;
using namespace Ort;

void CircularArresterCurrentALGO::Init()
{
	string model_path = "models/ArresterMonitor.onnx";
	std::wstring widestr = std::wstring(model_path.begin(), model_path.end());
	sessionOptions.SetGraphOptimizationLevel(ORT_ENABLE_BASIC);
	ort_session = new Session(env, widestr.c_str(), sessionOptions);
	size_t numInputNodes = ort_session->GetInputCount();
	size_t numOutputNodes = ort_session->GetOutputCount();
	AllocatorWithDefaultOptions allocator;
	for (int i = 0; i < numInputNodes; i++)
	{
		input_names.push_back(ort_session->GetInputName(i, allocator));
		Ort::TypeInfo input_type_info = ort_session->GetInputTypeInfo(i);
		auto input_tensor_info = input_type_info.GetTensorTypeAndShapeInfo();
		auto input_dims = input_tensor_info.GetShape();
		input_node_dims.push_back(input_dims);
	}
	for (int i = 0; i < numOutputNodes; i++)
	{
		output_names.push_back(ort_session->GetOutputName(i, allocator));
		Ort::TypeInfo output_type_info = ort_session->GetOutputTypeInfo(i);
		auto output_tensor_info = output_type_info.GetTensorTypeAndShapeInfo();
		auto output_dims = output_tensor_info.GetShape();
		output_node_dims.push_back(output_dims);
	}
	this->inpHeight = input_node_dims[0][2];
	this->inpWidth = input_node_dims[0][3];
	this->outHeight = output_node_dims[0][2];
	this->outWidth = output_node_dims[0][3];
}

float CircularArresterCurrentALGO::detect(Mat& srcimg)
{
	vector<float> input_image_ = { 1, 3, 512, 512 };   //1, 3, 512, 512
	Mat dstimg;
	Size resize_size(input_image_[2], input_image_[3]);
	resize(srcimg, dstimg, resize_size, 0, 0, cv::INTER_LINEAR);
	int channels = dstimg.channels();
	input_image_.resize((this->inpWidth * this->inpHeight * dstimg.channels()));

	for (int c = 0; c < channels; c++)
	{
		for (int i = 0; i < this->inpHeight; i++)
		{
			for (int j = 0; j < this->inpWidth; j++)
			{
				float pix = dstimg.ptr<uchar>(i)[j * 3 + 2 - c];
				input_image_[(c * this->inpHeight * this->inpWidth + i * this->inpWidth + j)] = (pix / 255.0 - mean[c]) / stds[c];
			}
		}
	}

	array<int64_t, 4> input_shape_{ 1, 3, this->inpHeight, this->inpWidth };
	auto allocator_info = MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeCPU);
	Value input_tensor_ = Value::CreateTensor<float>(allocator_info, input_image_.data(), input_image_.size(), input_shape_.data(), input_shape_.size());

	vector<Value> ort_outputs = ort_session->Run(RunOptions{ nullptr }, &input_names[0], &input_tensor_, 1, output_names.data(), output_names.size());
	float *pred = ort_outputs[0].GetTensorMutableData<float>();
	Mat result(outHeight, outWidth, CV_32FC1, pred);
	result = 6 - result;
	result *= 255;
	result.convertTo(result, CV_8UC1);

	//namedWindow("分割", WINDOW_NORMAL);
	//imshow("分割", result);
	//waitKey(1);

	Mat binary;
	threshold(result, binary, 150, 255, THRESH_BINARY);//二值化阈值处理

	//形态学变换
	Mat Sobel_Y_thres;
	Mat element = cv::getStructuringElement(MORPH_RECT, Size(5, 5));
	morphologyEx(binary, Sobel_Y_thres, cv::MORPH_OPEN, element, Point(-1, -1), 2);
	//cv::imshow("test", Sobel_Y_thres);
	//cv::waitKey(3000);
	//cv::destroyWindow("test");
	//查找边界轮廓
	vector<vector<Point>> contours;
	vector<Vec4i> hierarchy;
	float pointerArea = 0;
	float scaleArea = 0;
	vector<int>numArea;
	
	findContours(Sobel_Y_thres, contours, hierarchy, RETR_TREE, CHAIN_APPROX_SIMPLE, Point());
	for (size_t t = 0; t < contours.size(); t++)
	{
		drawContours(dstimg, contours, -1, Scalar(0, 0, 255), 2, 8);
		double area = contourArea(contours[t]);
		//计算指针和刻度关系
		numArea.push_back(area);
		float maxValue = 0;
		float minValue = 0;
		//scale
		maxValue = *max_element(numArea.begin(), numArea.end());
		scaleArea = maxValue;
		//pointer
		minValue = *min_element(numArea.begin(), numArea.end());
		if (numArea.size() == 1)
		{
			pointerArea = 0;
		}
		else
		{
			pointerArea = minValue;
		}
	}

	//计算表盘读数
	cout << "指针：" << pointerArea << "  刻度：" << scaleArea << endl;
	float result_ratio = (1.0 * pointerArea / (pointerArea + scaleArea));
	float resutValue = (result_ratio * METER_RANGE);
	if (resutValue>0.0f)
	{
		resutValue += 0.05;
	}
	cout << "result_ratio:" << result_ratio << "  result_value:" << resutValue << endl;
	//resutValue = result_value;
	//return result_value;
	return resutValue;
}