Merge pull request opencv#10446 from cabelo:style-dnn-yolo

cabelo · alalek · commit 7abaae39ecb5 · 2018-01-02T17:26:29.000+03:00
* add style draw in yolo

* fix sintaxe whitespace

* fix conversion from float to int

* sample refactoring

- minor code style fixes
- avoid confusing "bottom" names
- use cv::format
- always draw object detection/roi
diff --git a/samples/dnn/yolo_object_detection.cpp b/samples/dnn/yolo_object_detection.cpp
@@ -26,6 +26,7 @@ static const char* params =
 "{ model          |       | model weights       }"
 "{ camera_device  | 0     | camera device number}"
 "{ source         |       | video or image for detection}"
+"{ style          | box   | box or line style draw }"
 "{ min_confidence | 0.24  | min confidence      }"
 "{ class_names    |       | File with class names, [PATH-TO-DARKNET]/data/coco.names }";
 
@@ -78,7 +79,7 @@ int main(int argc, char** argv)
         }
     }
 
-    vector<string> classNamesVec;
+    vector<String> classNamesVec;
     ifstream classNamesFile(parser.get<String>("class_names").c_str());
     if (classNamesFile.is_open())
     {
@@ -87,6 +88,8 @@ int main(int argc, char** argv)
             classNamesVec.push_back(className);
     }
 
+    String object_roi_style = parser.get<String>("style");
+
     for(;;)
     {
         Mat frame;
@@ -114,11 +117,10 @@ int main(int argc, char** argv)
         //! [Make forward pass]
 
         vector<double> layersTimings;
-        double freq = getTickFrequency() / 1000;
-        double time = net.getPerfProfile(layersTimings) / freq;
-        ostringstream ss;
-        ss << "FPS: " << 1000/time << " ; time: " << time << " ms";
-        putText(frame, ss.str(), Point(20,20), 0, 0.5, Scalar(0,0,255));
+        double tick_freq = getTickFrequency();
+        double time_ms = net.getPerfProfile(layersTimings) / tick_freq * 1000;
+        putText(frame, format("FPS: %.2f ; time: %.2f ms", 1000.f / time_ms, time_ms),
+                Point(20, 20), 0, 0.5, Scalar(0, 0, 255));
 
         float confidenceThreshold = parser.get<float>("min_confidence");
         for (int i = 0; i < detectionMat.rows; i++)
@@ -132,44 +134,34 @@ int main(int argc, char** argv)
 
             if (confidence > confidenceThreshold)
             {
-                float x = detectionMat.at<float>(i, 0);
-                float y = detectionMat.at<float>(i, 1);
-                float width = detectionMat.at<float>(i, 2);
-                float height = detectionMat.at<float>(i, 3);
-                int xLeftBottom = static_cast<int>((x - width / 2) * frame.cols);
-                int yLeftBottom = static_cast<int>((y - height / 2) * frame.rows);
-                int xRightTop = static_cast<int>((x + width / 2) * frame.cols);
-                int yRightTop = static_cast<int>((y + height / 2) * frame.rows);
-
-                Rect object(xLeftBottom, yLeftBottom,
-                            xRightTop - xLeftBottom,
-                            yRightTop - yLeftBottom);
-
-                rectangle(frame, object, Scalar(0, 255, 0));
-
-                if (objectClass < classNamesVec.size())
+                float x_center = detectionMat.at<float>(i, 0) * frame.cols;
+                float y_center = detectionMat.at<float>(i, 1) * frame.rows;
+                float width = detectionMat.at<float>(i, 2) * frame.cols;
+                float height = detectionMat.at<float>(i, 3) * frame.rows;
+                Point p1(cvRound(x_center - width / 2), cvRound(y_center - height / 2));
+                Point p2(cvRound(x_center + width / 2), cvRound(y_center + height / 2));
+                Rect object(p1, p2);
+
+                Scalar object_roi_color(0, 255, 0);
+
+                if (object_roi_style == "box")
                 {
-                    ss.str("");
-                    ss << confidence;
-                    String conf(ss.str());
-                    String label = String(classNamesVec[objectClass]) + ": " + conf;
-                    int baseLine = 0;
-                    Size labelSize = getTextSize(label, FONT_HERSHEY_SIMPLEX, 0.5, 1, &baseLine);
-                    rectangle(frame, Rect(Point(xLeftBottom, yLeftBottom ),
-                                          Size(labelSize.width, labelSize.height + baseLine)),
-                              Scalar(255, 255, 255), CV_FILLED);
-                    putText(frame, label, Point(xLeftBottom, yLeftBottom+labelSize.height),
-                            FONT_HERSHEY_SIMPLEX, 0.5, Scalar(0,0,0));
+                    rectangle(frame, object, object_roi_color);
                 }
                 else
                 {
-                    cout << "Class: " << objectClass << endl;
-                    cout << "Confidence: " << confidence << endl;
-                    cout << " " << xLeftBottom
-                         << " " << yLeftBottom
-                         << " " << xRightTop
-                         << " " << yRightTop << endl;
+                    Point p_center(cvRound(x_center), cvRound(y_center));
+                    line(frame, object.tl(), p_center, object_roi_color, 1);
                 }
+
+                String className = objectClass < classNamesVec.size() ? classNamesVec[objectClass] : cv::format("unknown(%d)", objectClass);
+                String label = format("%s: %.2f", className.c_str(), confidence);
+                int baseLine = 0;
+                Size labelSize = getTextSize(label, FONT_HERSHEY_SIMPLEX, 0.5, 1, &baseLine);
+                rectangle(frame, Rect(p1, Size(labelSize.width, labelSize.height + baseLine)),
+                          object_roi_color, CV_FILLED);
+                putText(frame, label, p1 + Point(0, labelSize.height),
+                        FONT_HERSHEY_SIMPLEX, 0.5, Scalar(0,0,0));
             }
         }
 
