From 37ee101b639892b908cb6775e9a70c1d38321757 Mon Sep 17 00:00:00 2001
From: WJJ1995 <wjjisloser@163.com>
Date: Tue, 15 Nov 2022 17:06:16 +0800
Subject: [PATCH] [Bug Fix] Fixed MattingResult bug (#593)

* add onnx_ort_runtime demo

* rm in requirements

* support batch eval

* fixed MattingResults bug
---
 fastdeploy/vision/common/result.cc | 4 ++--
 tests/models/test_rvm.py           | 8 +++++++-
 2 files changed, 9 insertions(+), 3 deletions(-)

diff --git a/fastdeploy/vision/common/result.cc b/fastdeploy/vision/common/result.cc
index 3585713a67..4d746797f4 100755
--- a/fastdeploy/vision/common/result.cc
+++ b/fastdeploy/vision/common/result.cc
@@ -372,7 +372,7 @@ void MattingResult::Reserve(int size) {
   if (contain_foreground) {
     FDASSERT((shape.size() == 3),
              "Please initial shape (h,w,c) before call Reserve.");
-    int c = static_cast<int>(shape[3]);
+    int c = static_cast<int>(shape[2]);
     foreground.reserve(size * c);
   }
 }
@@ -382,7 +382,7 @@ void MattingResult::Resize(int size) {
   if (contain_foreground) {
     FDASSERT((shape.size() == 3),
              "Please initial shape (h,w,c) before call Resize.");
-    int c = static_cast<int>(shape[3]);
+    int c = static_cast<int>(shape[2]);
     foreground.resize(size * c);
   }
 }
diff --git a/tests/models/test_rvm.py b/tests/models/test_rvm.py
index 10d680948f..4fa3083e59 100644
--- a/tests/models/test_rvm.py
+++ b/tests/models/test_rvm.py
@@ -19,6 +19,7 @@
 import numpy as np
 import runtime_config as rc
 
+
 def test_matting_rvm_cpu():
     model_url = "https://bj.bcebos.com/paddlehub/fastdeploy/rvm.tgz"
     input_url = "https://bj.bcebos.com/paddlehub/fastdeploy/video.mp4"
@@ -38,7 +39,8 @@ def test_matting_rvm_cpu():
             break
         result = model.predict(frame)
         # compare diff
-        expect_alpha = np.load("resources/rvm/result_alpha_" + str(frame_id) + ".npy")
+        expect_alpha = np.load("resources/rvm/result_alpha_" + str(frame_id) +
+                               ".npy")
         result_alpha = np.array(result.alpha).reshape(1920, 1080)
         diff = np.fabs(expect_alpha - result_alpha)
         thres = 1e-05
@@ -51,3 +53,7 @@ def test_matting_rvm_cpu():
             cap.release()
             cv2.destroyAllWindows()
             break
+
+
+if __name__ == "__main__":
+    test_matting_rvm_cpu()