From 62daf2fbf9d12aa552f8e3c692ee93057059664f Mon Sep 17 00:00:00 2001
From: jack <136876878@qq.com>
Date: Mon, 13 Jul 2020 11:35:17 +0800
Subject: [PATCH] encapsulate GenerateColormap in Visualize

---
 deploy/cpp/demo/detector.cpp           | 4 ++--
 deploy/cpp/demo/segmenter.cpp          | 4 ++--
 deploy/cpp/include/paddlex/visualize.h | 7 ++-----
 deploy/cpp/src/visualize.cpp           | 6 +++---
 4 files changed, 9 insertions(+), 12 deletions(-)

diff --git a/deploy/cpp/demo/detector.cpp b/deploy/cpp/demo/detector.cpp
index 54f93d2..dc89396 100644
--- a/deploy/cpp/demo/detector.cpp
+++ b/deploy/cpp/demo/detector.cpp
@@ -124,7 +124,7 @@ int main(int argc, char** argv) {
       // 可视化
       for (int j = 0; j < im_vec_size - i; ++j) {
         cv::Mat vis_img = PaddleX::Visualize(
-            im_vec[j], results[j], model.labels, colormap, FLAGS_threshold);
+            im_vec[j], results[j], model.labels, FLAGS_threshold);
         std::string save_path =
             PaddleX::generate_save_path(FLAGS_save_dir, image_paths[i + j]);
         cv::imwrite(save_path, vis_img);
@@ -155,7 +155,7 @@ int main(int argc, char** argv) {
 
     // 可视化
     cv::Mat vis_img =
-        PaddleX::Visualize(im, result, model.labels, colormap, FLAGS_threshold);
+        PaddleX::Visualize(im, result, model.labels, FLAGS_threshold);
     std::string save_path =
         PaddleX::generate_save_path(FLAGS_save_dir, FLAGS_image);
     cv::imwrite(save_path, vis_img);
diff --git a/deploy/cpp/demo/segmenter.cpp b/deploy/cpp/demo/segmenter.cpp
index 90adb5a..4b3905a 100644
--- a/deploy/cpp/demo/segmenter.cpp
+++ b/deploy/cpp/demo/segmenter.cpp
@@ -106,7 +106,7 @@ int main(int argc, char** argv) {
       // 可视化
       for (int j = 0; j < im_vec_size - i; ++j) {
         cv::Mat vis_img =
-            PaddleX::Visualize(im_vec[j], results[j], model.labels, colormap);
+            PaddleX::Visualize(im_vec[j], results[j], model.labels);
         std::string save_path =
             PaddleX::generate_save_path(FLAGS_save_dir, image_paths[i + j]);
         cv::imwrite(save_path, vis_img);
@@ -124,7 +124,7 @@ int main(int argc, char** argv) {
                             microseconds::period::num /
                             microseconds::period::den;
     // 可视化
-    cv::Mat vis_img = PaddleX::Visualize(im, result, model.labels, colormap);
+    cv::Mat vis_img = PaddleX::Visualize(im, result, model.labels);
     std::string save_path =
         PaddleX::generate_save_path(FLAGS_save_dir, FLAGS_image);
     cv::imwrite(save_path, vis_img);
diff --git a/deploy/cpp/include/paddlex/visualize.h b/deploy/cpp/include/paddlex/visualize.h
index 9ddba53..9b80ca3 100644
--- a/deploy/cpp/include/paddlex/visualize.h
+++ b/deploy/cpp/include/paddlex/visualize.h
@@ -65,13 +65,12 @@ std::vector<int> GenerateColorMap(int num_class);
  * @param img: initial image matrix
  * @param results: the detection result
  * @param labels: label map
- * @param colormap: visualization color map
+ * @param threshold: minimum confidence to display
  * @return visualized image matrix
  * */
 cv::Mat Visualize(const cv::Mat& img,
                      const DetResult& results,
                      const std::map<int, std::string>& labels,
-                     const std::vector<int>& colormap,
                      float threshold = 0.5);
 
 /*
@@ -81,13 +80,11 @@ cv::Mat Visualize(const cv::Mat& img,
  * @param img: initial image matrix
  * @param results: the detection result
  * @param labels: label map
- * @param colormap: visualization color map
  * @return visualized image matrix
  * */
 cv::Mat Visualize(const cv::Mat& img,
                      const SegResult& result,
-                     const std::map<int, std::string>& labels,
-                     const std::vector<int>& colormap);
+                     const std::map<int, std::string>& labels);
 
 /*
  * @brief
diff --git a/deploy/cpp/src/visualize.cpp b/deploy/cpp/src/visualize.cpp
index 1511887..afc1733 100644
--- a/deploy/cpp/src/visualize.cpp
+++ b/deploy/cpp/src/visualize.cpp
@@ -34,8 +34,8 @@ std::vector<int> GenerateColorMap(int num_class) {
 cv::Mat Visualize(const cv::Mat& img,
                      const DetResult& result,
                      const std::map<int, std::string>& labels,
-                     const std::vector<int>& colormap,
                      float threshold) {
+  auto colormap = GenerateColorMap(labels.size());
   cv::Mat vis_img = img.clone();
   auto boxes = result.boxes;
   for (int i = 0; i < boxes.size(); ++i) {
@@ -107,8 +107,8 @@ cv::Mat Visualize(const cv::Mat& img,
 
 cv::Mat Visualize(const cv::Mat& img,
                      const SegResult& result,
-                     const std::map<int, std::string>& labels,
-                     const std::vector<int>& colormap) {
+                     const std::map<int, std::string>& labels) {
+  auto colormap = GenerateColorMap(labels.size());
   std::vector<uint8_t> label_map(result.label_map.data.begin(),
                                  result.label_map.data.end());
   cv::Mat mask(result.label_map.shape[0],
-- 
GitLab