syyxsxx 5 년 전
부모
커밋
2264b3a76c
3개의 변경된 파일9개의 추가작업 그리고 14개의 파일을 삭제
  1. 1 1
      deploy/cpp/include/paddlex/results.h
  2. 6 12
      deploy/cpp/src/paddlex.cpp
  3. 2 1
      deploy/cpp/src/visualize.cpp

+ 1 - 1
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;
 };
 
 /*

+ 6 - 12
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++;
       }
     }

+ 2 - 1
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];