From 197f0cabe05f3cadb100cbdb11ac6e4f3be1e0f7 Mon Sep 17 00:00:00 2001
From: jiangzhiwen <jiangzhiwen8@huawei.com>
Date: Tue, 29 Sep 2020 17:20:38 +0800
Subject: [PATCH] check affine input size

---
 .../kernels/image/lite_cv/image_process.cc    | 23 +++++++++++++++----
 .../kernels/image/lite_cv/image_process.h     |  4 ++--
 tests/ut/cpp/dataset/image_process_test.cc    | 11 ++++++++-
 3 files changed, 30 insertions(+), 8 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 4aa9e530cc..977cd993d5 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
@@ -370,6 +370,13 @@ static bool CheckZero(const std::vector<float> &vs) {
   return false;
 }
 
+static bool CheckZero(const std::vector<size_t> &vs) {
+  for (int i = 0; i < vs.size(); i++) {
+    if (vs[i] == 0) return true;
+  }
+  return false;
+}
+
 static bool CheckMeanAndStd(int channel, const std::vector<float> &mean, const std::vector<float> &std) {
   if (mean.size() == 0 && std.size() == 0) {
     return false;
@@ -632,7 +639,11 @@ std::vector<int> ApplyNms(const std::vector<std::vector<float>> &all_boxes, std:
 }
 
 template <typename Pixel_Type>
-void ImplementAffine(LiteMat &src, LiteMat &out_img, double M[6], std::vector<size_t> &dsize, Pixel_Type borderValue) {
+bool ImplementAffine(LiteMat &src, LiteMat &out_img, double M[6], std::vector<size_t> &dsize, Pixel_Type borderValue) {
+  if (dsize.size() != 2 || CheckZero(dsize)) {
+    return false;
+  }
+
   double IM[6];
   for (int i = 0; i < 6; i++) {
     IM[i] = M[i];
@@ -663,14 +674,16 @@ void ImplementAffine(LiteMat &src, LiteMat &out_img, double M[6], std::vector<si
       }
     }
   }
+
+  return true;
 }
 
-void Affine(LiteMat &src, LiteMat &out_img, double M[6], std::vector<size_t> dsize, UINT8_C1 borderValue) {
-  ImplementAffine(src, out_img, M, dsize, borderValue);
+bool Affine(LiteMat &src, LiteMat &out_img, double M[6], std::vector<size_t> dsize, UINT8_C1 borderValue) {
+  return ImplementAffine(src, out_img, M, dsize, borderValue);
 }
 
-void Affine(LiteMat &src, LiteMat &out_img, double M[6], std::vector<size_t> dsize, UINT8_C3 borderValue) {
-  ImplementAffine(src, out_img, M, dsize, borderValue);
+bool Affine(LiteMat &src, LiteMat &out_img, double M[6], std::vector<size_t> dsize, UINT8_C3 borderValue) {
+  return ImplementAffine(src, out_img, M, dsize, borderValue);
 }
 
 }  // namespace dataset
diff --git a/mindspore/ccsrc/minddata/dataset/kernels/image/lite_cv/image_process.h b/mindspore/ccsrc/minddata/dataset/kernels/image/lite_cv/image_process.h
index d03b1b28e2..11bdc76734 100644
--- a/mindspore/ccsrc/minddata/dataset/kernels/image/lite_cv/image_process.h
+++ b/mindspore/ccsrc/minddata/dataset/kernels/image/lite_cv/image_process.h
@@ -71,10 +71,10 @@ bool Pad(const LiteMat &src, LiteMat &dst, int top, int bottom, int left, int ri
          uint8_t fill_b_or_gray, uint8_t fill_g, uint8_t fill_r);
 
 /// \brief Apply affine transformation for 1 channel image
-void Affine(LiteMat &src, LiteMat &out_img, double M[6], std::vector<size_t> dsize, UINT8_C1 borderValue);
+bool Affine(LiteMat &src, LiteMat &out_img, double M[6], std::vector<size_t> dsize, UINT8_C1 borderValue);
 
 /// \brief Apply affine transformation for 3 channel image
-void Affine(LiteMat &src, LiteMat &out_img, double M[6], std::vector<size_t> dsize, UINT8_C3 borderValue);
+bool Affine(LiteMat &src, LiteMat &out_img, double M[6], std::vector<size_t> dsize, UINT8_C3 borderValue);
 
 /// \brief Get default anchor boxes for Faster R-CNN, SSD, YOLO etc
 std::vector<std::vector<float>> GetDefaultBoxes(const BoxesConfig config);
diff --git a/tests/ut/cpp/dataset/image_process_test.cc b/tests/ut/cpp/dataset/image_process_test.cc
index 0f21f5d7e9..cdefaf546e 100644
--- a/tests/ut/cpp/dataset/image_process_test.cc
+++ b/tests/ut/cpp/dataset/image_process_test.cc
@@ -271,6 +271,15 @@ TEST_F(MindDataImageProcess, TestApplyNms) {
   ASSERT_TRUE(keep[2] == 1);
 }
 
+TEST_F(MindDataImageProcess, TestAffineInput) {
+  LiteMat src(3, 3);
+  LiteMat dst;
+  double M[6] = {1};
+  EXPECT_FALSE(Affine(src, dst, M, {}, UINT8_C1(0)));
+  EXPECT_FALSE(Affine(src, dst, M, {3}, UINT8_C1(0)));
+  EXPECT_FALSE(Affine(src, dst, M, {0, 0}, UINT8_C1(0)));
+}
+
 TEST_F(MindDataImageProcess, TestAffine) {
   // The input matrix
   // 0 0 1 0 0
@@ -325,7 +334,7 @@ TEST_F(MindDataImageProcess, TestAffine) {
   }
   std::cout << std::endl;
   LiteMat dst;
-  Affine(src, dst, M, {rows, cols}, UINT8_C1(0));
+  EXPECT_TRUE(Affine(src, dst, M, {rows, cols}, UINT8_C1(0)));
 
   for (size_t i = 0; i < rows; i++) {
     for (size_t j = 0; j < cols; j++) {