fix 'cvtcolor' error in the preprocessing of single channel images (o…

…pen-mmlab#1666) (cherry picked from commit 8bb3fcc) (cherry picked from commit d91ba2f)
irexyc · Apr 4, 2023 · 26a65ff · 26a65ff
1 parent 5107242
commit 26a65ff
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 5 deletions.
diff --git a/csrc/mmdeploy/operation/cuda/cvtcolor.cpp b/csrc/mmdeploy/operation/cuda/cvtcolor.cpp
@@ -90,19 +90,20 @@ class CvtColorImpl : public CvtColor {
 
     auto height = src.height();
     auto width = src.width();
-    auto channels = src.channel();
-    auto stride = width * channels;
+    auto src_channels = src.channel();
+    auto src_stride = width * src_channels;
 
     Mat dst_mat(height, width, dst_fmt, src.type(), device());
+    auto dst_stride = width * dst_mat.channel();
 
     auto convert = [&](auto type) -> Result<void> {
       using T = typename decltype(type)::type;
       auto converter = GetConverter<T>(src.pixel_format(), dst_fmt);
       if (!converter) {
         return Status(eNotSupported);
       }
-      auto ret =
-          converter(cuda_stream, height, width, stride, src.data<T>(), stride, dst_mat.data<T>());
+      auto ret = converter(cuda_stream, height, width, src_stride, src.data<T>(), dst_stride,
+                           dst_mat.data<T>());
       if (ret != ppl::common::RC_SUCCESS) {
         return Status(eFail);
       }

diff --git a/csrc/mmdeploy/preprocess/transform/load.cpp b/csrc/mmdeploy/preprocess/transform/load.cpp
@@ -55,7 +55,7 @@ class PrepareImage : public Transform {
     if (color_type_ == "color" || color_type_ == "color_ignore_orientation") {
       OUTCOME_TRY(cvt_color_.Apply(src_mat, dst_mat, PixelFormat::kBGR));
     } else {
-      OUTCOME_TRY(cvt_color_.Apply(dst_mat, dst_mat, PixelFormat::kGRAYSCALE));
+      OUTCOME_TRY(cvt_color_.Apply(src_mat, dst_mat, PixelFormat::kGRAYSCALE));
     }
     auto tensor = to_tensor(dst_mat);
     if (to_float32_) {