#include <fstream>
#include <sstream>
#include <iostream>
#include <opencv2/dnn.hpp>
#include <opencv2/imgproc.hpp>
#include <opencv2/highgui.hpp>

using namespace cv;
using namespace dnn;
using namespace std;

struct Net_config
{
	float confThreshold; // Confidence threshold
	float nmsThreshold;  // Non-maximum suppression threshold
	float objThreshold;  //Object Confidence threshold
	string modelpath;
};

const float anchor[3][6] = { {10.0,  13.0, 16.0,  30.0,  33.0,  23.0},
								 {30.0,  61.0, 62.0,  45.0,  59.0,  119.0},
								 {116.0, 90.0, 156.0, 198.0, 373.0, 326.0} };
 
float* anchors;
int num_stride = 3;
int inpWidth = 1024;
int inpHeight= 1024;
vector<string> class_names;
int num_class;

float confThreshold = 0.2;
float nmsThreshold = 0.2;
float objThreshold = 0.2;
const bool keep_ratio = true;

 
Mat resize_image(Mat srcimg, int *newh, int *neww, int *top, int *left)
{
	int srch = srcimg.rows, srcw = srcimg.cols;
	*newh = inpHeight;
	*neww = inpWidth;
	Mat dstimg;
	if (keep_ratio && srch != srcw) {
		float hw_scale = (float)srch / srcw;
		if (hw_scale > 1) {
			*newh = inpHeight;
			*neww = int(inpWidth / hw_scale);
			resize(srcimg, dstimg, Size(*neww, *newh), INTER_AREA);
			*left = int((inpWidth - *neww) * 0.5);
			copyMakeBorder(dstimg, dstimg, 0, 0, *left,inpWidth - *neww - *left, BORDER_CONSTANT, 114);
		}
		else {
			*newh = (int)inpHeight * hw_scale;
			*neww = inpWidth;
			resize(srcimg, dstimg, Size(*neww, *newh), INTER_AREA);
			*top = (int)(inpHeight - *newh) * 0.5;
			copyMakeBorder(dstimg, dstimg, *top, inpHeight - *newh - *top, 0, 0, BORDER_CONSTANT, 114);
		}
	}
	else {
		resize(srcimg, dstimg, Size(*neww, *newh), INTER_AREA);
	}
	return dstimg;
}

void drawPred(Mat& frame, float conf, int classid, RotatedRect rrect)   // Draw the predicted bounding box
{

	cv::Point2f vertex[4];
	rrect.points(vertex);
	for (int i = 0; i < 4; i++)
	{
		cv::line(frame, vertex[i], vertex[(i + 1) % 4], cv::Scalar(0, 0,255), 1);
	}

	int xmin = (int)vertex[0].x;
	int ymin = (int)vertex[0].y - 10;
	string label = format("%.2f", conf);
	label =class_names[classid] + ":" + label;
	putText(frame, label, Point(xmin, ymin - 5), FONT_HERSHEY_SIMPLEX, 0.25, Scalar(0, 255, 0), 1);

}

void detect(Net net,Mat& frame)
{
	int newh = 0, neww = 0, padh = 0, padw = 0;
	Mat dstimg = resize_image(frame, &newh, &neww, &padh, &padw);
	Mat blob = blobFromImage(dstimg, 1 / 255.0, Size(inpWidth, inpHeight), Scalar(0, 0, 0), true, false);
	net.setInput(blob);
	vector<Mat> outs;
    net.forward(outs, net.getUnconnectedOutLayersNames());
	int num_proposal = outs[0].size[1];
	int nout = outs[0].size[2];

	if (outs[0].dims > 2)
	{
		outs[0] = outs[0].reshape(0, num_proposal);
	}
	/////generate proposals
	vector<float> confidences;
	vector<RotatedRect> boxes;
	vector<int> classIds;
	//vector<int> angles;
	float ratioh = (float)frame.rows / newh, ratiow = (float)frame.cols / neww;
	int n = 0, q = 0, i = 0, j = 0, row_ind = 0; ///xmin,ymin,xamx,ymax,box_score,class_score
	float* pdata = (float*)outs[0].data;
	for (n = 0; n < num_stride; n++)   ///特征图尺度
	{
		const float stride = pow(2, n + 3);
		int num_grid_x = (int)ceil(( inpWidth / stride));
		int num_grid_y = (int)ceil(( inpHeight / stride));
		for (q = 0; q < 3; q++)    ///anchor
		{
			const float anchor_w =  anchors[n * 6 + q * 2];
			const float anchor_h =  anchors[n * 6 + q * 2 + 1];
			for (i = 0; i < num_grid_y; i++)
			{
				for (j = 0; j < num_grid_x; j++)
				{
					float box_score = pdata[4];//box_score
					if (box_score >  objThreshold)
					{
						Mat tmp = outs[0].row(row_ind);    
						Mat scores = tmp.colRange(5, 21); 
  						Point classIdPoint;
						double max_class_socre;
						// Get the value and location of the maximum score
						minMaxLoc(scores, 0, &max_class_socre, 0, &classIdPoint);
						max_class_socre *= box_score;
						//21-200  角度,获得最大可能角度 
						Mat angleScores = tmp.colRange(21, nout);
						Point angleIdPoint;
						double max_angle_socre;
						minMaxLoc(angleScores, 0, &max_angle_socre, 0, &angleIdPoint);

						if (max_class_socre > confThreshold)
						{ 
							const int class_idx = classIdPoint.x;
							float cx = (pdata[0] * 2.f - 0.5f + j) * stride;  ///cx
							float cy = (pdata[1] * 2.f - 0.5f + i) * stride;   ///cy

							float w = powf(pdata[2] * 2.f, 2.f) * anchor_w;   ///w
							float h = powf(pdata[3] * 2.f, 2.f) * anchor_h;  ///h

							cx = (cx - padw)*ratiow;
							cy = (cy - padh)*ratioh;
							w *= ratiow;
							h *= ratioh;

							confidences.push_back((float)max_class_socre);
							classIds.push_back(class_idx);

							RotatedRect box = RotatedRect(Point2f(cx, cy), Size2f(w, h), 90-angleIdPoint.x);
							boxes.push_back(box);
			 
						}
					}
					row_ind++;
					pdata += nout; //下一个anchor
				}
			}
		}
	}

	// Perform non maximum suppression to eliminate redundant overlapping boxes with
	// lower confidences
	vector<int> indices;
	dnn::NMSBoxes(boxes, confidences,  confThreshold, nmsThreshold, indices);
	for (size_t i = 0; i < indices.size(); ++i)
	{
		int idx = indices[i];
		 drawPred(frame,confidences[idx], classIds[idx] , boxes[idx]);
	}
}

int main()
{
	dnn::Net net = readNet("weights1105/yolov5m_finetune_best_1114.onnx");
	ifstream ifs("dota.names");
	string line;
	while (getline(ifs, line)) class_names.push_back(line);
	num_class = class_names.size();


	string imgpath = "images/2022-10-29 15-19-25.png";
	Mat srcimg = imread(imgpath);
	anchors = (float*)anchor;

	detect(net,srcimg);

	static const string kWinName = "Deep learning object detection in OpenCV";
	namedWindow(kWinName, WINDOW_NORMAL);
	imshow(kWinName, srcimg);
	waitKey(0);
	destroyAllWindows();
}