You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

134 lines
2.3 KiB

/*********版权所有C2022武汉高德红外股份有限公司***************
* YOLO.h
* yolo
* yolo
*
* V1.0
* 04046wcw
* 20220408
*
* 1
*
*
*
*
* 2
*******************************************************************/
#ifndef YOLO_H
#define YOLO_H
#pragma once
#include "Arith_CommonDef.h"
#include <opencv2/opencv.hpp>
#include <opencv2/dnn.hpp>
#include <vector>
#include <string>
#include <thread>
using std::string;
#ifdef _WIN32
#include <windows.h>
#else
#include <unistd.h>
#endif
#include <ctime>
#include <iostream>
void SLEEP(int ms);
using cv::dnn::Net;
struct Detection
{
int class_id;
float confidence;
cv::Rect box;
};
struct YoloParam
{
float score_threshold;
float nms_threshold;
float conf_threshold;
bool bUseCuda;
};
class YOLO_ONNX_Detector
{
public:
YOLO_ONNX_Detector(){};
virtual ~YOLO_ONNX_Detector(){};
virtual int detect(cv::Mat& srcimg) = 0;
virtual void SetParam(YoloParam para){m_para = para;};
virtual std::vector<Detection>& getTargetVec(){return m_targetArray;};
virtual std::string getClassName(int ID){return g_YoloClassName[ID];};
protected:
cv::dnn::Net m_YoloNet;
std::vector<std::string> g_YoloClassName;
std::vector<Detection> m_targetArray;
YoloParam m_para;
std::string NetPath;
};
class yolox:public YOLO_ONNX_Detector
{
public:
yolox(std::string modelpath, std::string namesConfig, YoloParam para);
int detect(cv::Mat& srcimg);
private:
const int stride[3] = { 8, 16, 32 };
//const string classesFile = "coco.names";
const int input_shape[2] = { 640, 640 }; //// height, width
const double mean[3] = { 0.485, 0.456, 0.406 };
const double std[3] = { 0.229, 0.224, 0.225 };
int num_class;
cv::Mat resize_image(cv::Mat srcimg, float* scale);
void normalize(cv::Mat& srcimg);
int get_max_class(float* scores);
};
class yolo5 :public YOLO_ONNX_Detector
{
const int INPUT_WIDTH = 1024;
const int INPUT_HEIGHT = 1024;
public:
yolo5(std::string modelpath, std::string namesConfig, YoloParam para);
int detect(cv::Mat& srcimg);
private:
cv::Mat keepRatio(const cv::Mat &source);
};
#endif