check affine input size

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

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);

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++) {