fix(RVM): fixed the TNN preprocess for RVM (#240)

DefTruth · Mar 17, 2022 · 641e56b · 641e56b
1 parent 0989025
commit 641e56b
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 8 deletions.
diff --git a/examples/lite/cv/test_lite_rvm.cpp b/examples/lite/cv/test_lite_rvm.cpp
@@ -100,7 +100,7 @@ static void test_lite()
   test_default();
   test_onnxruntime();
   test_mnn();
-  test_ncnn();
+  // test_ncnn();
   test_tnn();
 }
 

diff --git a/lite/tnn/cv/tnn_rvm.cpp b/lite/tnn/cv/tnn_rvm.cpp
@@ -178,17 +178,18 @@ void TNNRobustVideoMatting::initialize_context()
   context_is_initialized = true;
 }
 
-void TNNRobustVideoMatting::transform(const cv::Mat &mat)
+void TNNRobustVideoMatting::transform(const cv::Mat &mat_rs)
 {
-  cv::Mat canvas;
-  cv::resize(mat, canvas, cv::Size(input_width, input_height));
-  cv::cvtColor(canvas, canvas, cv::COLOR_BGR2RGB);
+  //  cv::Mat canvas;
+  //  cv::resize(mat, canvas, cv::Size(input_width, input_height));
+  //  cv::cvtColor(canvas, canvas, cv::COLOR_BGR2RGB);
+  //  reference: https://github.com/DefTruth/lite.ai.toolkit/issues/240
   // push into src_mat
   src_mat = std::make_shared<tnn::Mat>(
       input_device_type,
       tnn::N8UC3,
       input_shapes.at("src"),
-      (void *) canvas.data
+      (void *) mat_rs.data
   );
   if (!src_mat->GetData())
   {
@@ -206,7 +207,12 @@ void TNNRobustVideoMatting::detect(const cv::Mat &mat, types::MattingContent &co
   if (!context_is_initialized) return;
 
   // 1. make input tensor
-  this->transform(mat);
+  cv::Mat mat_rs;
+  // resize mat outside 'transform' to prevent memory overflow
+  // reference: https://github.com/DefTruth/lite.ai.toolkit/issues/240
+  cv::resize(mat, mat_rs, cv::Size(input_width, input_height));
+  cv::cvtColor(mat_rs, mat_rs, cv::COLOR_BGR2RGB);
+  this->transform(mat_rs);
   // 2. set input_mat
   tnn::MatConvertParam src_cvt_param, ctx_cvt_param;
   src_cvt_param.scale = scale_vals;

diff --git a/lite/tnn/cv/tnn_rvm.h b/lite/tnn/cv/tnn_rvm.h
@@ -86,7 +86,7 @@ namespace tnncv
     void print_debug_string(); // debug information
 
   private:
-    void transform(const cv::Mat &mat); //
+    void transform(const cv::Mat &mat_rs); //
 
     void initialize_instance(); // init net & instance