[feat] [assistant] [I44TRS] extend invert_op's function

mindspore/ccsrc/minddata/dataset/kernels/image/invert_op.cc

@@ -18,6 +18,7 @@
 #include "minddata/dataset/core/cv_tensor.h"
 #include "minddata/dataset/util/status.h"
 
+
 namespace mindspore {
 namespace dataset {
 
@@ -35,15 +36,13 @@ Status InvertOp::Compute(const std::shared_ptr<Tensor> &input, std::shared_ptr<T
     if (input_cv->Rank() != 3 && input_cv->Rank() != 2) {
       RETURN_STATUS_UNEXPECTED("Invert: image shape is not <H,W,C> or <H,W>");
     }
-    int num_channels;
-    if (input_cv->Rank() == 2) {
-      num_channels = 1;
-    }
-    else{
+    int num_channels = 1;
+    if (input_cv->Rank() == 3) {
       num_channels = input_cv->shape()[2];
     }
     if (num_channels != 3 && num_channels != 1) {
-      RETURN_STATUS_UNEXPECTED("Invert: image shape is incorrect, expected num of channels is 1 or 3, but got:" + std::to_string(num_channels));
+      RETURN_STATUS_UNEXPECTED("Invert: image shape is incorrect, expected num of channels is 1 or 3, but got:" + 
+        std::to_string(num_channels));
     }
     std::shared_ptr<CVTensor> output_cv;
     RETURN_IF_NOT_OK(CVTensor::CreateEmpty(input_cv->shape(), input_cv->type(), &output_cv));
@@ -65,3 +64,4 @@ Status InvertOp::Compute(const std::shared_ptr<Tensor> &input, std::shared_ptr<T
 }
 }  // namespace dataset
 }  // namespace mindspore
+