Fix WarpBbox and memory leak in TextRecognizer (open-mmlab#1745)

lzhangzz · Feb 13, 2023 · cadc265 · cadc265
1 parent 0f5b149
commit cadc265
Show file tree

Hide file tree

Showing 4 changed files with 22 additions and 15 deletions.
diff --git a/csrc/mmdeploy/apis/cxx/mmdeploy/text_recognizer.hpp b/csrc/mmdeploy/apis/cxx/mmdeploy/text_recognizer.hpp
@@ -3,6 +3,8 @@
 #ifndef MMDEPLOY_CSRC_MMDEPLOY_APIS_CXX_TEXT_RECOGNIZER_HPP_
 #define MMDEPLOY_CSRC_MMDEPLOY_APIS_CXX_TEXT_RECOGNIZER_HPP_
 
+#include <numeric>
+
 #include "mmdeploy/common.hpp"
 #include "mmdeploy/text_detector.hpp"
 #include "mmdeploy/text_recognizer.h"
@@ -40,9 +42,12 @@ class TextRecognizer : public NonMovable {
     const TextDetection* p_bboxes{};
     const int* p_bbox_count{};
 
+    auto n_total_bboxes = static_cast<int>(images.size());
+
     if (!bboxes.empty()) {
       p_bboxes = bboxes.data();
       p_bbox_count = bbox_count.data();
+      n_total_bboxes = std::accumulate(bbox_count.begin(), bbox_count.end(), 0);
     }
 
     TextRecognition* results{};
@@ -53,7 +58,7 @@ class TextRecognizer : public NonMovable {
       throw_exception(static_cast<ErrorCode>(ec));
     }
 
-    std::shared_ptr<TextRecognition> data(results, [count = images.size()](auto p) {
+    std::shared_ptr<TextRecognition> data(results, [count = n_total_bboxes](auto p) {
       mmdeploy_text_recognizer_release_result(p, count);
     });
 

diff --git a/csrc/mmdeploy/codebase/mmocr/warp.cpp b/csrc/mmdeploy/codebase/mmocr/warp.cpp
@@ -21,7 +21,7 @@ class WarpBbox {
     if (det.is_object() && det.contains("bbox")) {
       auto bbox = from_value<std::vector<cv::Point>>(det["bbox"]);
       auto patch = warp(mmdeploy::cpu::Mat2CVMat(ori_img), bbox);
-      return Value{{"ori_img", cpu::CVMat2Mat(patch, PixelFormat::kBGR)}};
+      return Value{{"ori_img", cpu::CVMat2Mat(patch, ori_img.pixel_format())}};
     } else {  // whole image as a bbox
       return Value{{"ori_img", ori_img}};
     }

diff --git a/csrc/mmdeploy/utils/opencv/opencv_utils.cpp b/csrc/mmdeploy/utils/opencv/opencv_utils.cpp
@@ -38,11 +38,11 @@ Mat CVMat2Mat(const cv::Mat& mat, PixelFormat format) {
 }
 
 cv::Mat Mat2CVMat(const Mat& mat) {
-  std::map<DataType, int> type_mapper{{DataType::kFLOAT, CV_32F},
-                                      {DataType::kHALF, CV_16U},
-                                      {DataType::kINT8, CV_8U},
-                                      {DataType::kINT32, CV_32S}};
-  auto type = CV_MAKETYPE(type_mapper[mat.type()], mat.channel());
+  static const std::map<DataType, int> type_mapper{{DataType::kFLOAT, CV_32F},
+                                                   {DataType::kHALF, CV_16U},
+                                                   {DataType::kINT8, CV_8U},
+                                                   {DataType::kINT32, CV_32S}};
+  auto type = CV_MAKETYPE(type_mapper.at(mat.type()), mat.channel());
   auto format = mat.pixel_format();
   if (PixelFormat::kBGR == format || PixelFormat::kRGB == format) {
     return cv::Mat(mat.height(), mat.width(), type, mat.data<void>());

diff --git a/demo/csrc/cpp/utils/visualize.h b/demo/csrc/cpp/utils/visualize.h
@@ -64,9 +64,11 @@ class Visualize {
       cv::Rect rect(origin + cv::Point2f(0, text_size.height + 2 * thickness),
                     origin + cv::Point2f(text_size.width, 0));
       rect &= cv::Rect({}, img_.size());
-      img_(rect) *= .35f;
-      cv::putText(img_, text, origin + cv::Point2f(0, text_size.height), font_face, font_scale,
-                  cv::Scalar::all(255), thickness, cv::LINE_AA);
+      if (rect.area() > 0) {
+        img_(rect) *= .35f;
+        cv::putText(img_, text, origin + cv::Point2f(0, text_size.height), font_face, font_scale,
+                    cv::Scalar::all(255), thickness, cv::LINE_AA);
+      }
       return text_size.height;
     }
 
@@ -120,11 +122,11 @@ class Visualize {
 
     void add_text_det(mmdeploy_point_t bbox[4], float score, const char* text, size_t text_size,
                       int index) {
-      printf("bbox[%d]: (%.2f, %.2f), (%.2f, %.2f), (%.2f, %.2f), (%.2f, %.2f)\n", index,  //
-             bbox[0].x, bbox[0].y,                                                         //
-             bbox[1].x, bbox[1].y,                                                         //
-             bbox[2].x, bbox[2].y,                                                         //
-             bbox[3].x, bbox[3].y);
+      printf("bbox[%d]: (%.2f, %.2f), (%.2f, %.2f), (%.2f, %.2f), (%.2f, %.2f), %.2f\n", index,  //
+             bbox[0].x, bbox[0].y,                                                               //
+             bbox[1].x, bbox[1].y,                                                               //
+             bbox[2].x, bbox[2].y,                                                               //
+             bbox[3].x, bbox[3].y, score);
       std::vector<cv::Point> poly_points;
       cv::Point2f center{};
       for (int i = 0; i < 4; ++i) {