From 5dd651f1a5a9405faf4339759bdda9cfbd8cadaa Mon Sep 17 00:00:00 2001 From: xulei2020 Date: Thu, 14 Jan 2021 09:26:18 +0800 Subject: [PATCH] fix bug for extractchannel when type is UINT16 --- .../dataset/kernels/image/lite_cv/image_process.cc | 6 +++--- tests/ut/cpp/dataset/image_process_test.cc | 12 ++++++++++-- 2 files changed, 13 insertions(+), 5 deletions(-) diff --git a/mindspore/ccsrc/minddata/dataset/kernels/image/lite_cv/image_process.cc b/mindspore/ccsrc/minddata/dataset/kernels/image/lite_cv/image_process.cc index 7c13fb3f12..db7142b479 100644 --- a/mindspore/ccsrc/minddata/dataset/kernels/image/lite_cv/image_process.cc +++ b/mindspore/ccsrc/minddata/dataset/kernels/image/lite_cv/image_process.cc @@ -674,16 +674,16 @@ bool ExtractChannel(LiteMat &src, LiteMat &dst, int col) { return false; } - if (dst.data_type_ == LDataType::FLOAT32 || dst.data_type_ == LDataType::UINT8) { + if (src.data_type_ == LDataType::FLOAT32 || src.data_type_ == LDataType::UINT8) { if (dst.IsEmpty() || dst.width_ != src.width_ || dst.height_ != src.height_ || dst.channel_ != 1 || dst.data_type_ != src.data_type_) { dst.Init(src.width_, src.height_, 1, src.data_type_); } } - if (dst.data_type_ == LDataType::FLOAT32) { + if (src.data_type_ == LDataType::FLOAT32) { ExtractChannelImpl(src, dst, src.height_, src.width_, src.channel_, col); - } else if (dst.data_type_ == LDataType::UINT8) { + } else if (src.data_type_ == LDataType::UINT8) { ExtractChannelImpl(src, dst, src.height_, src.width_, src.channel_, col); } else { return false; diff --git a/tests/ut/cpp/dataset/image_process_test.cc b/tests/ut/cpp/dataset/image_process_test.cc index 1d6b9dcc67..a86a9e6bf8 100644 --- a/tests/ut/cpp/dataset/image_process_test.cc +++ b/tests/ut/cpp/dataset/image_process_test.cc @@ -408,7 +408,7 @@ TEST_F(MindDataImageProcess, TestPadd) { size_t total_size = makeborder.height_ * makeborder.width_ * makeborder.channel_; double distance = 0.0f; for (size_t i = 0; i < total_size; i++) { - distance += pow((uint8_t)b_image.data[i] - ((uint8_t*)makeborder)[i], 2); + distance += pow((uint8_t)b_image.data[i] - ((uint8_t *)makeborder)[i], 2); } distance = sqrt(distance / total_size); EXPECT_EQ(distance, 0.0f); @@ -439,7 +439,7 @@ TEST_F(MindDataImageProcess, TestPadZero) { size_t total_size = makeborder.height_ * makeborder.width_ * makeborder.channel_; double distance = 0.0f; for (size_t i = 0; i < total_size; i++) { - distance += pow((uint8_t)b_image.data[i] - ((uint8_t*)makeborder)[i], 2); + distance += pow((uint8_t)b_image.data[i] - ((uint8_t *)makeborder)[i], 2); } distance = sqrt(distance / total_size); EXPECT_EQ(distance, 0.0f); @@ -879,3 +879,11 @@ TEST_F(MindDataImageProcess, TestMultiplyFloat) { static_cast(dst_float.data_ptr_)[i].c1); } } + +TEST_F(MindDataImageProcess, TestExtractChannel) { + LiteMat lite_single; + LiteMat lite_mat = LiteMat(1, 4, 3, LDataType::UINT16); + + EXPECT_FALSE(ExtractChannel(lite_mat, lite_single, 0)); + EXPECT_TRUE(lite_single.IsEmpty()); +} \ No newline at end of file