add visualize

ttjjmm · Jan 18, 2022 · e3ec62b · e3ec62b
1 parent 4c98fa8
commit e3ec62b
Showing 1 changed file with 132 additions and 11 deletions.
diff --git a/main.cpp b/main.cpp
@@ -85,6 +85,133 @@ int resize_uniform(cv::Mat& src, cv::Mat& dst, const cv::Size& dst_size, object_
 }
 
 
+const int color_list[80][3] =
+        {
+                //{255 ,255 ,255}, //bg
+                {216 , 82 , 24},
+                {236 ,176 , 31},
+                {125 , 46 ,141},
+                {118 ,171 , 47},
+                { 76 ,189 ,237},
+                {238 , 19 , 46},
+                { 76 , 76 , 76},
+                {153 ,153 ,153},
+                {255 ,  0 ,  0},
+                {255 ,127 ,  0},
+                {190 ,190 ,  0},
+                {  0 ,255 ,  0},
+                {  0 ,  0 ,255},
+                {170 ,  0 ,255},
+                { 84 , 84 ,  0},
+                { 84 ,170 ,  0},
+                { 84 ,255 ,  0},
+                {170 , 84 ,  0},
+                {170 ,170 ,  0},
+                {170 ,255 ,  0},
+                {255 , 84 ,  0},
+                {255 ,170 ,  0},
+                {255 ,255 ,  0},
+                {  0 , 84 ,127},
+                {  0 ,170 ,127},
+                {  0 ,255 ,127},
+                { 84 ,  0 ,127},
+                { 84 , 84 ,127},
+                { 84 ,170 ,127},
+                { 84 ,255 ,127},
+                {170 ,  0 ,127},
+                {170 , 84 ,127},
+                {170 ,170 ,127},
+                {170 ,255 ,127},
+                {255 ,  0 ,127},
+                {255 , 84 ,127},
+                {255 ,170 ,127},
+                {255 ,255 ,127},
+                {  0 , 84 ,255},
+                {  0 ,170 ,255},
+                {  0 ,255 ,255},
+                { 84 ,  0 ,255},
+                { 84 , 84 ,255},
+                { 84 ,170 ,255},
+                { 84 ,255 ,255},
+                {170 ,  0 ,255},
+                {170 , 84 ,255},
+                {170 ,170 ,255},
+                {170 ,255 ,255},
+                {255 ,  0 ,255},
+                {255 , 84 ,255},
+                {255 ,170 ,255},
+                { 42 ,  0 ,  0},
+                { 84 ,  0 ,  0},
+                {127 ,  0 ,  0},
+                {170 ,  0 ,  0},
+                {212 ,  0 ,  0},
+                {255 ,  0 ,  0},
+                {  0 , 42 ,  0},
+                {  0 , 84 ,  0},
+                {  0 ,127 ,  0},
+                {  0 ,170 ,  0},
+                {  0 ,212 ,  0},
+                {  0 ,255 ,  0},
+                {  0 ,  0 , 42},
+                {  0 ,  0 , 84},
+                {  0 ,  0 ,127},
+                {  0 ,  0 ,170},
+                {  0 ,  0 ,212},
+                {  0 ,  0 ,255},
+                {  0 ,  0 ,  0},
+                { 36 , 36 , 36},
+                { 72 , 72 , 72},
+                {109 ,109 ,109},
+                {145 ,145 ,145},
+                {182 ,182 ,182},
+                {218 ,218 ,218},
+                {  0 ,113 ,188},
+                { 80 ,182 ,188},
+                {127 ,127 ,  0},
+        };
+
+void draw_bboxes(const cv::Mat& bgr, const std::vector<BoxInfo>& bboxes, object_rect effect_roi, std::vector<std::string> class_names) {
+    cv::Mat image = bgr.clone();
+    int src_w = image.cols;
+    int src_h = image.rows;
+    int dst_w = effect_roi.width;
+    int dst_h = effect_roi.height;
+    float width_ratio = (float)src_w / (float)dst_w;
+    float height_ratio = (float)src_h / (float)dst_h;
+
+    for (const auto & bbox : bboxes)
+    {
+        cv::Scalar color = cv::Scalar(color_list[bbox.label][0], color_list[bbox.label][1], color_list[bbox.label][2]);
+        //fprintf(stderr, "%d = %.5f at %.2f %.2f %.2f %.2f\n", bbox.label, bbox.score,
+        //    bbox.x1, bbox.y1, bbox.x2, bbox.y2);
+
+        cv::rectangle(image, cv::Rect(cv::Point((bbox.x1 - effect_roi.x) * width_ratio, (bbox.y1 - effect_roi.y) * height_ratio),
+                                      cv::Point((bbox.x2 - effect_roi.x) * width_ratio, (bbox.y2 - effect_roi.y) * height_ratio)), color);
+
+        char text[256];
+        sprintf(text, "%s %.1f%%", class_names[bbox.label].c_str(), bbox.score * 100);
+
+        int baseLine = 0;
+        cv::Size label_size = cv::getTextSize(text, cv::FONT_HERSHEY_SIMPLEX, 0.4, 1, &baseLine);
+
+        int x = (bbox.x1 - effect_roi.x) * width_ratio;
+        int y = (bbox.y1 - effect_roi.y) * height_ratio - label_size.height - baseLine;
+        if (y < 0)
+            y = 0;
+        if (x + label_size.width > image.cols)
+            x = image.cols - label_size.width;
+
+        cv::rectangle(image, cv::Rect(cv::Point(x, y), cv::Size(label_size.width, label_size.height + baseLine)),
+                      color, -1);
+
+        cv::putText(image, text, cv::Point(x, y + label_size.height),
+                    cv::FONT_HERSHEY_SIMPLEX, 0.4, cv::Scalar(255, 255, 255));
+    }
+
+    cv::imshow("image", image);
+}
+
+
 int main() {
     cv::Mat img;
     std::string path2 = "/home/ubuntu/Documents/cpp/yolov5/dog.jpg";
@@ -94,25 +221,19 @@ int main() {
         return -1;
     }
 
-    Yolov5 det("/home/ubuntu/Documents/cpp/yolov5/ncnn/yolov5s.param",
-               "/home/ubuntu/Documents/cpp/yolov5/ncnn/yolov5s.bin",
+    Yolov5 detector("/home/ubuntu/Documents/cpp/yolov5/ncnn/yolov5s_opt.param",
+               "/home/ubuntu/Documents/cpp/yolov5/ncnn/yolov5s_opt.bin",
                true);
 
     cv::Mat resize_img;
-    auto input_size = det.input_size;
+    auto input_size = detector.input_size;
     object_rect eff{};
     resize_uniform(img, resize_img, input_size, eff);
 
-    std::vector<BoxInfo> dets = det.detect(resize_img, 0.5, 0.5);
+    std::vector<BoxInfo> dets = detector.detect(resize_img, 0.5, 0.5);
 
-    for (auto& det_box: dets){
-        auto pt1 = cv::Point2i(int(det_box.x1), int(det_box.y1));
-        auto pt2 = cv::Point2i(int(det_box.x2), int(det_box.y2));
-        cv::rectangle(resize_img, pt1, pt2, cv::Scalar(0, 0, 255), 2, cv::LINE_4);
-    }
-    cv::imshow("res", resize_img);
+    draw_bboxes(img, dets, eff, detector.labels);
     cv::waitKey(0);
-
     return 0;
 }