From d5e8792f5502e451951b2ab5209838d753a7c369 Mon Sep 17 00:00:00 2001 From: Liubov Batanina Date: Mon, 25 May 2020 15:34:11 +0300 Subject: [PATCH] Merge pull request #17332 from l-bat:fix_nms Fixed NMSBoxes bug * Added NMS for each class * Updated cpp sample * Fixed errors * Refactoring * Added NMS for IE --- samples/dnn/object_detection.cpp | 56 +++++++++++++++++++++++++++----- samples/dnn/object_detection.py | 23 ++++++++++--- 2 files changed, 65 insertions(+), 14 deletions(-) diff --git a/samples/dnn/object_detection.cpp b/samples/dnn/object_detection.cpp index 83ed10db5d..c7e42430fe 100644 --- a/samples/dnn/object_detection.cpp +++ b/samples/dnn/object_detection.cpp @@ -45,7 +45,7 @@ std::vector classes; inline void preprocess(const Mat& frame, Net& net, Size inpSize, float scale, const Scalar& mean, bool swapRB); -void postprocess(Mat& frame, const std::vector& out, Net& net); +void postprocess(Mat& frame, const std::vector& out, Net& net, int backend); void drawPred(int classId, float conf, int left, int top, int right, int bottom, Mat& frame); @@ -148,7 +148,8 @@ int main(int argc, char** argv) // Load a model. Net net = readNet(modelPath, configPath, parser.get("framework")); - net.setPreferableBackend(parser.get("backend")); + int backend = parser.get("backend"); + net.setPreferableBackend(backend); net.setPreferableTarget(parser.get("target")); std::vector outNames = net.getUnconnectedOutLayersNames(); @@ -245,7 +246,7 @@ int main(int argc, char** argv) std::vector outs = predictionsQueue.get(); Mat frame = processedFramesQueue.get(); - postprocess(frame, outs, net); + postprocess(frame, outs, net, backend); if (predictionsQueue.counter > 1) { @@ -285,7 +286,7 @@ int main(int argc, char** argv) std::vector outs; net.forward(outs, outNames); - postprocess(frame, outs, net); + postprocess(frame, outs, net, backend); // Put efficiency information. std::vector layersTimes; @@ -319,7 +320,7 @@ inline void preprocess(const Mat& frame, Net& net, Size inpSize, float scale, } } -void postprocess(Mat& frame, const std::vector& outs, Net& net) +void postprocess(Mat& frame, const std::vector& outs, Net& net, int backend) { static std::vector outLayers = net.getUnconnectedOutLayers(); static std::string outLayerType = net.getLayer(outLayers[0])->type; @@ -396,11 +397,48 @@ void postprocess(Mat& frame, const std::vector& outs, Net& net) else CV_Error(Error::StsNotImplemented, "Unknown output layer type: " + outLayerType); - std::vector indices; - NMSBoxes(boxes, confidences, confThreshold, nmsThreshold, indices); - for (size_t i = 0; i < indices.size(); ++i) + // NMS is used inside Region layer only on DNN_BACKEND_OPENCV for another backends we need NMS in sample + // or NMS is required if number of outputs > 1 + if (outLayers.size() > 1 || (outLayerType == "Region" && backend != DNN_BACKEND_OPENCV)) + { + std::map > class2indices; + for (size_t i = 0; i < classIds.size(); i++) + { + if (confidences[i] >= confThreshold) + { + class2indices[classIds[i]].push_back(i); + } + } + std::vector nmsBoxes; + std::vector nmsConfidences; + std::vector nmsClassIds; + for (std::map >::iterator it = class2indices.begin(); it != class2indices.end(); ++it) + { + std::vector localBoxes; + std::vector localConfidences; + std::vector classIndices = it->second; + for (size_t i = 0; i < classIndices.size(); i++) + { + localBoxes.push_back(boxes[classIndices[i]]); + localConfidences.push_back(confidences[classIndices[i]]); + } + std::vector nmsIndices; + NMSBoxes(localBoxes, localConfidences, confThreshold, nmsThreshold, nmsIndices); + for (size_t i = 0; i < nmsIndices.size(); i++) + { + size_t idx = nmsIndices[i]; + nmsBoxes.push_back(localBoxes[idx]); + nmsConfidences.push_back(localConfidences[idx]); + nmsClassIds.push_back(it->first); + } + } + boxes = nmsBoxes; + classIds = nmsClassIds; + confidences = nmsConfidences; + } + + for (size_t idx = 0; idx < boxes.size(); ++idx) { - int idx = indices[i]; Rect box = boxes[idx]; drawPred(classIds[idx], confidences[idx], box.x, box.y, box.x + box.width, box.y + box.height, frame); diff --git a/samples/dnn/object_detection.py b/samples/dnn/object_detection.py index d4ea40f935..babac0dbe8 100644 --- a/samples/dnn/object_detection.py +++ b/samples/dnn/object_detection.py @@ -141,9 +141,6 @@ def postprocess(frame, outs): # Network produces output blob with a shape NxC where N is a number of # detected objects and C is a number of classes + 4 where the first 4 # numbers are [center_x, center_y, width, height] - classIds = [] - confidences = [] - boxes = [] for out in outs: for detection in out: scores = detection[5:] @@ -163,9 +160,25 @@ def postprocess(frame, outs): print('Unknown output layer type: ' + lastLayer.type) exit() - indices = cv.dnn.NMSBoxes(boxes, confidences, confThreshold, nmsThreshold) + # NMS is used inside Region layer only on DNN_BACKEND_OPENCV for another backends we need NMS in sample + # or NMS is required if number of outputs > 1 + if len(outNames) > 1 or lastLayer.type == 'Region' and args.backend != cv.dnn.DNN_BACKEND_OPENCV: + indices = [] + classIds = np.array(classIds) + boxes = np.array(boxes) + confidences = np.array(confidences) + unique_classes = set(classIds) + for cl in unique_classes: + class_indices = np.where(classIds == cl)[0] + conf = confidences[class_indices] + box = boxes[class_indices].tolist() + nms_indices = cv.dnn.NMSBoxes(box, conf, confThreshold, nmsThreshold) + nms_indices = nms_indices[:, 0] if len(nms_indices) else [] + indices.extend(class_indices[nms_indices]) + else: + indices = np.arange(0, len(classIds)) + for i in indices: - i = i[0] box = boxes[i] left = box[0] top = box[1]