// 自定义配置结构
struct Configuration
{
public:
	float confThreshold; // Confidence threshold
	float nmsThreshold;  // Non-maximum suppression threshold
	float objThreshold;  // Object Confidence threshold
	string modelpath;
};

// 模型
class YOLOv5
{
public:
	// 初始化
	YOLOv5(Configuration config, bool isCuda);
	void detect(Mat &frame);  // 检测函数
private:
	float confThreshold;
	float nmsThreshold;
	float objThreshold;
	int inpWidth;
	int inpHeight;
	int num_classes;
	string classes[1] = { "eleb" };
	//vector<string> class_names;
	const bool keep_ratio = true;
	Net net;   // dnn里的
	void drawPred(float conf, int left, int top, int right, int bottom, Mat &frame, int classid);
	Mat resize_image(Mat srcimg, int *newh, int *neww, int *top, int *left);
};

Mat dstimg = this->resize_image(frame, &newh, &neww, &padh, &padw);
Mat blob = blobFromImage(dstimg, 1 / 255.0, Size(this->inpWidth, this->inpHeight), Scalar(0, 0, 0), true, false);
this->net.setInput(blob);

vector<Mat> outs;
this->net.forward(outs, this->net.getUnconnectedOutLayersNames());

for (int i = 0; i < num_proposal; ++i) 
{
    // ...
    if (obj_conf > this->objThreshold)
    {
        // ...
        if (max_class_socre > this->confThreshold)
        {
            // ...
            confidences.push_back((float) max_class_socre);
            boxes.push_back(Rect(left, top, (int) (w * ratiow), (int) (h * ratioh)));
            classIds.push_back(class_idx);
        }
    }
}

vector<int> indices;
dnn::NMSBoxes(boxes, confidences, this->confThreshold, this->nmsThreshold, indices);
for (size_t i = 0; i < indices.size(); ++i)
{
    // ...
    this->drawPred(confidences[idx], box.x, box.y, box.x + box.width, box.y + box.height, frame, classIds[idx]);
}

// 绘制预测目标的边界框
void YOLOv5::drawPred(float conf, int left, int top, int right, int bottom, Mat &frame, int classid)
{
	// 绘制一个显示边界框的矩形
	rectangle(frame, Point(left, top), Point(right, bottom), Scalar(0, 0, 255), 10);

	// 得到目标的类名和置信度
	string label = cv::format("%.2f", conf);
	label = this->classes[classid] + ":" + label;

	// 绘制目标的类名和置信度文本
	int baseLine;
	Size labelSize = getTextSize(label, FONT_HERSHEY_SIMPLEX, 0.5, 1, &baseLine);
	top = max(top, labelSize.height);
	putText(frame, label, Point(left, top), FONT_HERSHEY_SIMPLEX, 0.75, Scalar(0, 255, 0), 1);
}

//yolov5-onnx.cpp
int main(int argc, const char *argv[])
{
	if (argc != 3)
	{
		std::cout << "Run like this:\n    " << argv[0] << " yolov5s.onnx input.jpg" << std::endl;
		return -1;
	}
	
	Configuration yolo_nets = { 0.3, 0.5, 0.3, argv[1] }; // 初始化数据
	YOLOv5 yolo_model(yolo_nets, false);
	Mat srcimg = imread(argv[2]);  // 读入图片
	yolo_model.detect(srcimg);     // 进行推理
	imwrite("result.jpg", srcimg); // 输出图片
	return 0;
}

struct Result
{
	float score;
	cv::Rect box;
	int class_id;
};

// TensorRT需要日志
// 参考: https://github.com/cyrusbehr/tensorrt-cpp-api
using Severity = nvinfer1::ILogger::Severity;
class Logger : public nvinfer1::ILogger
{
	void log(Severity severity, const char *msg) noexcept override;
};

void Logger::log(Severity severity, const char *msg) noexcept
{
	if (severity <= Severity::kWARNING)
	{
		std::cout << msg << std::endl;
	}
}

Logger logger;

size_t get_memory_size(const nvinfer1::Dims &dims, const int32_t elem_size)
{
	return std::accumulate(dims.d, dims.d + dims.nbDims, 1, std::multiplies<int64_t>()) * static_cast<size_t>(elem_size);
}

// 读取TensorRT引擎
std::ifstream engine_file(engine_file_path, std::ios::binary);
if (engine_file.fail())
{
	std::cout << "Failed to read model file." << std::endl;
}
engine_file.seekg(0, std::ifstream::end);
auto fsize = engine_file.tellg();
engine_file.seekg(0, std::ifstream::beg);
std::vector<char> engineData(fsize);
engine_file.read(engineData.data(), fsize);

// 初始化IRuntime和ICudaEngine
std::unique_ptr<nvinfer1::IRuntime> runtime{ nvinfer1::createInferRuntime(logger) };
std::unique_ptr<nvinfer1::ICudaEngine> mEngine{ runtime->deserializeCudaEngine(engineData.data(), fsize) };
if (mEngine.get() == nullptr)
{
	std::cout << "Failed to deserialize CUDA engine." << std::endl;
}

std::unique_ptr<nvinfer1::IExecutionContext> context{ mEngine->createExecutionContext() };
if (context.get() == nullptr)
{
	std::cout << "Failed to create CUDA context." << std::endl;
}

// 记录与推演相关的内存区域
void *bindings[5];
const char *input_name = "images";
int32_t input_idx = mEngine->getBindingIndex(input_name);
if (input_idx == -1)
{
	std::cout << "ERROR: failed to get input by name: " << input_name << std::endl;
}

// 计算显存大小并开辟
int32_t channels = 3, height = 640, width = 640;
nvinfer1::Dims4 input_dims { 1, channels, height, width };
size_t input_mem_size = get_memory_size(input_dims, sizeof(float));
void* cuda_mem_input { nullptr };
if (cudaMalloc(&cuda_mem_input, input_mem_size) != cudaSuccess) 
{
    std::cout << "ERROR: input cuda memory allocation failed, size = " << input_mem_size << " bytes" << std::endl;
}

bindings[0] = cuda_mem_input;

std::vector<std::string> output_node_names{ "output0" };
std::vector<size_t> output_mem_sizes;
bool output_mem_initialized = true;
for (size_t i = 0; i < output_node_names.size(); i++)
{
	int32_t output_idx = mEngine->getBindingIndex(output_node_names[i].c_str());
	if (output_idx == -1)
	{
		std::cout << "ERROR: failed to get output by name: " << output_node_names[i] << std::endl;
		output_mem_initialized = false;
		break;
	}
	auto output_dims = context->getBindingDimensions(output_idx);
	auto output_size = get_memory_size(output_dims, sizeof(float));
	output_mem_sizes.push_back(output_size);
	void *cuda_mem_output{ nullptr };
	if (cudaMalloc(&cuda_mem_output, output_size) != cudaSuccess)
	{
		std::cout << "ERROR: output cuda memory allocation failed, size = " << output_size << " bytes" << std::endl;
		output_mem_initialized = false;
		break;
	}
	else
	{
		bindings[1 + i] = cuda_mem_output;
	}
}

cudaStream_t stream{ nullptr };
if (cudaStreamCreate(&stream) != cudaSuccess)
{
	std::cout << "ERROR: cuda stream creation failed." << std::endl;
}

cv::Mat img_bgr = cv::imread(image_file_path);

// 预处理
cv::resize(img_bgr, img_bgr, cv::Size(width, height));
int32_t tmp = height * width * channels;
float *input_buffer = new float[tmp]();
for (int c = 0; c < channels; c++)
{
	for (int j = 0, HW = height * width; j < HW; ++j)
	{
		input_buffer[c * HW + j] = static_cast<float>(img_bgr.data[j * channels + 2 - c]) / 255.0f;
		// 其中 img_bgr.data 为指向像素存储区域的指针，而括号内下表运算中的 +2 实现了BGR数据到RGB的转换。
	}
}

if (cudaMemcpyAsync(cuda_mem_input, input_buffer, input_mem_size, cudaMemcpyHostToDevice, stream) != cudaSuccess)
{
	std::cout << "ERROR: CUDA memory copy of input failed, size = " << input_mem_size << " bytes" << std::endl;
}

bool status = context->enqueueV2(bindings, stream, nullptr);
if (!status)
{
	std::cout << "ERROR: TensorRT inference failed." << std::endl;
}

std::vector<float *> output_buffers;
for (size_t i = 0; i < output_mem_sizes.size(); i++)
{
	float *buf = new float[output_mem_sizes[i] / sizeof(float)];
	output_buffers.push_back(buf);
}

for (size_t i = 0; i < output_mem_sizes.size(); i++)
{
	auto mem_to_host_result = cudaMemcpyAsync(output_buffers[i], bindings[1 + i], output_mem_sizes[i], cudaMemcpyDeviceToHost, stream);
	if (mem_to_host_result != cudaSuccess)
	{
		std::cout << "ERROR: CUDA memory copy of output " << i << " failed, size = " << output_mem_sizes[i] << " bytes" << std::endl;
	}
}

cudaStreamSynchronize(stream);

std::vector<float> scores;
std::vector<cv::Rect> boxes;
std::vector<int> class_ids;

float *p = output_buffers[0]; // 指向output0的第一个元素
int step = 6, proposal_count = 25200;
float score_threshold = 0.5, nms_threshold = 0.45; // 分支阈值和NMS阈值
float scale = 1.0;
for (size_t i = 0; i < proposal_count; i++)
{
	// 置信度低于分数阈值
	float obj_score = p[4]; // 第4个元素是预测框分值
	if (obj_score < score_threshold)
	{
		p += step;
		continue;
	}

	// 判断对象分类
	int c_id = -1;
	float c_score = 0;
	for (size_t j = 1; j < step; j++)
	{
		if (p[j] > c_score)
		{
			c_score = p[j];
			c_id = j - 2;
		}
	}
	
	scores.push_back(c_score * obj_score);
	class_ids.push_back(c_id);
	boxes.push_back(cv::Rect((p[0] - p[2] / 2) / scale, (p[1] - p[3] / 2) / scale, p[2] / scale, p[3] / scale));
	p += step; // 指向下一行
}

// 执行非最大值抑制以消除具有较低置信度的冗余重叠框
std::vector<int> indices;
cv::dnn::NMSBoxes(boxes, scores, score_threshold, nms_threshold, indices);

// 收集检测结果
std::vector<Result> results;
for (auto i : indices) 
{
    results.push_back(Result { scores[i], boxes[i], class_ids[i] });
}

for (auto &r : results)
{
	cv::rectangle(img_bgr, r.box, cv::Scalar(0, 255, 255), 5);
}

cv::imwrite("result.jpg", img_bgr);

//yolov5-engine.cpp
int main(int argc, char const *argv[])
{
	if (argc != 3)
	{
		std::cout << "Run like this:\n    " << argv[0] << " yolov5s.engine input.jpg" << std::endl;
		return -1;
	}
	
	//...
	run(argv[1], argv[2]);

	return 0;
}

//导出名冲突
string getUniqueFileName(const std::string &fileName, const std::string &outputDirectory)
{
	string baseName, extension;
	size_t dotIndex = fileName.rfind('.');
	if (dotIndex != string::npos)
	{
		baseName = fileName.substr(0, dotIndex);
		extension = fileName.substr(dotIndex);
	}
	else
	{
		baseName = fileName;
	}

	string newFileName = fileName;
	int counter = 1;
	while (ifstream(outputDirectory + newFileName).good())
	{
		newFileName = baseName + "_" + to_string(counter) + extension;
		counter++;
	}

	return newFileName;
}

// 备份存储
string dir = "output/"; //文件夹路径
if (_access(dir.c_str(), 0) == -1)//判断该文件夹是否存在
{
	_mkdir(dir.c_str());
}

imwrite("result.jpg", srcimg);
string newname = getUniqueFileName("result.jpg", "output/");
imwrite("output/" + newname, srcimg);