Merge pull request #11322 from dkurt:dnn_yolov3

samples/dnn/object_detection.cpp

@@ -35,12 +35,14 @@ using namespace dnn;
 float confThreshold;
 std::vector<std::string> classes;
 
-void postprocess(Mat& frame, const Mat& out, Net& net);
+void postprocess(Mat& frame, const std::vector<Mat>& out, Net& net);
 
 void drawPred(int classId, float conf, int left, int top, int right, int bottom, Mat& frame);
 
 void callback(int pos, void* userdata);
 
+std::vector<String> getOutputsNames(const Net& net);
+
 int main(int argc, char** argv)
 {
     CommandLineParser parser(argc, argv, keys);
@@ -115,9 +117,10 @@ int main(int argc, char** argv)
             Mat imInfo = (Mat_<float>(1, 3) << inpSize.height, inpSize.width, 1.6f);
             net.setInput(imInfo, "im_info");
         }
-        Mat out = net.forward();
+        std::vector<Mat> outs;
+        net.forward(outs, getOutputsNames(net));
 
-        postprocess(frame, out, net);
+        postprocess(frame, outs, net);
 
         // Put efficiency information.
         std::vector<double> layersTimes;
@@ -131,18 +134,19 @@ int main(int argc, char** argv)
     return 0;
 }
 
-void postprocess(Mat& frame, const Mat& out, Net& net)
+void postprocess(Mat& frame, const std::vector<Mat>& outs, Net& net)
 {
     static std::vector<int> outLayers = net.getUnconnectedOutLayers();
     static std::string outLayerType = net.getLayer(outLayers[0])->type;
 
-    float* data = (float*)out.data;
     if (net.getLayer(0)->outputNameToIndex("im_info") != -1)  // Faster-RCNN or R-FCN
     {
         // Network produces output blob with a shape 1x1xNx7 where N is a number of
         // detections and an every detection is a vector of values
         // [batchId, classId, confidence, left, top, right, bottom]
-        for (size_t i = 0; i < out.total(); i += 7)
+        CV_Assert(outs.size() == 1);
+        float* data = (float*)outs[0].data;
+        for (size_t i = 0; i < outs[0].total(); i += 7)
         {
             float confidence = data[i + 2];
             if (confidence > confThreshold)
@@ -161,7 +165,9 @@ void postprocess(Mat& frame, const Mat& out, Net& net)
         // Network produces output blob with a shape 1x1xNx7 where N is a number of
         // detections and an every detection is a vector of values
         // [batchId, classId, confidence, left, top, right, bottom]
-        for (size_t i = 0; i < out.total(); i += 7)
+        CV_Assert(outs.size() == 1);
+        float* data = (float*)outs[0].data;
+        for (size_t i = 0; i < outs[0].total(); i += 7)
         {
             float confidence = data[i + 2];
             if (confidence > confThreshold)
@@ -177,27 +183,45 @@ void postprocess(Mat& frame, const Mat& out, Net& net)
     }
     else if (outLayerType == "Region")
     {
-        // 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]
-        for (int i = 0; i < out.rows; ++i, data += out.cols)
+        std::vector<int> classIds;
+        std::vector<float> confidences;
+        std::vector<Rect> boxes;
+        for (size_t i = 0; i < outs.size(); ++i)
         {
-            Mat confidences = out.row(i).colRange(5, out.cols);
-            Point classIdPoint;
-            double confidence;
-            minMaxLoc(confidences, 0, &confidence, 0, &classIdPoint);
-            if (confidence > confThreshold)
+            // 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]
+            float* data = (float*)outs[i].data;
+            for (int j = 0; j < outs[i].rows; ++j, data += outs[i].cols)
             {
-                int classId = classIdPoint.x;
-                int centerX = (int)(data[0] * frame.cols);
-                int centerY = (int)(data[1] * frame.rows);
-                int width = (int)(data[2] * frame.cols);
-                int height = (int)(data[3] * frame.rows);
-                int left = centerX - width / 2;
-                int top = centerY - height / 2;
-                drawPred(classId, (float)confidence, left, top, left + width, top + height, frame);
+                Mat scores = outs[i].row(j).colRange(5, outs[i].cols);
+                Point classIdPoint;
+                double confidence;
+                minMaxLoc(scores, 0, &confidence, 0, &classIdPoint);
+                if (confidence > confThreshold)
+                {
+                    int centerX = (int)(data[0] * frame.cols);
+                    int centerY = (int)(data[1] * frame.rows);
+                    int width = (int)(data[2] * frame.cols);
+                    int height = (int)(data[3] * frame.rows);
+                    int left = centerX - width / 2;
+                    int top = centerY - height / 2;
+
+                    classIds.push_back(classIdPoint.x);
+                    confidences.push_back((float)confidence);
+                    boxes.push_back(Rect(left, top, width, height));
+                }
             }
         }
+        std::vector<int> indices;
+        NMSBoxes(boxes, confidences, confThreshold, 0.4, indices);
+        for (size_t i = 0; i < indices.size(); ++i)
+        {
+            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);
+        }
     }
     else
         CV_Error(Error::StsNotImplemented, "Unknown output layer type: " + outLayerType);
@@ -227,3 +251,17 @@ void callback(int pos, void*)
 {
     confThreshold = pos * 0.01f;
 }
+
+std::vector<String> getOutputsNames(const Net& net)
+{
+    static std::vector<String> names;
+    if (names.empty())
+    {
+        std::vector<int> outLayers = net.getUnconnectedOutLayers();
+        std::vector<String> layersNames = net.getLayerNames();
+        names.resize(outLayers.size());
+        for (size_t i = 0; i < outLayers.size(); ++i)
+            names[i] = layersNames[outLayers[i] - 1];
+    }
+    return names;
+}

samples/dnn/object_detection.py

@@ -55,7 +55,11 @@ net.setPreferableTarget(args.target)
 
 confThreshold = args.thr
 
-def postprocess(frame, out):
+def getOutputsNames(net):
+    layersNames = net.getLayerNames()
+    return [layersNames[i[0] - 1] for i in net.getUnconnectedOutLayers()]
+
+def postprocess(frame, outs):
     frameHeight = frame.shape[0]
     frameWidth = frame.shape[1]
 
@@ -63,7 +67,7 @@ def postprocess(frame, out):
         # Draw a bounding box.
         cv.rectangle(frame, (left, top), (right, bottom), (0, 255, 0))
 
-        label = '%.2f' % confidence
+        label = '%.2f' % conf
 
         # Print a label of class.
         if classes:
@@ -83,6 +87,8 @@ def postprocess(frame, out):
         # Network produces output blob with a shape 1x1xNx7 where N is a number of
         # detections and an every detection is a vector of values
         # [batchId, classId, confidence, left, top, right, bottom]
+        assert(len(outs) == 1)
+        out = outs[0]
         for detection in out[0, 0]:
             confidence = detection[2]
             if confidence > confThreshold:
@@ -96,6 +102,8 @@ def postprocess(frame, out):
         # Network produces output blob with a shape 1x1xNx7 where N is a number of
         # detections and an every detection is a vector of values
         # [batchId, classId, confidence, left, top, right, bottom]
+        assert(len(outs) == 1)
+        out = outs[0]
         for detection in out[0, 0]:
             confidence = detection[2]
             if confidence > confThreshold:
@@ -109,18 +117,33 @@ def postprocess(frame, out):
         # 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]
-        for detection in out:
-            confidences = detection[5:]
-            classId = np.argmax(confidences)
-            confidence = confidences[classId]
-            if confidence > confThreshold:
-                center_x = int(detection[0] * frameWidth)
-                center_y = int(detection[1] * frameHeight)
-                width = int(detection[2] * frameWidth)
-                height = int(detection[3] * frameHeight)
-                left = center_x - width / 2
-                top = center_y - height / 2
-                drawPred(classId, confidence, left, top, left + width, top + height)
+        classIds = []
+        confidences = []
+        boxes = []
+        for out in outs:
+            for detection in out:
+                scores = detection[5:]
+                classId = np.argmax(scores)
+                confidence = scores[classId]
+                if confidence > confThreshold:
+                    center_x = int(detection[0] * frameWidth)
+                    center_y = int(detection[1] * frameHeight)
+                    width = int(detection[2] * frameWidth)
+                    height = int(detection[3] * frameHeight)
+                    left = center_x - width / 2
+                    top = center_y - height / 2
+                    classIds.append(classId)
+                    confidences.append(float(confidence))
+                    boxes.append([left, top, width, height])
+        indices = cv.dnn.NMSBoxes(boxes, confidences, confThreshold, 0.4)
+        for i in indices:
+            i = i[0]
+            box = boxes[i]
+            left = box[0]
+            top = box[1]
+            width = box[2]
+            height = box[3]
+            drawPred(classIds[i], confidences[i], left, top, left + width, top + height)
 
 # Process inputs
 winName = 'Deep learning object detection in OpenCV'
@@ -152,9 +175,9 @@ while cv.waitKey(1) < 0:
     if net.getLayer(0).outputNameToIndex('im_info') != -1:  # Faster-RCNN or R-FCN
         frame = cv.resize(frame, (inpWidth, inpHeight))
         net.setInput(np.array([inpHeight, inpWidth, 1.6], dtype=np.float32), 'im_info');
-    out = net.forward()
+    outs = net.forward(getOutputsNames(net))
 
-    postprocess(frame, out)
+    postprocess(frame, outs)
 
     # Put efficiency information.
     t, _ = net.getPerfProfile()