#include <iostream>
#include<opencv2/opencv.hpp>

#include<math.h>
#include "yolov8_onnx.h"
#include "yolov8_pose_onnx.h"
#include<time.h>
//#define  VIDEO_OPENCV //if define, use opencv for video.

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


template<typename _Tp>
int yolov8_onnx(_Tp& task, cv::Mat& img, std::string& model_path)
{

    if (task.ReadModel(model_path, true)) {
		std::cout << "read net ok!" << std::endl;
	}
	else {
        std::cerr << "fail to read net!" << std::endl;
		return -1;
	}
	//生成随机颜色
	std::vector<cv::Scalar> color;
	srand(time(0));
	for (int i = 0; i < 80; i++) {
		int b = rand() % 256;
		int g = rand() % 256;
		int r = rand() % 256;
		color.push_back(cv::Scalar(b, g, r));
	}
	bool isPose = false;
	if (typeid(task) == typeid(Yolov8PoseOnnx)) {
		isPose = true;
	}
	PoseParams poseParams;

	std::vector<OutputParams> result;
	if (task.OnnxDetect(img, result)) {
		if (isPose)
			DrawPredPose(img, result, poseParams);
		else
			DrawPred(img, result, task._className, color);
	}
	else {
		std::cout << "Detect Failed!" << std::endl;
	}

    cv::imshow("1", img);
    cv::waitKey(0);
    system("pause");
	return 0;
}


template<typename _Tp>
int video_demo(_Tp& task, std::string& model_path)
{
	std::vector<cv::Scalar> color;
	srand(time(0));
	for (int i = 0; i < 80; i++) {
		int b = rand() % 256;
		int g = rand() % 256;
		int r = rand() % 256;
		color.push_back(cv::Scalar(b, g, r));
	}
	std::vector<OutputParams> result;
	cv::VideoCapture cap(0);
	if (!cap.isOpened())
	{
		std::cout << "open capture failured!" << std::endl;
		return -1;
	}
	cv::Mat frame;
	bool isPose = false;
	PoseParams poseParams;
#ifdef VIDEO_OPENCV
	cv::dnn::Net net;
	if (typeid(task) == typeid(Yolov8Pose)) {
		isPose = true;
	}
	if (task.ReadModel(net, model_path, true)) {
		std::cout << "read net ok!" << std::endl;
	}
	else {
		std::cout << "read net failured!" << std::endl;
		return -1;
	}

#else
	if (typeid(task) == typeid(Yolov8PoseOnnx)) {
		isPose = true;
	}
	if (task.ReadModel(model_path, true)) {
		std::cout << "read net ok!" << std::endl;
	}
	else {
		std::cout << "read net failured!" << std::endl;
		return -1;
	}

#endif

	while (true)
	{

		cap.read(frame);
		if (frame.empty())
		{
			std::cout << "read to end" << std::endl;
			break;
		}
		result.clear();
#ifdef VIDEO_OPENCV

		if (task.Detect(frame, net, result)) {

			if (isPose)
				DrawPredPose(frame, result, poseParams,true);
			else
				DrawPred(frame, result, task._className, color,true);
	
		}
#else
		if (task.OnnxDetect(frame, result)) {
			if (isPose)
                DrawPredPose(frame, result, poseParams);
			else
                DrawPred(frame, result, task._className, color);
		}
#endif
            cv::imshow("1", frame);
		int k = waitKey(10);
		if (k == 27) { //esc 
			break;
		}

	}
	cap.release();

	system("pause");

	return 0;
}

#define Pos_Flag

int main() {

#ifdef Pos_Flag
     std::string img_path = "/home/sxyx2026/cdq_data/imgs/2/2_2026-01-22_21-48-33-291.jpg";
    std::string model_path = "/home/sxyx2026/cdq_arrow.onnx";

    cv::Mat src = imread(img_path);
    cv::Mat img = src.clone();

    Yolov8PoseOnnx		task_pose_ort;
    yolov8_onnx(task_pose_ort, img, model_path); //yolov8 onnxruntime pose

#else

    std::string img_path = "/home/sxyx2026/cdq_data/imgs/0/0_2026-01-22_14-35-48-249.jpg";

    std::string model_path = "/home/sxyx2026/cdq_grab.onnx";

    cv::Mat src = imread(img_path);
    cv::Mat img = src.clone();

    Yolov8Onnx			task_detect_ort;
    yolov8_onnx(task_detect_ort,img,model_path);  //yoolov8 onnxruntime detect
#endif
	return 0;
}