yoloWrapper.cpp
3.7 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
#include "yoloWrapper.h"
int YoloWrapper::init(const std::string& detector_config_path,
const std::string& detector_model_path, const cv::Size& input_size) {
if (!detector_config_path.empty() && !detector_model_path.empty()) {
//加载并初始化网络
net_ = cv::dnn::readNet(detector_config_path, detector_model_path);
net_.setPreferableBackend(cv::dnn::DNN_BACKEND_OPENCV);
net_.setPreferableBackend(cv::dnn::DNN_TARGET_CPU);
//设置输入尺寸
Size_ = cv::Size(input_size.width, input_size.height);
}
return 0;
}
int YoloWrapper::setPreferableParams(float fConfidence, float fNMSThreshold)
{
this->confidence = fConfidence; this->NMSThreshold = fNMSThreshold;
return 0;
}
std::vector<cv::Rect> YoloWrapper::forward(cv::Mat img) {
//获取输出层名称
auto layerNames = net_.getLayerNames();
std::vector<int> outLayers = net_.getUnconnectedOutLayers();
std::vector<std::string> outBlobNames(outLayers.size());
for (int i = 0; i < outLayers.size(); i++) {
outBlobNames[i] = layerNames[outLayers[i] - 1];
}
//预处理图像
cv::Mat blob;
cv::dnn::blobFromImage(img, blob, 1.0 / 255.0, Size_);
//为网络输入新值
net_.setInput(blob);
//获取预测结果
std::vector<cv::Mat> outputBlobs;
net_.forward(outputBlobs, outBlobNames);
std::vector<float> confidences;
std::vector<cv::Rect> bboxes;
for (int n = 0; n < outputBlobs.size(); n++) {
for (int row = 0; row < outputBlobs[n].rows; row++) {
//置信度
cv::Mat prob = outputBlobs[n](cv::Range(row, row + 1), cv::Range(5, outputBlobs[n].cols));
//取得最大分数值与索引
double score = .0; cv::Point classIdPoint;
cv::minMaxLoc(prob, 0, &score, 0, &classIdPoint);
//把可能的区域都算作进去
if (score > confidence) {
cv::Mat dt = outputBlobs[n](cv::Range(row, row + 1), cv::Range(0, 5));
int cx = cvRound(dt.ptr<float>(0)[0] * (float)img.cols);
int cy = cvRound(dt.ptr<float>(0)[1] * (float)img.rows);
int w = cvRound(dt.ptr<float>(0)[2] * (float)img.cols);
int h = cvRound(dt.ptr<float>(0)[3] * (float)img.rows);
int left = cx - w / 2;
int top = cy - h / 2;
confidences.push_back((float)score);
bboxes.push_back(cv::Rect(left, top, w, h));
}
}
}
std::vector<int> indices;
//非极大值抑制
cv::dnn::NMSBoxes(bboxes, confidences, confidence, NMSThreshold, indices);
//裁剪区域
auto points = std::vector<cv::Rect>();
for (int i = 0; i < indices.size(); i++) {
int idx = indices[i];
cv::Rect bbox = bboxes[idx] & cv::Rect(0, 0, img.cols, img.rows);
points.push_back(bbox);
}
return points;
}
#ifdef _DEBUG
//class YoloDarknet::Impl {
//public:
// Impl() {}
// ~Impl() {}
//
// std::shared_ptr<DarkNet> net_;
//};
//
//YoloDarknet::YoloDarknet(const std::string& config_path, const std::string& model_path, const int ntype)
//{
// p = cv::makePtr<YoloDarknet::Impl>();
// if (!config_path.empty() && !model_path.empty()) {
// p->net_ = std::make_shared<DarkNet>();
// p->net_->init(config_path, model_path, ntype);
// }
// else {
// p->net_ = NULL;
// }
//}
//
//void YoloDarknet::setInput(cv::Mat& img, float threshold, int topk) {
// topk_ = topk;
// p->net_->setPreferableParams(threshold, topk);
// p->net_->forward(img);
//}
//
//void YoloDarknet::forward(std::vector<int> &outputPredict, std::vector<float> &outputConfidence, std::vector<cv::Rect> &outputBoxes)
//{
// auto bbox = std::vector<cv::Rect>();
// int *predicts = new int[topk_]; float *confidence = new float[topk_];
// p->net_->getResult(predicts, confidence, &bbox);
// for (int top = 0; top < topk_; top++) outputPredict.push_back(predicts[top]), outputConfidence.push_back(confidence[top]);
// for (auto&box : bbox) {
// outputBoxes.push_back(box);
// }
// delete[] predicts; predicts = NULL; delete[] confidence; confidence = NULL;
//}
#endif