opencv4部署darknet训练的yolo模型
记录一下代码,方便以后回看
darknet框架训练的yolov4会生成 weights权重文件,加上配置的.cfg文件和names分类文件用于配置
DraknetYolo.h
// An highlighted block
class DraknetYolo
{
public:
DraknetYolo();
~DraknetYolo();
/*
这个函数必须要opencv4版本以上才可以使用
初始化YOLO库,在检测之前必须进行初始化
@param modelConfiguration 模型cfg文件
@param modelBinary 模型weights权重文件
@param classNames 模型分类names文件
@res 返回0为初始化成功,其他为初始化失败
*/
int Init_YOLOV4(std::string modelConfiguration, std::string modelBinary, std::string classNames);
/*
检测一帧图像
@param image------输入图像,类型为Mat
@param confidenceThreshold------输入判断阈值,大于阈值的才会被检测到
@param res_classID------输出,返回检测到的每个对象的对应的分类ID是多少
@param res_confidences------输出,返回检测到的每个对象的对应的评分是多少
@param res_boxes------输出,返回检测到的每个对象的对应的在输入图像中的框选位置
@param indices------输出,返回经过极大值抑制后的对象的索引index
@param isDrawRect------是否显示显示框选图像(设置为true会弹出一个窗口)
@res 返回0表示检测成功,其他为检测失败
*/
int Detect_YOLOV4(cv::Mat image, float confidenceThreshold, std::vector<int>& res_classID, std::vector<float>& res_confidences, std::vector<cv::Rect>& res_boxes, std::vector<int>& indices, bool isDrawRect);
private:
//画出检测到的指定区域(内部调用函数)
void drawPred(int classId, float conf, int left, int top, int right, int bottom, cv::Mat& frame);
std::vector<cv::String> getOutputsNames(const cv::dnn::Net& net);
//YOLO param
std::vector<std::string> YOLOV4classNamesVec;
cv::dnn::Net* YOLOV4net;
bool YOLO_init;
};
DraknetYolo.cpp 具体实现
#include"DraknetYolo.h"
DraknetYolo::DraknetYolo()
{
}
DraknetYolo::~DraknetYolo()
{
}
int DraknetYolo::Init_YOLOV4(std::string modelConfiguration, std::string modelBinary, std::string classNames)
{
if (modelConfiguration.empty() || modelBinary.empty() || classNames.empty()) {
std::cout << "input modelfile is empty" << std::endl;
return -1;
}
YOLOV4net = new cv::dnn::Net();
YOLOV4classNamesVec.clear();
YOLOV4classNamesVec.resize(0);
YOLO_init = false; //是否已初始化
*YOLOV4net = cv::dnn::readNetFromDarknet(modelConfiguration, modelBinary);
if (YOLOV4net->empty())
{
std::cout << "Could not load net..." << std::endl;
return -1;
}
#ifdef CUDA
//使用CUDA加速
//YOLOV4net.setPreferableBackend(cv::dnn::DNN_BACKEND_CUDA);
//YOLOV4net.setPreferableTarget(cv::dnn::DNN_TARGET_CUDA);
#else
YOLOV4net->setPreferableTarget(cv::dnn::DNN_TARGET_OPENCL_FP16);
#endif //CUDA
std::ifstream classNamesFile(classNames);
if (classNamesFile.is_open())
{
std::string className = "";
while (std::getline(classNamesFile, className))
YOLOV4classNamesVec.push_back(className);
}
std::cout << "Init YOLOV4 suc" << std::endl;
YOLO_init = true;
return 0;
}
int DraknetYolo::Detect_YOLOV4(cv::Mat image, float confidenceThreshold, vector<int>& res_classID, vector<float>& res_confidences, vector<cv::Rect>& res_boxes, vector<int>& indices, bool isDrawRect)
{
if (!YOLO_init) {
std::cout << "must be init first" << std::endl;
return -1;
}
if (image.empty()) {
std::cout << "input image is empty________in func Detect_YOLOV4" << std::endl;
return -1;
}
cv::Mat tempimg = image.clone();
//416
cv::Mat inputBlob = cv::dnn::blobFromImage(tempimg, 1 / 255.F, cv::Size(416, 416), cv::Scalar(0, 0, 0), true, false);
if (inputBlob.empty())
return -1;
YOLOV4net->setInput(inputBlob.clone());
cv::Mat detectionMat;
std::vector<cv::Mat> outs;
YOLOV4net->forward(outs, getOutputsNames(*YOLOV4net));
vector<double> layersTimings;
vector<int> classIds; //保存每个检测到的对象的分类的ID(第几个分类)
vector<float> confidences; //保存每个检测到的对象的分数
vector<cv::Rect> boxes; //保存每个检测到的对象的框选区域
for (int i = 0; i < outs.size(); i++) {
if (outs[i].empty())
break;
detectionMat = outs[i].clone();
float* data = (float*)detectionMat.data;
for (int j = 0; j < detectionMat.rows; ++j, data += detectionMat.cols)
{
cv::Mat scores = detectionMat.row(j).colRange(5, detectionMat.cols);
cv::Point classIdPoint;
double confidence;
minMaxLoc(scores, 0, &confidence, 0, &classIdPoint);
if (confidence > confidenceThreshold)
{
int centerX = (int)(data[0] * tempimg.cols);
int centerY = (int)(data[1] * tempimg.rows);
int width = (int)(data[2] * tempimg.cols);
int height = (int)(data[3] * tempimg.rows);
int left = centerX - width / 2;
int top = centerY - height / 2;
classIds.push_back(classIdPoint.x);
confidences.push_back((float)confidence);
boxes.push_back(cv::Rect(left, top, width, height));
}
}
}
//非极大值抑制
if (boxes.empty()) {
return -1;
}
cv::dnn::NMSBoxes(boxes, confidences, confidenceThreshold, confidenceThreshold, indices);
res_boxes = boxes;
res_confidences = confidences;
res_classID = classIds;
if (isDrawRect) {
for (size_t i = 0; i < indices.size(); ++i)
{
int idx = indices[i];
cv::Rect box = boxes[idx];
drawPred(classIds[idx], confidences[idx], box.x, box.y,
box.x + box.width, box.y + box.height, tempimg);
}
cv::imshow("YOLO-Detections", tempimg);
cv::waitKey(10);
}
return 0;
}
std::vector<cv::String> DraknetYolo::getOutputsNames(const cv::dnn::Net& net)
{
vector<cv::String> names;
if (names.empty())
{
//得到输出层的索引号
std::vector<int> out_layer_indx = net.getUnconnectedOutLayers();
//得到网络中所有层的名称
std::vector<cv::String> all_layers_names = net.getLayerNames();
//在名称中获取输出层的名称
names.resize(out_layer_indx.size());
for (int i = 0; i < out_layer_indx.size(); i++)
{
names[i] = all_layers_names[out_layer_indx[i] - 1];
}
}
return names;
}
void DraknetYolo::drawPred(int classId, float conf, int left, int top, int right, int bottom, cv::Mat& frame)
{
rectangle(frame, cv::Point(left, top), cv::Point(right, bottom), cv::Scalar(0, 255, 0));
std::string label = cv::format("%.2f", conf);
if (!YOLOV4classNamesVec.empty())
{
CV_Assert(classId < (int)YOLOV4classNamesVec.size());
label = YOLOV4classNamesVec[classId] + ": " + label;
}
int baseLine;
cv::Size labelSize = cv::getTextSize(label, cv::FONT_HERSHEY_SIMPLEX, 0.5, 1, &baseLine);
top = cv::max(top, labelSize.height);
rectangle(frame, cv::Point(left, top - labelSize.height),
cv::Point(left + labelSize.width, top + baseLine), cv::Scalar::all(255), cv::FILLED);
putText(frame, label, cv::Point(left, top), cv::FONT_HERSHEY_SIMPLEX, 0.5, cv::Scalar());
}
具体调用方法
。
main.cpp
#include"Darknet.h"
int main(){
cvsa::DraknetYolo yolov4net;
yolov4net.Init_YOLOV4("yolov2-tiny-custom.cfg","yolov2-tiny-custom_final.weights","obj.names");
cv::Mat tempmat = cv::imread("1.jpg");
std::vector<int> res_classID; //属于哪一类(目前模型只有一类
std::vector<float> res_confidences; //检测分数是多少
std::vector<cv::Rect> res_boxes; //检测到的ROI区域在哪
std::vector<int> indices;
yolov4net.Detect_YOLOV4(tempmat, 0.75, res_classID, res_confidences, res_boxes, indices, false);
return 0;
}