diff --git a/modules/dnn/src/darknet/darknet_io.cpp b/modules/dnn/src/darknet/darknet_io.cpp index 815b84f651..54a53fd867 100644 --- a/modules/dnn/src/darknet/darknet_io.cpp +++ b/modules/dnn/src/darknet/darknet_io.cpp @@ -371,7 +371,7 @@ namespace cv { fused_layer_names.push_back(last_layer); } - void setYolo(int classes, const std::vector& mask, const std::vector& anchors) + void setYolo(int classes, const std::vector& mask, const std::vector& anchors, float thresh, float nms_threshold) { cv::dnn::LayerParams region_param; region_param.name = "Region-name"; @@ -382,6 +382,8 @@ namespace cv { region_param.set("classes", classes); region_param.set("anchors", numAnchors); region_param.set("logistic", true); + region_param.set("thresh", thresh); + region_param.set("nms_threshold", nms_threshold); std::vector usedAnchors(numAnchors * 2); for (int i = 0; i < numAnchors; ++i) @@ -646,6 +648,8 @@ namespace cv { { int classes = getParam(layer_params, "classes", -1); int num_of_anchors = getParam(layer_params, "num", -1); + float thresh = getParam(layer_params, "thresh", 0.2); + float nms_threshold = getParam(layer_params, "nms_threshold", 0.4); std::string anchors_values = getParam(layer_params, "anchors", std::string()); CV_Assert(!anchors_values.empty()); @@ -658,7 +662,7 @@ namespace cv { CV_Assert(classes > 0 && num_of_anchors > 0 && (num_of_anchors * 2) == anchors_vec.size()); setParams.setPermute(false); - setParams.setYolo(classes, mask_vec, anchors_vec); + setParams.setYolo(classes, mask_vec, anchors_vec, thresh, nms_threshold); } else { CV_Error(cv::Error::StsParseError, "Unknown layer type: " + layer_type);