Bläddra i källkod

encapsulate GenerateColormap in Visualize

jack 5 år sedan
förälder
incheckning
62daf2fbf9

+ 2 - 2
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);

+ 2 - 2
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);

+ 2 - 5
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

+ 3 - 3
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],