From 0959c54ae90e68a3bb665fc4a25721a2fe484c22 Mon Sep 17 00:00:00 2001
From: liyong <liyong126@huawei.com>
Date: Thu, 9 Sep 2021 19:50:57 +0800
Subject: [PATCH] fix bug in random posterize

---
 .../ccsrc/minddata/dataset/core/cv_tensor.cc  |  3 +-
 .../dataset/kernels/image/posterize_op.cc     |  1 +
 .../cpp/dataset/random_posterize_op_test.cc   | 33 +++++++++++++++++--
 3 files changed, 34 insertions(+), 3 deletions(-)

diff --git a/mindspore/ccsrc/minddata/dataset/core/cv_tensor.cc b/mindspore/ccsrc/minddata/dataset/core/cv_tensor.cc
index 48980fb929a..424f3905239 100644
--- a/mindspore/ccsrc/minddata/dataset/core/cv_tensor.cc
+++ b/mindspore/ccsrc/minddata/dataset/core/cv_tensor.cc
@@ -56,7 +56,8 @@ Status CVTensor::CreateFromMat(const cv::Mat &mat, const dsize_t rank, CVTensorP
   } else if (mat.dims == 2 && rank == 3) {
     shape = TensorShape({mat.rows, mat.cols, mat.channels()});
   } else {
-    RETURN_STATUS_UNEXPECTED("Error in creating CVTensor: Invalid input rank or cv::mat dimension.");
+    // the info of <C, H, W> tensor is: dims = 3, size = (C, H, W), channels = 1
+    RETURN_STATUS_UNEXPECTED("CreateFromMat: tensor should be in shape of <H,W,C> or <H,W>.");
   }
   DataType type = DataType::FromCVType(mat_local.type());
   RETURN_IF_NOT_OK(CreateFromMemory(shape, type, mat_local.data, &out_tensor));
diff --git a/mindspore/ccsrc/minddata/dataset/kernels/image/posterize_op.cc b/mindspore/ccsrc/minddata/dataset/kernels/image/posterize_op.cc
index 9757ee1c5a3..b4e29b8d798 100644
--- a/mindspore/ccsrc/minddata/dataset/kernels/image/posterize_op.cc
+++ b/mindspore/ccsrc/minddata/dataset/kernels/image/posterize_op.cc
@@ -39,6 +39,7 @@ Status PosterizeOp::Compute(const std::shared_ptr<Tensor> &input, std::shared_pt
     lut_vector.push_back(i & mask_value);
   }
   cv::Mat in_image = input_cv->mat();
+
   cv::Mat output_img;
   CHECK_FAIL_RETURN_UNEXPECTED(in_image.depth() == CV_8U || in_image.depth() == CV_8S,
                                "Posterize: input image data type can not be float, "
diff --git a/tests/ut/cpp/dataset/random_posterize_op_test.cc b/tests/ut/cpp/dataset/random_posterize_op_test.cc
index 8107e282135..1e8ff6ed596 100644
--- a/tests/ut/cpp/dataset/random_posterize_op_test.cc
+++ b/tests/ut/cpp/dataset/random_posterize_op_test.cc
@@ -16,6 +16,8 @@
 #include "common/common.h"
 #include "common/cvop_common.h"
 #include "minddata/dataset/kernels/image/random_posterize_op.h"
+#include "minddata/dataset/include/dataset/execute.h"
+#include "minddata/dataset/include/dataset/vision.h"
 #include "minddata/dataset/core/cv_tensor.h"
 #include "utils/log_adapter.h"
 
@@ -25,8 +27,6 @@ using mindspore::ExceptionType::NoExceptionType;
 using mindspore::MsLogLevel::INFO;
 
 class MindDataTestRandomPosterizeOp : public UT::CVOP::CVOpCommon {
- public:
-  MindDataTestRandomPosterizeOp() : CVOpCommon() {}
 };
 
 TEST_F(MindDataTestRandomPosterizeOp, TestOp1) {
@@ -39,3 +39,32 @@ TEST_F(MindDataTestRandomPosterizeOp, TestOp1) {
   EXPECT_TRUE(s.IsOk());
   CheckImageShapeAndData(output_tensor, kRandomPosterize);
 }
+
+TEST_F(MindDataTestRandomPosterizeOp, TestOp2) {
+   // Test Eager RandomPosterize image = (h, w, c)
+  MS_LOG(INFO) << "Doing VisionRandomPosterizeTest.";
+  std::shared_ptr<Tensor> de_tensor;
+  std::string dataset_root_path = "data/dataset";
+  Tensor::CreateFromFile(dataset_root_path + "/testPK/data/class1/0.jpg", &de_tensor);
+  auto image = mindspore::MSTensor(std::make_shared<DETensor>(de_tensor));
+  std::shared_ptr<TensorTransform>  decode_op = std::make_shared<vision::Decode>();
+  std::shared_ptr<TensorTransform> randomposterize_op(new vision::RandomPosterize({3, 5}));
+  auto transform = Execute({decode_op, randomposterize_op});
+  Status rc = transform(image, &image);
+  EXPECT_TRUE(rc.IsOk());
+  EXPECT_EQ(image.Shape().size(), 3);
+  EXPECT_EQ(image.Shape()[2], 3);
+}
+
+TEST_F(MindDataTestRandomPosterizeOp, TestOp3) {
+  // Test Eager RandomSolarize image.size = {2, 2, 2, 1}
+  MS_LOG(INFO) << "Doing VisionRandomSolarizeTest.";
+
+  std::shared_ptr<Tensor> de_tensor;
+  Tensor::CreateFromVector(std::vector<uint8_t>({0, 25, 120, 0, 38, 2, 10, 13}), TensorShape({2, 2, 2, 1}), &de_tensor);
+  auto image = mindspore::MSTensor(std::make_shared<DETensor>(de_tensor));
+  std::shared_ptr<TensorTransform> randomsolarize_op(new vision::RandomSolarize({12, 25}));
+  auto transform = Execute({randomsolarize_op});
+  Status rc = transform(image, &image);
+  EXPECT_TRUE(rc.IsError());
+}