syyxsxx пре 5 година
родитељ
комит
d10c89e9d6
3 измењених фајлова са 32 додато и 16 уклоњено
  1. 1 1
      deploy/cpp/include/paddlex/results.h
  2. 28 8
      deploy/cpp/src/paddlex.cpp
  3. 3 7
      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<float> mask;
+  Mask<int> mask;
 };
 
 /*

+ 28 - 8
deploy/cpp/src/paddlex.cpp

@@ -345,12 +345,22 @@ bool Model::predict(const cv::Mat& im, DetResult* result) {
     result->mask_resolution = output_mask_shape[2];
     for (int i = 0; i < result->boxes.size(); ++i) {
       Box* box = &result->boxes[i];
-      auto begin_mask =
-          output_mask.begin() + (i * classes + box->category_id) * mask_pixels;
-      auto end_mask = begin_mask + mask_pixels;
-      box->mask.data.assign(begin_mask, end_mask);
       box->mask.shape = {static_cast<int>(box->coordinate[2]),
                          static_cast<int>(box->coordinate[3])};
+      auto begin_mask =
+          output_mask.begin() + (i * classes + box->category_id) * mask_pixels;
+      cv::Mat bin_mask(result->mask_resolution,
+                     result->mask_resolution,
+                     CV_32FC1,
+                     begin_mask);
+      cv::resize(bin_mask,
+               bin_mask,
+               cv::Size(box->mask.shape[0], box->mask.shape[1]));
+      cv::threshold(bin_mask, bin_mask, 0.5, 1, cv::THRESH_BINARY);
+      auto mask_int_begin = bin_mask.data;
+      auto mask_int_end =
+        mask_int_begin + box->mask.shape[0] * box->mask.shape[1];
+      box->mask.data.assign(mask_int_begin, mask_int_end);
     }
   }
   return true;
@@ -519,12 +529,22 @@ bool Model::predict(const std::vector<cv::Mat>& im_batch,
       for (int j = 0; j < (*results)[i].boxes.size(); ++j) {
         Box* box = &result->boxes[i];
         int category_id = box->category_id;
+        box->mask.shape = {static_cast<int>(box->coordinate[2]),
+                          static_cast<int>(box->coordinate[3])};
         auto begin_mask =
           output_mask.begin() + (i * classes + box->category_id) * mask_pixels;
-        auto end_mask = begin_mask + mask_pixels;
-        box->mask.data.assign(begin_mask, end_mask);
-        box->mask.shape = {static_cast<int>(box->coordinate[2]),
-                         static_cast<int>(box->coordinate[3])};
+        cv::Mat bin_mask(result->mask_resolution,
+                      result->mask_resolution,
+                      CV_32FC1,
+                      begin_mask);
+        cv::resize(bin_mask,
+                bin_mask,
+                cv::Size(box->mask.shape[0], box->mask.shape[1]));
+        cv::threshold(bin_mask, bin_mask, 0.5, 1, cv::THRESH_BINARY);
+        auto mask_int_begin = bin_mask.data;
+        auto mask_int_end =
+          mask_int_begin + box->mask.shape[0] * box->mask.shape[1];
+        box->mask.data.assign(mask_int_begin, mask_int_end);
         mask_idx++;
       }
     }

+ 3 - 7
deploy/cpp/src/visualize.cpp

@@ -84,14 +84,10 @@ cv::Mat Visualize(const cv::Mat& img,
     if (boxes[i].mask.data.size() == 0) {
       continue;
     }
-    cv::Mat bin_mask(result.mask_resolution,
-                     result.mask_resolution,
-                     CV_32FC1
+    cv::Mat bin_mask(boxes[i].mask.shape[0],
+                     boxes[i].mask.shape[1],
+                     CV_8UC1,
                      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];