From 2264b3a76c9bb75c90619cc6c82739dbbea476cc Mon Sep 17 00:00:00 2001
From: syyxsxx <csuwsy@163.com>
Date: Mon, 31 Aug 2020 01:16:49 +0800
Subject: [PATCH] fix mask

---
 deploy/cpp/include/paddlex/results.h |  2 +-
 deploy/cpp/src/paddlex.cpp           | 18 ++++++------------
 deploy/cpp/src/visualize.cpp         |  3 ++-
 3 files changed, 9 insertions(+), 14 deletions(-)

diff --git a/deploy/cpp/include/paddlex/results.h b/deploy/cpp/include/paddlex/results.h
index e3526bf..75a857d 100644
--- a/deploy/cpp/include/paddlex/results.h
+++ b/deploy/cpp/include/paddlex/results.h
@@ -47,7 +47,7 @@ struct Box {
   // confidence score
   float score;
   std::vector<float> coordinate;
-  Mask<int> mask;
+  Mask<float> mask;
 };
 
 /*
diff --git a/deploy/cpp/src/paddlex.cpp b/deploy/cpp/src/paddlex.cpp
index a3bf1dc..fd82792 100644
--- a/deploy/cpp/src/paddlex.cpp
+++ b/deploy/cpp/src/paddlex.cpp
@@ -348,10 +348,7 @@ bool Model::predict(const cv::Mat& im, DetResult* result) {
       auto begin_mask =
           output_mask.begin() + (i * classes + box->category_id) * mask_pixels;
       auto end_mask = begin_mask + mask_pixels;
-      for (auto iter = begin_mask; iter != end_mask; iter++) {
-        int mask_int = floor((*iter) + 0.5);
-        box->mask.data.push_back(mask_int);
-      }
+      box->mask.data.assign(begin_mask, end_mask);
       box->mask.shape = {static_cast<int>(box->coordinate[2]),
                          static_cast<int>(box->coordinate[3])};
     }
@@ -520,17 +517,14 @@ bool Model::predict(const std::vector<cv::Mat>& im_batch,
     for (int i = 0; i < lod_vector[0].size() - 1; ++i) {
       (*results)[i].mask_resolution = output_mask_shape[2];
       for (int j = 0; j < (*results)[i].boxes.size(); ++j) {
-        Box* box = &(*results)[i].boxes[j];
+        Box* box = &result->boxes[i];
         int category_id = box->category_id;
-        auto begin_mask = output_mask.begin() +
-                          (mask_idx * classes + category_id) * mask_pixels;
+        auto begin_mask =
+          output_mask.begin() + (i * classes + box->category_id) * mask_pixels;
         auto end_mask = begin_mask + mask_pixels;
-        for (auto iter = begin_mask; iter != end_mask; iter++) {
-          int mask_int = floor((*iter) + 0.5);
-          box->mask.data.push_back(mask_int);
-        }
+        box->mask.data.assign(begin_mask, end_mask);
         box->mask.shape = {static_cast<int>(box->coordinate[2]),
-                           static_cast<int>(box->coordinate[3])};
+                         static_cast<int>(box->coordinate[3])};
         mask_idx++;
       }
     }
diff --git a/deploy/cpp/src/visualize.cpp b/deploy/cpp/src/visualize.cpp
index 7beba65..4eda4a3 100644
--- a/deploy/cpp/src/visualize.cpp
+++ b/deploy/cpp/src/visualize.cpp
@@ -86,11 +86,12 @@ cv::Mat Visualize(const cv::Mat& img,
     }
     cv::Mat bin_mask(result.mask_resolution,
                      result.mask_resolution,
-                     CV_8UC1,
+                     CV_32FC1
                      boxes[i].mask.data.data());
     cv::resize(bin_mask,
                bin_mask,
                cv::Size(boxes[i].mask.shape[0], boxes[i].mask.shape[1]));
+    cv::threshold(bin_mask, bin_mask, 0.5, 1, cv::THRESH_BINARY);
     cv::Mat full_mask = cv::Mat::zeros(vis_img.size(), CV_8UC1);
     bin_mask.copyTo(full_mask(roi));
     cv::Mat mask_ch[3];
-- 
GitLab