From e3ec62bd8d7818e21a0120d00d66168250306e3e Mon Sep 17 00:00:00 2001 From: ttjjmm Date: Tue, 18 Jan 2022 21:01:23 +0800 Subject: [PATCH] add visualize --- main.cpp | 143 ++++++++++++++++++++++++++++++++++++++++++++++++++----- 1 file changed, 132 insertions(+), 11 deletions(-) diff --git a/main.cpp b/main.cpp index 9a94e47..0bd7fe9 100644 --- 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& bboxes, object_rect effect_roi, std::vector 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 dets = det.detect(resize_img, 0.5, 0.5); + std::vector 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; }