From 89f109a3f3eddaccdaa3c92cfb6e07e1ba54e255 Mon Sep 17 00:00:00 2001
From: shenwei41 <shenwei41@huawei.com>
Date: Wed, 17 Mar 2021 20:28:29 +0800
Subject: [PATCH] add check of operation

---
 .../kernels/image/lite_cv/image_process.cc    |  6 ++++
 tests/ut/cpp/dataset/image_process_test.cc    | 32 +++++++++++++++++++
 2 files changed, 38 insertions(+)

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 bb3afa8f39..817e84aff1 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
@@ -1560,6 +1560,9 @@ bool GetPerspectiveTransformImpl(const LiteMat &src1, const LiteMat &src2, LiteM
 }
 
 bool GetPerspectiveTransform(std::vector<Point> src_point, std::vector<Point> dst_point, LiteMat &M) {
+  if (src_point.size() != 4 || dst_point.size() != 4) {
+    return false;
+  }
   double m[8][8];
   double n[8];
   LiteMat src1(8, 8, m, LDataType(LDataType::DOUBLE));
@@ -1645,6 +1648,9 @@ bool GetAffineTransformImpl(LiteMat src, LiteMat dst) {
 }
 
 bool GetAffineTransform(std::vector<Point> src_point, std::vector<Point> dst_point, LiteMat &M) {
+  if (src_point.size() != 3 || dst_point.size() != 3) {
+    return false;
+  }
   double m[6 * 6];
   double n[6];
   LiteMat src1(6, 6, m, LDataType(LDataType::DOUBLE));
diff --git a/tests/ut/cpp/dataset/image_process_test.cc b/tests/ut/cpp/dataset/image_process_test.cc
index dc83ba4521..5e024fdb04 100644
--- a/tests/ut/cpp/dataset/image_process_test.cc
+++ b/tests/ut/cpp/dataset/image_process_test.cc
@@ -1350,6 +1350,22 @@ TEST_F(MindDataImageProcess, testGetPerspectiveTransform) {
   AccuracyComparison(expect_matrix, M);
 }
 
+TEST_F(MindDataImageProcess, testGetPerspectiveTransformFail) {
+  std::vector<Point> src = {Point(165, 270), Point(835, 270), Point(360, 125), Point(615, 125)};
+  std::vector<Point> dst = {Point(100, 100), Point(500, 30)};
+
+  LiteMat M;
+  bool ret = GetPerspectiveTransform(src, dst, M);
+  EXPECT_FALSE(ret);
+
+  std::vector<Point> src1 = {Point(360, 125), Point(615, 125)};
+  std::vector<Point> dst2 = {Point(165, 270), Point(835, 270), Point(100, 100), Point(500, 30)};
+
+  LiteMat M1;
+  bool ret1 = GetPerspectiveTransform(src, dst, M1);
+  EXPECT_FALSE(ret1);
+}
+
 TEST_F(MindDataImageProcess, testGetAffineTransform) {
   std::vector<std::vector<double>> expect_matrix = {{0.400000, 0.066667, 16.666667}, {0.000000, 0.333333, 23.333333}};
 
@@ -1363,6 +1379,22 @@ TEST_F(MindDataImageProcess, testGetAffineTransform) {
   AccuracyComparison(expect_matrix, M);
 }
 
+TEST_F(MindDataImageProcess, testGetAffineTransformFail) {
+  std::vector<Point> src = {Point(50, 50), Point(200, 50)};
+  std::vector<Point> dst = {Point(40, 40), Point(100, 40), Point(50, 90)};
+
+  LiteMat M;
+  bool ret = GetAffineTransform(src, dst, M);
+  EXPECT_FALSE(ret);
+
+  std::vector<Point> src1 = {Point(50, 50), Point(200, 50), Point(50, 200)};
+  std::vector<Point> dst1 = {Point(40, 40), Point(100, 40)};
+
+  LiteMat M1;
+  bool ret1 = GetAffineTransform(src1, dst1, M1);
+  EXPECT_FALSE(ret1);
+}
+
 TEST_F(MindDataImageProcess, TestConv2D8U) {
   LiteMat lite_mat_src;
   lite_mat_src.Init(3, 3, 1, LDataType::UINT8);