generated from zhangwei/Matrixai
99 lines
3.9 KiB
C++
99 lines
3.9 KiB
C++
#ifndef YOLOV5_COMMON_H_
|
|
#define YOLOV5_COMMON_H_
|
|
|
|
#include <fstream>
|
|
#include <map>
|
|
#include <sstream>
|
|
#include <vector>
|
|
#include <tuple>
|
|
#include <opencv2/opencv.hpp>
|
|
#include "NvInfer.h"
|
|
#include "yololayer.h"
|
|
|
|
using namespace std;
|
|
using namespace nvinfer1;
|
|
|
|
#define LOCATIONS 4
|
|
|
|
struct alignas(float) stDetection { //alignas(float)
|
|
float bbox[LOCATIONS];
|
|
float class_conf;
|
|
float clear_conf;
|
|
int class_id;
|
|
int clear_id;
|
|
};
|
|
|
|
struct alignas(float) yolov5Detection { //alignas(float)
|
|
float bbox[LOCATIONS];
|
|
float class_conf;
|
|
int class_id;
|
|
};
|
|
|
|
|
|
//HSV2BGR
|
|
std::tuple<uint8_t, uint8_t, uint8_t> hsv2bgr(float h, float s, float v);
|
|
|
|
//HSV2RGB
|
|
std::tuple<uint8_t, uint8_t, uint8_t> hsv2rgb(float h, float s, float v);
|
|
|
|
//随机颜色生成
|
|
std::tuple<uint8_t, uint8_t, uint8_t> random_color(int id);
|
|
|
|
//center x, y, w, h to x1,y1,x2,y2
|
|
void xywh2xyxy(float *xywh, float * xyxy);
|
|
|
|
//获取区域框1
|
|
cv::Rect get_rect(cv::Mat& img, float bbox[4]);
|
|
|
|
//获取区域框2
|
|
cv::Rect get_rectangle(cv::Mat& img, float bbox[4]);
|
|
|
|
//计算IOU
|
|
float iou(float lbox[4], float rbox[4]);
|
|
|
|
//比较两者置信度
|
|
bool cmp(const Yolo::Detection& a, const Yolo::Detection& b);
|
|
|
|
//非极大值抑制,默认0.5
|
|
void nms(std::vector<Yolo::Detection>& res, float *output, float conf_thresh, float nms_thresh);
|
|
|
|
//获取缩放比例
|
|
float GetResizeRatio(unsigned int img_width, unsigned int img_height, unsigned int model_width, unsigned int model_height);
|
|
|
|
//坐标还原
|
|
void UpperVertexResetLocation(float fResizeRatio, unsigned int orig_width, unsigned int orig_height, stDetection &detection); //左上顶点补边方式坐标还原
|
|
void CenterResetLocation(float fResizeRatio, unsigned int orig_width, unsigned int orig_height, unsigned int input_w, unsigned int input_h, stDetection &detection); //中心补边方式坐标还原
|
|
|
|
//yolov5带清晰度后处理
|
|
void yolov5_clearity_decode_opencv_nms(std::vector<stDetection>& res, float *output, int outSize, unsigned int detSize, unsigned int classNum, unsigned int clearNum, float confThresh, float nmsThresh);
|
|
|
|
//yolov5后处理
|
|
void yolov5_decode_opencv_nms(std::vector<stDetection>& res, float *output, int outSize, unsigned int detSize, unsigned int classNum, float confThresh, float nmsThresh);
|
|
|
|
// TensorRT weight files have a simple space delimited format:
|
|
// [type] [size] <data x size in hex>
|
|
std::map<std::string, Weights> loadWeights(const std::string file);
|
|
|
|
IScaleLayer* addBatchNorm2d(INetworkDefinition *network, std::map<std::string, Weights>& weightMap, ITensor& input, std::string lname, float eps) ;
|
|
|
|
ILayer* convBlock(INetworkDefinition *network, std::map<std::string, Weights>& weightMap, ITensor& input, int outch, int ksize, int s, int g, std::string lname);
|
|
|
|
ILayer* focus(INetworkDefinition *network, std::map<std::string, Weights>& weightMap, ITensor& input, int inch, int outch, int ksize, std::string lname);
|
|
|
|
ILayer* bottleneck(INetworkDefinition *network, std::map<std::string, Weights>& weightMap, ITensor& input, int c1, int c2, bool shortcut, int g, float e, std::string lname);
|
|
|
|
ILayer* bottleneckCSP(INetworkDefinition *network, std::map<std::string, Weights>& weightMap, ITensor& input, int c1, int c2, int n, bool shortcut, int g, float e, std::string lname);
|
|
|
|
ILayer* C3(INetworkDefinition *network, std::map<std::string, Weights>& weightMap, ITensor& input, int c1, int c2, int n, bool shortcut, int g, float e, std::string lname);
|
|
|
|
ILayer* SPP(INetworkDefinition *network, std::map<std::string, Weights>& weightMap, ITensor& input, int c1, int c2, int k1, int k2, int k3, std::string lname);
|
|
|
|
ILayer* SPPF(INetworkDefinition *network, std::map<std::string, Weights>& weightMap, ITensor& input, int c1, int c2, int k, std::string lname);
|
|
|
|
std::vector<std::vector<float>> getAnchors(std::map<std::string, Weights>& weightMap, std::string lname);
|
|
|
|
IPluginV2Layer* addYoLoLayer(INetworkDefinition *network, std::map<std::string, Weights>& weightMap, std::string lname, std::vector<IConvolutionLayer*> dets);
|
|
|
|
#endif
|
|
|