opencv4部署darknet训练的yolo模型
记录一下代码，方便以后回看
darknet框架训练的yolov4会生成 weights权重文件，加上配置的.cfg文件和names分类文件用于配置
yolov4在github上的使用darknet的训练教程
查看显卡对应驱动配置CUDNN
yolov4加CUDA的配置
DraknetYolo.h
class DraknetYolo
{
public:
DraknetYolo();
~DraknetYolo();
int Init_YOLOV4(std::string modelConfiguration, std::string modelBinary, std::string classNames);
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);
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
#else
YOLOV4net->setPreferableTarget(cv::dnn::DNN_TARGET_OPENCL_FP16);
#endif
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();
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;
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;
std::vector<int> indices;
yolov4net.Detect_YOLOV4(tempmat, 0.75, res_classID, res_confidences, res_boxes, indices, false);
return 0;
}
本文内容由网友自发贡献，版权归原作者所有，本站不承担相应法律责任。如您发现有涉嫌抄袭侵权的内容，请联系:hwhale#tublm.com(使用前将#替换为@)
