Yolov5的烟火检测——C++实现模型部署

#include <opencv2/opencv.hpp> #include <ncnn/net.h> struct ObjectFlag { cv::Rect_<float> rect; int label; float prob; }; static inline float intersectionArea(const ObjectFlag& a, const ObjectFlag& b) { cv::Rect_<float> inter = a.rect & b.rect; return inter.area(); } static void qsortDescentInplace(std::vector<ObjectFlag>& faceobjects, int left, int right) { int i = left; int j = right; float p = faceobjects[(left + right) / 2].prob; while (i <= j) { while (faceobjects[i].prob > p) i++; while (faceobjects[j].prob < p) j--; if (i <= j) { // swap std::swap(faceobjects[i], faceobjects[j]); i++; j--; } } #pragma omp parallel sections { #pragma omp section { if (left < j) qsortDescentInplace(faceobjects, left, j); } #pragma omp section { if (i < right) qsortDescentInplace(faceobjects, i, right); } } } static void qsortDescentInplace(std::vector<ObjectFlag>& faceobjects) { if (faceobjects.empty()) return; qsortDescentInplace(faceobjects, 0, int(faceobjects.size()) - 1); } static void nmsSortedBboxes(const std::vector<ObjectFlag>& faceobjects, std::vector<int>& picked, float nms_threshold) { picked.clear(); const int n = int(faceobjects.size()); std::vector<float> areas(n); for (int i = 0; i < n; i++) { areas[i] = faceobjects[i].rect.area(); } for (int i = 0; i < n; i++) { const ObjectFlag& a = faceobjects[i]; int keep = 1; for (int j = 0; j < (int)picked.size(); j++) { const ObjectFlag& b = faceobjects[picked[j]]; // intersection over union float inter_area = intersectionArea(a, b); float union_area = areas[i] + areas[picked[j]] - inter_area; // float IoU = inter_area / union_area if (inter_area / union_area > nms_threshold) keep = 0; } if (keep) picked.push_back(i); } } static inline float sigmoid(float x) { return static_cast<float>(1.f / (1.f + exp(-x))); } static void generateProposals(const ncnn::Mat& anchors, int stride, const ncnn::Mat& in_pad, const ncnn::Mat& feat_blob, float prob_threshold, std::vector<ObjectFlag>& objects) { const int num_grid = feat_blob.h; int num_grid_x; int num_grid_y; if (in_pad.w > in_pad.h) { num_grid_x = in_pad.w / stride; num_grid_y = num_grid / num_grid_x; } else { num_grid_y = in_pad.h / stride; num_grid_x = num_grid / num_grid_y; } const int num_class = feat_blob.w - 5; const int num_anchors = anchors.w / 2; for (int q = 0; q < num_anchors; q++) { const float anchor_w = anchors[q * 2]; const float anchor_h = anchors[q * 2 + 1]; const ncnn::Mat feat = feat_blob.channel(q); for (int i = 0; i < num_grid_y; i++) { for (int j = 0; j < num_grid_x; j++) { const float* featptr = feat.row(i * num_grid_x + j); // find class index with max class score int class_index = 0; float class_score = -FLT_MAX; for (int k = 0; k < num_class; k++) { float score = featptr[5 + k]; if (score > class_score) { class_index = k; class_score = score; } } float box_score = featptr[4]; float confidence = sigmoid(box_score) * sigmoid(class_score); if (confidence >= prob_threshold) { float dx = sigmoid(featptr[0]); float dy = sigmoid(featptr[1]); float dw = sigmoid(featptr[2]); float dh = sigmoid(featptr[3]); float pb_cx = (dx * 2.f - 0.5f + j) * stride; float pb_cy = (dy * 2.f - 0.5f + i) * stride; float pb_w = pow(dw * 2.f, 2) * anchor_w; float pb_h = pow(dh * 2.f, 2) * anchor_h; float x0 = pb_cx - pb_w * 0.5f; float y0 = pb_cy - pb_h * 0.5f; float x1 = pb_cx + pb_w * 0.5f; float y1 = pb_cy + pb_h * 0.5f; ObjectFlag obj; obj.rect.x = x0; obj.rect.y = y0; obj.rect.width = x1 - x0; obj.rect.height = y1 - y0; obj.label = class_index; obj.prob = confidence; objects.push_back(obj); } } } } } int detectYolov5(const cv::Mat& cv_src, ncnn::Net& ncnn_net, std::vector<ObjectFlag>& objects, int threads, int target_size, float prob_threshold, float nms_threshold) { int img_w = cv_src.cols; int img_h = cv_src.rows; // letterbox pad to multiple of 32 int w = img_w; int h = img_h; float scale = 1.f; if (w > h) { scale = (float)target_size / w; w = target_size; h = h * scale; } else { scale = (float)target_size / h; h = target_size; w = w * scale; } ncnn::Mat in = ncnn::Mat::from_pixels_resize(cv_src.data, ncnn::Mat::PIXEL_BGR2RGB, img_w, img_h, w, h); int wpad = (w + 31) / 32 * 32 - w; int hpad = (h + 31) / 32 * 32 - h; ncnn::Mat in_pad; ncnn::copy_make_border(in, in_pad, hpad / 2, hpad - hpad / 2, wpad / 2, wpad - wpad / 2, ncnn::BORDER_CONSTANT, 114.f); const float norm_vals[3] = { 1 / 255.f, 1 / 255.f, 1 / 255.f }; in_pad.substract_mean_normalize(0, norm_vals); ncnn::Extractor ex = ncnn_net.create_extractor(); //设置线程数，这将覆盖全局设置，依赖于系统能调用的最大个数 ex.set_num_threads(threads); ex.input("images", in_pad); std::vector<ObjectFlag> proposals; // stride 8 { ncnn::Mat out; ex.extract("output", out); ncnn::Mat anchors(6); anchors[0] = 10.f; anchors[1] = 13.f; anchors[2] = 16.f; anchors[3] = 30.f; anchors[4] = 33.f; anchors[5] = 23.f; std::vector<ObjectFlag> objects8; generateProposals(anchors, 8, in_pad, out, prob_threshold, objects8); proposals.insert(proposals.end(), objects8.begin(), objects8.end()); } // stride 16 { ncnn::Mat out; ex.extract("402", out); ncnn::Mat anchors(6); anchors[0] = 30.f; anchors[1] = 61.f; anchors[2] = 62.f; anchors[3] = 45.f; anchors[4] = 59.f; anchors[5] = 119.f; std::vector<ObjectFlag> objects16; generateProposals(anchors, 16, in_pad, out, prob_threshold, objects16); proposals.insert(proposals.end(), objects16.begin(), objects16.end()); } // stride 32 { ncnn::Mat out; ex.extract("416", out); ncnn::Mat anchors(6); anchors[0] = 116.f; anchors[1] = 90.f; anchors[2] = 156.f; anchors[3] = 198.f; anchors[4] = 373.f; anchors[5] = 326.f; std::vector<ObjectFlag> objects32; generateProposals(anchors, 32, in_pad, out, prob_threshold, objects32); proposals.insert(proposals.end(), objects32.begin(), objects32.end()); } qsortDescentInplace(proposals); std::vector<int> picked; nmsSortedBboxes(proposals,