63 lines
1.4 KiB
C++
63 lines
1.4 KiB
C++
#ifndef YOLOPV2_H
|
|
#define YOLOPV2_H
|
|
|
|
#include <fstream>
|
|
#include <sstream>
|
|
#include <iostream>
|
|
#include <opencv2/imgproc.hpp>
|
|
#include <opencv2/highgui.hpp>
|
|
#include <onnxruntime_cxx_api.h>
|
|
|
|
using namespace cv;
|
|
using namespace Ort;
|
|
using namespace std;
|
|
|
|
struct Net_config
|
|
{
|
|
float confThreshold; // Confidence threshold
|
|
float nmsThreshold; // Non-maximum suppression threshold
|
|
string modelpath;
|
|
};
|
|
|
|
typedef struct BoxInfo
|
|
{
|
|
float x1;
|
|
float y1;
|
|
float x2;
|
|
float y2;
|
|
float score;
|
|
int label;
|
|
} BoxInfo;
|
|
|
|
class YOLOPv2
|
|
{
|
|
public:
|
|
YOLOPv2(Net_config config);
|
|
Mat detect(Mat& frame);
|
|
private:
|
|
int inpWidth;
|
|
int inpHeight;
|
|
int nout;
|
|
int num_proposal;
|
|
vector<string> class_names;
|
|
int num_class;
|
|
|
|
float confThreshold;
|
|
float nmsThreshold;
|
|
vector<float> input_image_;
|
|
void normalize_(Mat img);
|
|
void nms(vector<BoxInfo>& input_boxes);
|
|
const float anchors[3][6] = { {12, 16, 19, 36, 40, 28}, {36, 75, 76, 55, 72, 146},{142, 110, 192, 243, 459, 401} };
|
|
const float stride[3] = { 8.0, 16.0, 32.0 };
|
|
|
|
Env env = Env(ORT_LOGGING_LEVEL_ERROR, "YOLOPv2");
|
|
Ort::Session *ort_session = nullptr;
|
|
SessionOptions sessionOptions = SessionOptions();
|
|
vector<char*> input_names;
|
|
vector<char*> output_names;
|
|
vector<vector<int64_t>> input_node_dims; // >=1 outputs
|
|
vector<vector<int64_t>> output_node_dims; // >=1 outputs
|
|
};
|
|
|
|
#endif // YOLOPV2_H
|