!2142 Cleanup dataset UT: refactor visualize()

Merge pull request !2142 from Tinazhang/reorg-visualize

tests/ut/python/dataset/test_HWC2CHW.py

@@ -12,12 +12,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
+"""
+Testing HWC2CHW op in DE
+"""
 import numpy as np
 import mindspore.dataset as ds
 import mindspore.dataset.transforms.vision.c_transforms as c_vision
 import mindspore.dataset.transforms.vision.py_transforms as py_vision
 from mindspore import log as logger
-from util import diff_mse, visualize, save_and_check_md5
+from util import diff_mse, visualize_list, save_and_check_md5
 
 GENERATE_GOLDEN = False
 
@@ -55,7 +58,7 @@ def test_HWC2CHW(plot=False):
         mse = diff_mse(transposed_item, original_item.transpose(2, 0, 1))
         assert mse == 0
     if plot:
-        visualize(image, image_transposed)
+        visualize_list(image, image_transposed)
 
 
 def test_HWC2CHW_md5():
@@ -112,7 +115,7 @@ def test_HWC2CHW_comp(plot=False):
         image_c_transposed.append(c_image.transpose(1, 2, 0))
         image_py_transposed.append(py_image.transpose(1, 2, 0))
     if plot:
-        visualize(image_c_transposed, image_py_transposed)
+        visualize_list(image_c_transposed, image_py_transposed, visualize_mode=2)
 
 
 if __name__ == '__main__':

tests/ut/python/dataset/test_autocontrast.py

@@ -12,34 +12,19 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-
-import matplotlib.pyplot as plt
+"""
+Testing AutoContrast op in DE
+"""
 import numpy as np
 
 import mindspore.dataset.engine as de
 import mindspore.dataset.transforms.vision.py_transforms as F
 from mindspore import log as logger
+from util import visualize_list
 
 DATA_DIR = "../data/dataset/testImageNetData/train/"
 
 
-def visualize(image_original, image_auto_contrast):
-    """
-    visualizes the image using DE op and Numpy op
-    """
-    num = len(image_auto_contrast)
-    for i in range(num):
-        plt.subplot(2, num, i + 1)
-        plt.imshow(image_original[i])
-        plt.title("Original image")
-
-        plt.subplot(2, num, i + num + 1)
-        plt.imshow(image_auto_contrast[i])
-        plt.title("DE AutoContrast image")
-
-    plt.show()
-
-
 def test_auto_contrast(plot=False):
     """
     Test AutoContrast
@@ -94,7 +79,7 @@ def test_auto_contrast(plot=False):
     logger.info("MSE= {}".format(str(np.mean(mse))))
 
     if plot:
-        visualize(images_original, images_auto_contrast)
+        visualize_list(images_original, images_auto_contrast)
 
 
 if __name__ == "__main__":

tests/ut/python/dataset/test_center_crop.py

@@ -12,12 +12,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
+"""
+Testing CenterCrop op in DE
+"""
 import numpy as np
 import mindspore.dataset as ds
 import mindspore.dataset.transforms.vision.c_transforms as vision
 import mindspore.dataset.transforms.vision.py_transforms as py_vision
 from mindspore import log as logger
-from util import diff_mse, visualize, save_and_check_md5
+from util import diff_mse, visualize_list, save_and_check_md5
 
 GENERATE_GOLDEN = False
 
@@ -49,7 +52,7 @@ def test_center_crop_op(height=375, width=375, plot=False):
         image_cropped.append(item1["image"].copy())
         image.append(item2["image"].copy())
     if plot:
-        visualize(image, image_cropped)
+        visualize_list(image, image_cropped)
 
 
 def test_center_crop_md5(height=375, width=375):
@@ -93,17 +96,17 @@ def test_center_crop_comp(height=375, width=375, plot=False):
     transform = py_vision.ComposeOp(transforms)
     data2 = data2.map(input_columns=["image"], operations=transform())
 
-    image_cropped = []
-    image = []
+    image_c_cropped = []
+    image_py_cropped = []
     for item1, item2 in zip(data1.create_dict_iterator(), data2.create_dict_iterator()):
         c_image = item1["image"]
         py_image = (item2["image"].transpose(1, 2, 0) * 255).astype(np.uint8)
         # Note: The images aren't exactly the same due to rounding error
         assert diff_mse(py_image, c_image) < 0.001
-        image_cropped.append(c_image.copy())
-        image.append(py_image.copy())
+        image_c_cropped.append(c_image.copy())
+        image_py_cropped.append(py_image.copy())
     if plot:
-        visualize(image, image_cropped)
+        visualize_list(image_c_cropped, image_py_cropped, visualize_mode=2)
 
 
 # pylint: disable=unnecessary-lambda

tests/ut/python/dataset/test_cut_out.py

@@ -15,37 +15,19 @@
 """
 Testing CutOut op in DE
 """
-import matplotlib.pyplot as plt
 import numpy as np
 
 import mindspore.dataset as ds
 import mindspore.dataset.transforms.vision.c_transforms as c
 import mindspore.dataset.transforms.vision.py_transforms as f
 from mindspore import log as logger
+from util import visualize_image, diff_mse
 
 DATA_DIR = ["../data/dataset/test_tf_file_3_images/train-0000-of-0001.data"]
 SCHEMA_DIR = "../data/dataset/test_tf_file_3_images/datasetSchema.json"
 
 
-def visualize(image_1, image_2):
-    """
-    visualizes the image using RandomErasing and Cutout
-    """
-    plt.subplot(141)
-    plt.imshow(image_1)
-    plt.title("RandomErasing")
-
-    plt.subplot(142)
-    plt.imshow(image_2)
-    plt.title("Cutout")
-
-    plt.subplot(143)
-    plt.imshow(image_1 - image_2)
-    plt.title("Difference image")
-    plt.show()
-
-
-def test_cut_out_op():
+def test_cut_out_op(plot=False):
     """
     Test Cutout
     """
@@ -87,7 +69,9 @@ def test_cut_out_op():
         logger.info("dtype of image_1: {}".format(image_1.dtype))
         logger.info("dtype of image_2: {}".format(image_2.dtype))
 
-        # visualize(image_1, image_2)
+        mse = diff_mse(image_1, image_2)
+        if plot:
+            visualize_image(image_1, image_2, mse)
 
 
 def test_cut_out_op_multicut():
@@ -134,5 +118,5 @@ def test_cut_out_op_multicut():
 
 
 if __name__ == "__main__":
-    test_cut_out_op()
+    test_cut_out_op(plot=True)
     test_cut_out_op_multicut()

tests/ut/python/dataset/test_decode.py

@@ -15,12 +15,12 @@
 """
 Testing Decode op in DE
 """
-import numpy as np
 import cv2
 
 import mindspore.dataset as ds
 import mindspore.dataset.transforms.vision.c_transforms as vision
 from mindspore import log as logger
+from util import diff_mse
 
 DATA_DIR = ["../data/dataset/test_tf_file_3_images/train-0000-of-0001.data"]
 SCHEMA_DIR = "../data/dataset/test_tf_file_3_images/datasetSchema.json"
@@ -45,8 +45,7 @@ def test_decode_op():
         expected = cv2.imdecode(item2["image"], cv2.IMREAD_COLOR)
         expected = cv2.cvtColor(expected, cv2.COLOR_BGR2RGB)
         assert actual.shape == expected.shape
-        diff = actual - expected
-        mse = np.sum(np.power(diff, 2))
+        mse = diff_mse(actual, expected)
         assert mse == 0
 
 
@@ -71,8 +70,7 @@ def test_decode_op_tf_file_dataset():
         expected = cv2.imdecode(item2["image"], cv2.IMREAD_COLOR)
         expected = cv2.cvtColor(expected, cv2.COLOR_BGR2RGB)
         assert actual.shape == expected.shape
-        diff = actual - expected
-        mse = np.sum(np.power(diff, 2))
+        mse = diff_mse(actual, expected)
         assert mse == 0
 
 

tests/ut/python/dataset/test_equalize.py

@@ -12,34 +12,19 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-
-import matplotlib.pyplot as plt
+"""
+Testing Equalize op in DE
+"""
 import numpy as np
 
 import mindspore.dataset.engine as de
 import mindspore.dataset.transforms.vision.py_transforms as F
 from mindspore import log as logger
+from util import visualize_list
 
 DATA_DIR = "../data/dataset/testImageNetData/train/"
 
 
-def visualize(image_original, image_equalize):
-    """
-    visualizes the image using DE op and Numpy op
-    """
-    num = len(image_equalize)
-    for i in range(num):
-        plt.subplot(2, num, i + 1)
-        plt.imshow(image_original[i])
-        plt.title("Original image")
-
-        plt.subplot(2, num, i + num + 1)
-        plt.imshow(image_equalize[i])
-        plt.title("DE Color Equalized image")
-
-    plt.show()
-
-
 def test_equalize(plot=False):
     """
     Test Equalize
@@ -94,7 +79,7 @@ def test_equalize(plot=False):
     logger.info("MSE= {}".format(str(np.mean(mse))))
 
     if plot:
-        visualize(images_original, images_equalize)
+        visualize_list(images_original, images_equalize)
 
 
 if __name__ == "__main__":

tests/ut/python/dataset/test_five_crop.py

@@ -20,7 +20,7 @@ import numpy as np
 import mindspore.dataset as ds
 import mindspore.dataset.transforms.vision.py_transforms as vision
 from mindspore import log as logger
-from util import visualize
+from util import visualize_list
 
 DATA_DIR = ["../data/dataset/test_tf_file_3_images/train-0000-of-0001.data"]
 SCHEMA_DIR = "../data/dataset/test_tf_file_3_images/datasetSchema.json"
@@ -63,7 +63,7 @@ def test_five_crop_op(plot=False):
         logger.info("dtype of image_1: {}".format(image_1.dtype))
         logger.info("dtype of image_2: {}".format(image_2.dtype))
         if plot:
-            visualize(np.array([image_1]*10), (image_2 * 255).astype(np.uint8).transpose(0, 2, 3, 1))
+            visualize_list(np.array([image_1]*10), (image_2 * 255).astype(np.uint8).transpose(0, 2, 3, 1))
 
         # The output data should be of a 4D tensor shape, a stack of 5 images.
         assert len(image_2.shape) == 4

tests/ut/python/dataset/test_invert.py

@@ -12,34 +12,19 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-
-import matplotlib.pyplot as plt
+"""
+Testing Invert op in DE
+"""
 import numpy as np
 
 import mindspore.dataset.engine as de
 import mindspore.dataset.transforms.vision.py_transforms as F
 from mindspore import log as logger
+from util import visualize_list
 
 DATA_DIR = "../data/dataset/testImageNetData/train/"
 
 
-def visualize(image_original, image_invert):
-    """
-    visualizes the image using DE op and Numpy op
-    """
-    num = len(image_invert)
-    for i in range(num):
-        plt.subplot(2, num, i + 1)
-        plt.imshow(image_original[i])
-        plt.title("Original image")
-
-        plt.subplot(2, num, i + num + 1)
-        plt.imshow(image_invert[i])
-        plt.title("DE Color Inverted image")
-
-    plt.show()
-
-
 def test_invert(plot=False):
     """
     Test Invert
@@ -94,7 +79,7 @@ def test_invert(plot=False):
     logger.info("MSE= {}".format(str(np.mean(mse))))
 
     if plot:
-        visualize(images_original, images_invert)
+        visualize_list(images_original, images_invert)
 
 
 if __name__ == "__main__":

tests/ut/python/dataset/test_linear_transformation.py

@@ -19,7 +19,7 @@ import numpy as np
 import mindspore.dataset as ds
 import mindspore.dataset.transforms.vision.py_transforms as py_vision
 from mindspore import log as logger
-from util import diff_mse, visualize, save_and_check_md5
+from util import diff_mse, visualize_list, save_and_check_md5
 
 GENERATE_GOLDEN = False
 
@@ -71,7 +71,7 @@ def test_linear_transformation_op(plot=False):
         mse = diff_mse(image1, image2)
         assert mse == 0
     if plot:
-        visualize(image, image_transformed)
+        visualize_list(image, image_transformed)
 
 def test_linear_transformation_md5_01():
     """

tests/ut/python/dataset/test_normalizeOp.py

@@ -16,12 +16,11 @@
 Testing Normalize op in DE
 """
 import numpy as np
-import matplotlib.pyplot as plt
 import mindspore.dataset as ds
 import mindspore.dataset.transforms.vision.c_transforms as c_vision
 import mindspore.dataset.transforms.vision.py_transforms as py_vision
 from mindspore import log as logger
-from util import diff_mse, save_and_check_md5
+from util import diff_mse, save_and_check_md5, visualize_image
 
 DATA_DIR = ["../data/dataset/test_tf_file_3_images/train-0000-of-0001.data"]
 SCHEMA_DIR = "../data/dataset/test_tf_file_3_images/datasetSchema.json"
@@ -29,28 +28,6 @@ SCHEMA_DIR = "../data/dataset/test_tf_file_3_images/datasetSchema.json"
 GENERATE_GOLDEN = False
 
 
-def visualize_mse(image_de_normalized, image_np_normalized, mse, image_original):
-    """
-    visualizes the image using DE op and Numpy op
-    """
-    plt.subplot(141)
-    plt.imshow(image_original)
-    plt.title("Original image")
-
-    plt.subplot(142)
-    plt.imshow(image_de_normalized)
-    plt.title("DE normalized image")
-
-    plt.subplot(143)
-    plt.imshow(image_np_normalized)
-    plt.title("Numpy normalized image")
-
-    plt.subplot(144)
-    plt.imshow(image_de_normalized - image_np_normalized)
-    plt.title("Difference image, mse : {}".format(mse))
-    plt.show()
-
-
 def normalize_np(image, mean, std):
     """
     Apply the normalization
@@ -137,7 +114,7 @@ def test_normalize_op_c(plot=False):
         logger.info("image_{}, mse: {}".format(num_iter + 1, mse))
         assert mse < 0.01
         if plot:
-            visualize_mse(image_de_normalized, image_np_normalized, mse, image_original)
+            visualize_image(image_original, image_de_normalized, mse, image_np_normalized)
         num_iter += 1
 
 
@@ -174,7 +151,7 @@ def test_normalize_op_py(plot=False):
         logger.info("image_{}, mse: {}".format(num_iter + 1, mse))
         assert mse < 0.01
         if plot:
-            visualize_mse(image_de_normalized, image_np_normalized, mse, image_original)
+            visualize_image(image_original, image_de_normalized, mse, image_np_normalized)
         num_iter += 1
 
 

tests/ut/python/dataset/test_onehot_op.py

@@ -20,6 +20,7 @@ import numpy as np
 import mindspore.dataset as ds
 import mindspore.dataset.transforms.c_transforms as data_trans
 from mindspore import log as logger
+from util import diff_mse
 
 DATA_DIR = ["../data/dataset/test_tf_file_3_images/train-0000-of-0001.data"]
 SCHEMA_DIR = "../data/dataset/test_tf_file_3_images/datasetSchema.json"
@@ -55,7 +56,7 @@ def test_one_hot():
         assert len(item1) == len(item2)
         label1 = item1["label"]
         label2 = one_hot(item2["label"][0], depth)
-        mse = np.sum(label1 - label2)
+        mse = diff_mse(label1, label2)
         logger.info("DE one_hot: {}, Numpy one_hot: {}, diff: {}".format(label1, label2, mse))
         num_iter += 1
 

tests/ut/python/dataset/test_random_affine.py

@@ -19,7 +19,7 @@ import numpy as np
 import mindspore.dataset as ds
 import mindspore.dataset.transforms.vision.py_transforms as py_vision
 from mindspore import log as logger
-from util import visualize, save_and_check_md5, \
+from util import visualize_list, save_and_check_md5, \
     config_get_set_seed, config_get_set_num_parallel_workers
 
 GENERATE_GOLDEN = False
@@ -62,7 +62,7 @@ def test_random_affine_op(plot=False):
         image_affine.append(image1)
         image_original.append(image2)
     if plot:
-        visualize(image_original, image_affine)
+        visualize_list(image_original, image_affine)
 
 
 def test_random_affine_md5():

tests/ut/python/dataset/test_random_apply.py

@@ -19,7 +19,7 @@ import numpy as np
 import mindspore.dataset as ds
 import mindspore.dataset.transforms.vision.py_transforms as py_vision
 from mindspore import log as logger
-from util import visualize, config_get_set_seed, \
+from util import visualize_list, config_get_set_seed, \
     config_get_set_num_parallel_workers, save_and_check_md5
 
 GENERATE_GOLDEN = False
@@ -63,7 +63,7 @@ def test_random_apply_op(plot=False):
         image_apply.append(image1)
         image_original.append(image2)
     if plot:
-        visualize(image_original, image_apply)
+        visualize_list(image_original, image_apply)
 
 
 def test_random_apply_md5():

tests/ut/python/dataset/test_random_choice.py

@@ -19,7 +19,7 @@ import numpy as np
 import mindspore.dataset as ds
 import mindspore.dataset.transforms.vision.py_transforms as py_vision
 from mindspore import log as logger
-from util import visualize, diff_mse
+from util import visualize_list, diff_mse
 
 DATA_DIR = ["../data/dataset/test_tf_file_3_images/train-0000-of-0001.data"]
 SCHEMA_DIR = "../data/dataset/test_tf_file_3_images/datasetSchema.json"
@@ -60,7 +60,7 @@ def test_random_choice_op(plot=False):
         image_choice.append(image1)
         image_original.append(image2)
     if plot:
-        visualize(image_original, image_choice)
+        visualize_list(image_original, image_choice)
 
 
 def test_random_choice_comp(plot=False):
@@ -102,7 +102,7 @@ def test_random_choice_comp(plot=False):
         mse = diff_mse(image1, image2)
         assert mse == 0
     if plot:
-        visualize(image_original, image_choice)
+        visualize_list(image_original, image_choice)
 
 
 def test_random_choice_exception_random_crop_badinput():

tests/ut/python/dataset/test_random_color.py

@@ -12,34 +12,19 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-
-import matplotlib.pyplot as plt
+"""
+Testing RandomColor op in DE
+"""
 import numpy as np
 
 import mindspore.dataset.engine as de
 import mindspore.dataset.transforms.vision.py_transforms as F
 from mindspore import log as logger
+from util import visualize_list
 
 DATA_DIR = "../data/dataset/testImageNetData/train/"
 
 
-def visualize(image_original, image_random_color):
-    """
-    visualizes the image using DE op and Numpy op
-    """
-    num = len(image_random_color)
-    for i in range(num):
-        plt.subplot(2, num, i + 1)
-        plt.imshow(image_original[i])
-        plt.title("Original image")
-
-        plt.subplot(2, num, i + num + 1)
-        plt.imshow(image_random_color[i])
-        plt.title("DE Random Color image")
-
-    plt.show()
-
-
 def test_random_color(degrees=(0.1, 1.9), plot=False):
     """
     Test RandomColor
@@ -94,7 +79,7 @@ def test_random_color(degrees=(0.1, 1.9), plot=False):
     logger.info("MSE= {}".format(str(np.mean(mse))))
 
     if plot:
-        visualize(images_original, images_random_color)
+        visualize_list(images_original, images_random_color)
 
 
 if __name__ == "__main__":

tests/ut/python/dataset/test_random_color_adjust.py

@@ -16,37 +16,18 @@
 Testing RandomColorAdjust op in DE
 """
 import pytest
-import matplotlib.pyplot as plt
 import numpy as np
-from util import diff_mse
 
 import mindspore.dataset as ds
 import mindspore.dataset.transforms.vision.c_transforms as c_vision
 import mindspore.dataset.transforms.vision.py_transforms as py_vision
 from mindspore import log as logger
+from util import diff_mse, visualize_image
 
 DATA_DIR = ["../data/dataset/test_tf_file_3_images/train-0000-of-0001.data"]
 SCHEMA_DIR = "../data/dataset/test_tf_file_3_images/datasetSchema.json"
 
 
-def visualize(first, mse, second):
-    """
-    visualizes the image using DE op and OpenCV
-    """
-    plt.subplot(141)
-    plt.imshow(first)
-    plt.title("c transformed image")
-
-    plt.subplot(142)
-    plt.imshow(second)
-    plt.title("py random_color_adjust image")
-
-    plt.subplot(143)
-    plt.imshow(first - second)
-    plt.title("Difference image, mse : {}".format(mse))
-    plt.show()
-
-
 def util_test_random_color_adjust_error(brightness=(1, 1), contrast=(1, 1), saturation=(1, 1), hue=(0, 0)):
     """
     Util function that tests the error message in case of grayscale images
@@ -126,7 +107,7 @@ def util_test_random_color_adjust_op(brightness=(1, 1), contrast=(1, 1), saturat
         assert mse < 0.01
 
         if plot:
-            visualize(c_image, mse, py_image)
+            visualize_image(c_image, py_image, mse)
 
 
 def test_random_color_adjust_op_brightness(plot=False):

tests/ut/python/dataset/test_random_crop.py

@@ -21,7 +21,7 @@ import mindspore.dataset.transforms.vision.py_transforms as py_vision
 import mindspore.dataset.transforms.vision.utils as mode
 import mindspore.dataset as ds
 from mindspore import log as logger
-from util import save_and_check_md5, visualize, config_get_set_seed, \
+from util import save_and_check_md5, visualize_list, config_get_set_seed, \
     config_get_set_num_parallel_workers
 
 
@@ -57,7 +57,7 @@ def test_random_crop_op_c(plot=False):
         image_cropped.append(image1)
         image.append(image2)
     if plot:
-        visualize(image, image_cropped)
+        visualize_list(image, image_cropped)
 
 def test_random_crop_op_py(plot=False):
     """
@@ -91,7 +91,7 @@ def test_random_crop_op_py(plot=False):
         crop_images.append(crop)
         original_images.append(original)
     if plot:
-        visualize(original_images, crop_images)
+        visualize_list(original_images, crop_images)
 
 def test_random_crop_01_c():
     """
@@ -534,7 +534,7 @@ def test_random_crop_comp(plot=False):
         image_c_cropped.append(c_image)
         image_py_cropped.append(py_image)
     if plot:
-        visualize(image_c_cropped, image_py_cropped)
+        visualize_list(image_c_cropped, image_py_cropped, visualize_mode=2)
 
 
 if __name__ == "__main__":

tests/ut/python/dataset/test_random_crop_and_resize.py

@@ -23,7 +23,7 @@ import mindspore.dataset.transforms.vision.py_transforms as py_vision
 import mindspore.dataset.transforms.vision.utils as mode
 import mindspore.dataset as ds
 from mindspore import log as logger
-from util import diff_mse, save_and_check_md5, visualize, \
+from util import diff_mse, save_and_check_md5, visualize_list, \
     config_get_set_seed, config_get_set_num_parallel_workers
 
 DATA_DIR = ["../data/dataset/test_tf_file_3_images/train-0000-of-0001.data"]
@@ -64,7 +64,7 @@ def test_random_crop_and_resize_op_c(plot=False):
         crop_and_resize_images.append(crop_and_resize)
         original_images.append(original)
     if plot:
-        visualize(original_images, crop_and_resize_images)
+        visualize_list(original_images, crop_and_resize_images)
 
 
 def test_random_crop_and_resize_op_py(plot=False):
@@ -106,7 +106,7 @@ def test_random_crop_and_resize_op_py(plot=False):
         crop_and_resize_images.append(crop_and_resize)
         original_images.append(original)
     if plot:
-        visualize(original_images, crop_and_resize_images)
+        visualize_list(original_images, crop_and_resize_images)
 
 
 def test_random_crop_and_resize_01():
@@ -332,7 +332,7 @@ def test_random_crop_and_resize_comp(plot=False):
         image_c_cropped.append(c_image)
         image_py_cropped.append(py_image)
     if plot:
-        visualize(image_c_cropped, image_py_cropped)
+        visualize_list(image_c_cropped, image_py_cropped, visualize_mode=2)
 
 
 if __name__ == "__main__":

tests/ut/python/dataset/test_random_crop_decode_resize.py

@@ -15,37 +15,18 @@
 """
 Testing RandomCropDecodeResize op in DE
 """
-import matplotlib.pyplot as plt
-import numpy as np
 import cv2
 
 import mindspore.dataset as ds
 import mindspore.dataset.transforms.vision.c_transforms as vision
 from mindspore import log as logger
+from util import diff_mse, visualize_image
 
 DATA_DIR = ["../data/dataset/test_tf_file_3_images/train-0000-of-0001.data"]
 SCHEMA_DIR = "../data/dataset/test_tf_file_3_images/datasetSchema.json"
 
 
-def visualize(a, mse, original):
-    """
-    visualizes the image using DE op and Numpy Op
-    """
-    plt.subplot(141)
-    plt.imshow(original)
-    plt.title("Original image")
-
-    plt.subplot(142)
-    plt.imshow(a)
-    plt.title("DE random_crop_decode_resize image")
-
-    plt.subplot(143)
-    plt.imshow(a - original)
-    plt.title("Difference image, mse : {}".format(mse))
-    plt.show()
-
-
-def test_random_crop_decode_resize_op():
+def test_random_crop_decode_resize_op(plot=False):
     """
     Test RandomCropDecodeResize op
     """
@@ -66,16 +47,15 @@ def test_random_crop_decode_resize_op():
 
         if num_iter > 0:
             break
-        crop_and_resize = item1["image"]
+        crop_and_resize_de = item1["image"]
         original = item2["image"]
-        original = cv2.resize(original, (512, 256))
-        diff = crop_and_resize - original
-        mse = np.sum(np.power(diff, 2))
+        crop_and_resize_cv = cv2.resize(original, (512, 256))
+        mse = diff_mse(crop_and_resize_de, crop_and_resize_cv)
         logger.info("random_crop_decode_resize_op_{}, mse: {}".format(num_iter + 1, mse))
-        # Uncomment below line if you want to visualize images
-        # visualize(crop_and_resize, mse, original)
+        if plot:
+            visualize_image(original, crop_and_resize_de, mse, crop_and_resize_cv)
         num_iter += 1
 
 
 if __name__ == "__main__":
-    test_random_crop_decode_resize_op()
+    test_random_crop_decode_resize_op(plot=True)

tests/ut/python/dataset/test_random_erasing.py

@@ -15,36 +15,18 @@
 """
 Testing RandomErasing op in DE
 """
-import matplotlib.pyplot as plt
 import numpy as np
 
 import mindspore.dataset as ds
 import mindspore.dataset.transforms.vision.py_transforms as vision
 from mindspore import log as logger
+from util import diff_mse, visualize_image
 
 DATA_DIR = ["../data/dataset/test_tf_file_3_images/train-0000-of-0001.data"]
 SCHEMA_DIR = "../data/dataset/test_tf_file_3_images/datasetSchema.json"
 
 
-def visualize(image_1, image_2):
-    """
-    visualizes the image using RandomErasing and Cutout
-    """
-    plt.subplot(141)
-    plt.imshow(image_1)
-    plt.title("RandomErasing")
-
-    plt.subplot(142)
-    plt.imshow(image_2)
-    plt.title("Cutout")
-
-    plt.subplot(143)
-    plt.imshow(image_1 - image_2)
-    plt.title("Difference image")
-    plt.show()
-
-
-def test_random_erasing_op():
+def test_random_erasing_op(plot=False):
     """
     Test RandomErasing and Cutout
     """
@@ -82,8 +64,10 @@ def test_random_erasing_op():
         logger.info("dtype of image_1: {}".format(image_1.dtype))
         logger.info("dtype of image_2: {}".format(image_2.dtype))
 
-        # visualize(image_1, image_2)
+        mse = diff_mse(image_1, image_2)
+        if plot:
+            visualize_image(image_1, image_2, mse)
 
 
 if __name__ == "__main__":
-    test_random_erasing_op()
+    test_random_erasing_op(plot=True)

tests/ut/python/dataset/test_random_grayscale.py

@@ -19,7 +19,7 @@ import numpy as np
 import mindspore.dataset.transforms.vision.py_transforms as py_vision
 import mindspore.dataset as ds
 from mindspore import log as logger
-from util import save_and_check_md5, visualize, \
+from util import save_and_check_md5, visualize_list, \
     config_get_set_seed, config_get_set_num_parallel_workers
 
 GENERATE_GOLDEN = False
@@ -61,7 +61,7 @@ def test_random_grayscale_valid_prob(plot=False):
         image_gray.append(image1)
         image.append(image2)
     if plot:
-        visualize(image, image_gray)
+        visualize_list(image, image_gray)
 
 def test_random_grayscale_input_grayscale_images():
     """

tests/ut/python/dataset/test_random_horizontal_flip.py

@@ -15,13 +15,12 @@
 """
 Testing the random horizontal flip op in DE
 """
-import matplotlib.pyplot as plt
 import numpy as np
 import mindspore.dataset as ds
 import mindspore.dataset.transforms.vision.c_transforms as c_vision
 import mindspore.dataset.transforms.vision.py_transforms as py_vision
 from mindspore import log as logger
-from util import save_and_check_md5, visualize, diff_mse, \
+from util import save_and_check_md5, visualize_list, visualize_image, diff_mse, \
     config_get_set_seed, config_get_set_num_parallel_workers
 
 GENERATE_GOLDEN = False
@@ -41,29 +40,7 @@ def h_flip(image):
     return image
 
 
-def visualize_mse(image_de_random_horizontal, image_pil_random_horizontal, mse, image_original):
-    """
-    visualizes the image using DE op and Numpy op
-    """
-    plt.subplot(141)
-    plt.imshow(image_original)
-    plt.title("Original image")
-
-    plt.subplot(142)
-    plt.imshow(image_de_random_horizontal)
-    plt.title("DE random_horizontal image")
-
-    plt.subplot(143)
-    plt.imshow(image_pil_random_horizontal)
-    plt.title("Horizontally flipped image")
-
-    plt.subplot(144)
-    plt.imshow(image_de_random_horizontal - image_pil_random_horizontal)
-    plt.title("Difference image, mse : {}".format(mse))
-    plt.show()
-
-
-def test_random_horizontal_op():
+def test_random_horizontal_op(plot=False):
     """
     Test RandomHorizontalFlip op
     """
@@ -93,9 +70,10 @@ def test_random_horizontal_op():
 
         mse = diff_mse(image_h_flipped, image_h_flipped_2)
         logger.info("image_{}, mse: {}".format(num_iter + 1, mse))
-        # Uncomment below line if you want to visualize images
-        # visualize_mse(image_h_flipped, image_h_flipped_2, mse, image)
         num_iter += 1
+        if plot:
+            visualize_image(image, image_h_flipped, mse, image_h_flipped_2)
+
 
 def test_random_horizontal_valid_prob_c():
     """
@@ -119,6 +97,7 @@ def test_random_horizontal_valid_prob_c():
     ds.config.set_seed(original_seed)
     ds.config.set_num_parallel_workers(original_num_parallel_workers)
 
+
 def test_random_horizontal_valid_prob_py():
     """
     Test RandomHorizontalFlip op with py_transforms: valid non-default input, expect to pass
@@ -144,6 +123,7 @@ def test_random_horizontal_valid_prob_py():
     ds.config.set_seed(original_seed)
     ds.config.set_num_parallel_workers(original_num_parallel_workers)
 
+
 def test_random_horizontal_invalid_prob_c():
     """
     Test RandomHorizontalFlip op in c_transforms: invalid input, expect to raise error
@@ -162,6 +142,7 @@ def test_random_horizontal_invalid_prob_c():
         logger.info("Got an exception in DE: {}".format(str(e)))
         assert "Input is not" in str(e)
 
+
 def test_random_horizontal_invalid_prob_py():
     """
     Test RandomHorizontalFlip op in py_transforms: invalid input, expect to raise error
@@ -184,6 +165,7 @@ def test_random_horizontal_invalid_prob_py():
         logger.info("Got an exception in DE: {}".format(str(e)))
         assert "Input is not" in str(e)
 
+
 def test_random_horizontal_comp(plot=False):
     """
     Test test_random_horizontal_flip and compare between python and c image augmentation ops
@@ -220,12 +202,13 @@ def test_random_horizontal_comp(plot=False):
         mse = diff_mse(image_c, image_py)
         assert mse < 0.001
     if plot:
-        visualize(images_list_c, images_list_py)
+        visualize_list(images_list_c, images_list_py, visualize_mode=2)
+
 
 if __name__ == "__main__":
-    test_random_horizontal_op()
+    test_random_horizontal_op(plot=True)
     test_random_horizontal_valid_prob_c()
     test_random_horizontal_valid_prob_py()
     test_random_horizontal_invalid_prob_c()
     test_random_horizontal_invalid_prob_py()
-    test_random_horizontal_comp(True)
+    test_random_horizontal_comp(plot=True)

tests/ut/python/dataset/test_random_order.py

@@ -19,7 +19,7 @@ import numpy as np
 import mindspore.dataset as ds
 import mindspore.dataset.transforms.vision.py_transforms as py_vision
 from mindspore import log as logger
-from util import visualize, config_get_set_seed, \
+from util import visualize_list, config_get_set_seed, \
     config_get_set_num_parallel_workers, save_and_check_md5
 
 GENERATE_GOLDEN = False
@@ -63,7 +63,7 @@ def test_random_order_op(plot=False):
         image_order.append(image1)
         image_original.append(image2)
     if plot:
-        visualize(image_original, image_order)
+        visualize_list(image_original, image_order)
 
 
 def test_random_order_md5():

tests/ut/python/dataset/test_random_perspective.py

@@ -20,7 +20,7 @@ import mindspore.dataset as ds
 import mindspore.dataset.transforms.vision.py_transforms as py_vision
 from mindspore.dataset.transforms.vision.utils import Inter
 from mindspore import log as logger
-from util import visualize, save_and_check_md5, \
+from util import visualize_list, save_and_check_md5, \
     config_get_set_seed, config_get_set_num_parallel_workers
 
 
@@ -64,7 +64,7 @@ def test_random_perspective_op(plot=False):
         image_perspective.append(image1)
         image_original.append(image2)
     if plot:
-        visualize(image_original, image_perspective)
+        visualize_list(image_original, image_perspective)
 
 
 def skip_test_random_perspective_md5():

tests/ut/python/dataset/test_random_resize.py

@@ -15,35 +15,16 @@
 """
 Testing the resize op in DE
 """
-import matplotlib.pyplot as plt
-
 import mindspore.dataset as ds
 import mindspore.dataset.transforms.vision.c_transforms as vision
 from mindspore import log as logger
+from util import visualize_list
 
 DATA_DIR = ["../data/dataset/test_tf_file_3_images/train-0000-of-0001.data"]
 SCHEMA_DIR = "../data/dataset/test_tf_file_3_images/datasetSchema.json"
 
 
-def visualize(image_de_resized, image_np_resized, mse):
-    """
-    visualizes the image using DE op and Numpy op
-    """
-    plt.subplot(131)
-    plt.imshow(image_de_resized)
-    plt.title("DE resize image")
-
-    plt.subplot(132)
-    plt.imshow(image_np_resized)
-    plt.title("Numpy resized image")
-
-    plt.subplot(133)
-    plt.imshow(image_de_resized - image_np_resized)
-    plt.title("Difference image, mse : {}".format(mse))
-    plt.show()
-
-
-def test_random_resize_op():
+def test_random_resize_op(plot=False):
     """
     Test random_resize_op
     """
@@ -56,15 +37,20 @@ def test_random_resize_op():
 
     # apply map operations on images
     data1 = data1.map(input_columns=["image"], operations=decode_op)
-    data1 = data1.map(input_columns=["image"], operations=resize_op)
 
+    data2 = data1.map(input_columns=["image"], operations=resize_op)
+    image_original = []
+    image_resized = []
     num_iter = 0
-    for item in data1.create_dict_iterator():
-        _ = item["image"]
-        # Uncomment below line if you want to visualize images
-        # visualize(image_de_resized, image_np_resized, mse)
+    for item1, item2 in zip(data1.create_dict_iterator(), data2.create_dict_iterator()):
+        image_1 = item1["image"]
+        image_2 = item2["image"]
+        image_original.append(image_1)
+        image_resized.append(image_2)
         num_iter += 1
+    if plot:
+        visualize_list(image_original, image_resized)
 
 
 if __name__ == "__main__":
-    test_random_resize_op()
+    test_random_resize_op(plot=True)

tests/ut/python/dataset/test_random_rotation.py

@@ -15,7 +15,6 @@
 """
 Testing RandomRotation op in DE
 """
-import matplotlib.pyplot as plt
 import numpy as np
 import cv2
 
@@ -23,35 +22,13 @@ import mindspore.dataset as ds
 import mindspore.dataset.transforms.vision.c_transforms as c_vision
 import mindspore.dataset.transforms.vision.py_transforms as py_vision
 from mindspore import log as logger
+from util import visualize_image, diff_mse
 
 DATA_DIR = ["../data/dataset/test_tf_file_3_images/train-0000-of-0001.data"]
 SCHEMA_DIR = "../data/dataset/test_tf_file_3_images/datasetSchema.json"
 
 
-def visualize(a, mse, original):
-    """
-    visualizes the image using DE op and enCV
-    """
-    plt.subplot(141)
-    plt.imshow(original)
-    plt.title("Original image")
-
-    plt.subplot(142)
-    plt.imshow(a)
-    plt.title("DE random_crop_and_resize image")
-
-    plt.subplot(143)
-    plt.imshow(a - original)
-    plt.title("Difference image, mse : {}".format(mse))
-    plt.show()
-
-
-def diff_mse(in1, in2):
-    mse = (np.square(in1.astype(float) / 255 - in2.astype(float) / 255)).mean()
-    return mse * 100
-
-
-def test_random_rotation_op():
+def test_random_rotation_op(plot=False):
     """
     Test RandomRotation op
     """
@@ -73,17 +50,16 @@ def test_random_rotation_op():
     for item1, item2 in zip(data1.create_dict_iterator(), data2.create_dict_iterator()):
         if num_iter > 0:
             break
-        rotation = item1["image"]
+        rotation_de = item1["image"]
         original = item2["image"]
         logger.info("shape before rotate: {}".format(original.shape))
-        original = cv2.rotate(original, cv2.ROTATE_90_COUNTERCLOCKWISE)
-        diff = rotation - original
-        mse = np.sum(np.power(diff, 2))
+        rotation_cv = cv2.rotate(original, cv2.ROTATE_90_COUNTERCLOCKWISE)
+        mse = diff_mse(rotation_de, rotation_cv)
         logger.info("random_rotation_op_{}, mse: {}".format(num_iter + 1, mse))
         assert mse == 0
-        # Uncomment below line if you want to visualize images
-        # visualize(rotation, mse, original)
         num_iter += 1
+    if plot:
+        visualize_image(original, rotation_de, mse, rotation_cv)
 
 
 def test_random_rotation_expand():
@@ -148,6 +124,6 @@ def test_rotation_diff():
 
 
 if __name__ == "__main__":
-    test_random_rotation_op()
+    test_random_rotation_op(True)
     test_random_rotation_expand()
     test_rotation_diff()

tests/ut/python/dataset/test_random_sharpness.py

@@ -12,34 +12,19 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-
-import matplotlib.pyplot as plt
+"""
+Testing RandomSharpness op in DE
+"""
 import numpy as np
 
 import mindspore.dataset.engine as de
 import mindspore.dataset.transforms.vision.py_transforms as F
 from mindspore import log as logger
+from util import visualize_list
 
 DATA_DIR = "../data/dataset/testImageNetData/train/"
 
 
-def visualize(image_original, image_random_sharpness):
-    """
-    visualizes the image using DE op and Numpy op
-    """
-    num = len(image_random_sharpness)
-    for i in range(num):
-        plt.subplot(2, num, i + 1)
-        plt.imshow(image_original[i])
-        plt.title("Original image")
-
-        plt.subplot(2, num, i + num + 1)
-        plt.imshow(image_random_sharpness[i])
-        plt.title("DE Random Sharpness image")
-
-    plt.show()
-
-
 def test_random_sharpness(degrees=(0.1, 1.9), plot=False):
     """
     Test RandomSharpness
@@ -94,7 +79,7 @@ def test_random_sharpness(degrees=(0.1, 1.9), plot=False):
     logger.info("MSE= {}".format(str(np.mean(mse))))
 
     if plot:
-        visualize(images_original, images_random_sharpness)
+        visualize_list(images_original, images_random_sharpness)
 
 
 if __name__ == "__main__":

tests/ut/python/dataset/test_random_vertical_flip.py

@@ -15,14 +15,12 @@
 """
 Testing the random vertical flip op in DE
 """
-import matplotlib.pyplot as plt
 import numpy as np
-
 import mindspore.dataset as ds
 import mindspore.dataset.transforms.vision.c_transforms as c_vision
 import mindspore.dataset.transforms.vision.py_transforms as py_vision
 from mindspore import log as logger
-from util import save_and_check_md5, visualize, diff_mse, \
+from util import save_and_check_md5, visualize_list, visualize_image, diff_mse, \
     config_get_set_seed, config_get_set_num_parallel_workers
 
 GENERATE_GOLDEN = False
@@ -42,29 +40,7 @@ def v_flip(image):
     return image
 
 
-def visualize_with_mse(image_de_random_vertical, image_pil_random_vertical, mse, image_original):
-    """
-    visualizes the image using DE op and Numpy op
-    """
-    plt.subplot(141)
-    plt.imshow(image_original)
-    plt.title("Original image")
-
-    plt.subplot(142)
-    plt.imshow(image_de_random_vertical)
-    plt.title("DE random_vertical image")
-
-    plt.subplot(143)
-    plt.imshow(image_pil_random_vertical)
-    plt.title("vertically flipped image")
-
-    plt.subplot(144)
-    plt.imshow(image_de_random_vertical - image_pil_random_vertical)
-    plt.title("Difference image, mse : {}".format(mse))
-    plt.show()
-
-
-def test_random_vertical_op():
+def test_random_vertical_op(plot=False):
     """
     Test random_vertical with default probability
     """
@@ -96,9 +72,10 @@ def test_random_vertical_op():
         diff = image_v_flipped - image_v_flipped_2
         mse = np.sum(np.power(diff, 2))
         logger.info("image_{}, mse: {}".format(num_iter + 1, mse))
-        # Uncomment below line if you want to visualize images
-        # visualize_with_mse(image_v_flipped, image_v_flipped_2, mse, image)
         num_iter += 1
+        if plot:
+            visualize_image(image, image_v_flipped, mse, image_v_flipped_2)
+
 
 def test_random_vertical_valid_prob_c():
     """
@@ -122,6 +99,7 @@ def test_random_vertical_valid_prob_c():
     ds.config.set_seed(original_seed)
     ds.config.set_num_parallel_workers(original_num_parallel_workers)
 
+
 def test_random_vertical_valid_prob_py():
     """
     Test RandomVerticalFlip op with py_transforms: valid non-default input, expect to pass
@@ -147,6 +125,7 @@ def test_random_vertical_valid_prob_py():
     ds.config.set_seed(original_seed)
     ds.config.set_num_parallel_workers(original_num_parallel_workers)
 
+
 def test_random_vertical_invalid_prob_c():
     """
     Test RandomVerticalFlip op in c_transforms: invalid input, expect to raise error
@@ -165,6 +144,7 @@ def test_random_vertical_invalid_prob_c():
         logger.info("Got an exception in DE: {}".format(str(e)))
         assert "Input is not" in str(e)
 
+
 def test_random_vertical_invalid_prob_py():
     """
     Test RandomVerticalFlip op in py_transforms: invalid input, expect to raise error
@@ -186,6 +166,7 @@ def test_random_vertical_invalid_prob_py():
         logger.info("Got an exception in DE: {}".format(str(e)))
         assert "Input is not" in str(e)
 
+
 def test_random_vertical_comp(plot=False):
     """
     Test test_random_vertical_flip and compare between python and c image augmentation ops
@@ -223,13 +204,13 @@ def test_random_vertical_comp(plot=False):
         mse = diff_mse(image_c, image_py)
         assert mse < 0.001
     if plot:
-        visualize(images_list_c, images_list_py)
+        visualize_list(images_list_c, images_list_py, visualize_mode=2)
 
 
 if __name__ == "__main__":
-    test_random_vertical_op()
+    test_random_vertical_op(plot=True)
     test_random_vertical_valid_prob_c()
     test_random_vertical_valid_prob_py()
     test_random_vertical_invalid_prob_c()
     test_random_vertical_invalid_prob_py()
-    test_random_vertical_comp(True)
+    test_random_vertical_comp(plot=True)

tests/ut/python/dataset/test_rescale_op.py

@@ -15,12 +15,10 @@
 """
 Testing the rescale op in DE
 """
-import matplotlib.pyplot as plt
-import numpy as np
-
 import mindspore.dataset as ds
 import mindspore.dataset.transforms.vision.c_transforms as vision
 from mindspore import log as logger
+from util import visualize_image, diff_mse
 
 DATA_DIR = ["../data/dataset/test_tf_file_3_images/train-0000-of-0001.data"]
 SCHEMA_DIR = "../data/dataset/test_tf_file_3_images/datasetSchema.json"
@@ -52,25 +50,7 @@ def get_rescaled(image_id):
     return None
 
 
-def visualize(image_de_rescaled, image_np_rescaled, mse):
-    """
-    visualizes the image using DE op and Numpy op
-    """
-    plt.subplot(131)
-    plt.imshow(image_de_rescaled)
-    plt.title("DE rescale image")
-
-    plt.subplot(132)
-    plt.imshow(image_np_rescaled)
-    plt.title("Numpy rescaled image")
-
-    plt.subplot(133)
-    plt.imshow(image_de_rescaled - image_np_rescaled)
-    plt.title("Difference image, mse : {}".format(mse))
-    plt.show()
-
-
-def test_rescale_op():
+def test_rescale_op(plot=False):
     """
     Test rescale
     """
@@ -83,19 +63,20 @@ def test_rescale_op():
 
     # apply map operations on images
     data1 = data1.map(input_columns=["image"], operations=decode_op)
-    data1 = data1.map(input_columns=["image"], operations=rescale_op)
+
+    data2 = data1.map(input_columns=["image"], operations=rescale_op)
 
     num_iter = 0
-    for item in data1.create_dict_iterator():
-        image_de_rescaled = item["image"]
+    for item1, item2 in zip(data1.create_dict_iterator(), data2.create_dict_iterator()):
+        image_original = item1["image"]
+        image_de_rescaled = item2["image"]
         image_np_rescaled = get_rescaled(num_iter)
-        diff = image_de_rescaled - image_np_rescaled
-        mse = np.sum(np.power(diff, 2))
+        mse = diff_mse(image_de_rescaled, image_np_rescaled)
         logger.info("image_{}, mse: {}".format(num_iter + 1, mse))
-        # Uncomment below line if you want to visualize images
-        # visualize(image_de_rescaled, image_np_rescaled, mse)
         num_iter += 1
+        if plot:
+            visualize_image(image_original, image_de_rescaled, mse, image_np_rescaled)
 
 
 if __name__ == "__main__":
-    test_rescale_op()
+    test_rescale_op(plot=True)

tests/ut/python/dataset/test_ten_crop.py

@@ -17,10 +17,10 @@ Testing TenCrop in DE
 import pytest
 import numpy as np
 
-from util import visualize, save_and_check_md5
 import mindspore.dataset as ds
 import mindspore.dataset.transforms.vision.py_transforms as vision
 from mindspore import log as logger
+from util import visualize_list, save_and_check_md5
 
 GENERATE_GOLDEN = False
 
@@ -62,7 +62,7 @@ def util_test_ten_crop(crop_size, vertical_flip=False, plot=False):
         logger.info("dtype of image_2: {}".format(image_2.dtype))
 
         if plot:
-            visualize(np.array([image_1]*10), (image_2 * 255).astype(np.uint8).transpose(0, 2, 3, 1))
+            visualize_list(np.array([image_1]*10), (image_2 * 255).astype(np.uint8).transpose(0, 2, 3, 1))
 
         # The output data should be of a 4D tensor shape, a stack of 10 images.
         assert len(image_2.shape) == 4

tests/ut/python/dataset/test_uniform_augment.py

@@ -12,35 +12,20 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-
-import matplotlib.pyplot as plt
+"""
+Testing UniformAugment in DE
+"""
 import numpy as np
 
 import mindspore.dataset.engine as de
 import mindspore.dataset.transforms.vision.c_transforms as C
 import mindspore.dataset.transforms.vision.py_transforms as F
 from mindspore import log as logger
+from util import visualize_list
 
 DATA_DIR = "../data/dataset/testImageNetData/train/"
 
 
-def visualize(image_original, image_ua):
-    """
-    visualizes the image using DE op and Numpy op
-    """
-    num = len(image_ua)
-    for i in range(num):
-        plt.subplot(2, num, i + 1)
-        plt.imshow(image_original[i])
-        plt.title("Original image")
-
-        plt.subplot(2, num, i + num + 1)
-        plt.imshow(image_ua[i])
-        plt.title("DE UniformAugment image")
-
-    plt.show()
-
-
 def test_uniform_augment(plot=False, num_ops=2):
     """
     Test UniformAugment
@@ -102,7 +87,7 @@ def test_uniform_augment(plot=False, num_ops=2):
     logger.info("MSE= {}".format(str(np.mean(mse))))
 
     if plot:
-        visualize(images_original, images_ua)
+        visualize_list(images_original, images_ua)
 
 
 def test_cpp_uniform_augment(plot=False, num_ops=2):
@@ -157,7 +142,7 @@ def test_cpp_uniform_augment(plot=False, num_ops=2):
                                   np.transpose(image, (0, 2, 3, 1)),
                                   axis=0)
     if plot:
-        visualize(images_original, images_ua)
+        visualize_list(images_original, images_ua)
 
     num_samples = images_original.shape[0]
     mse = np.zeros(num_samples)
@@ -226,6 +211,7 @@ def test_cpp_uniform_augment_exception_nonpositive_numops(num_ops=0):
         logger.info("Got an exception in DE: {}".format(str(e)))
         assert "num_ops" in str(e)
 
+
 def test_cpp_uniform_augment_exception_float_numops(num_ops=2.5):
     """
     Test UniformAugment invalid float number of ops
@@ -245,6 +231,7 @@ def test_cpp_uniform_augment_exception_float_numops(num_ops=2.5):
         logger.info("Got an exception in DE: {}".format(str(e)))
         assert "integer" in str(e)
 
+
 def test_cpp_uniform_augment_random_crop_badinput(num_ops=1):
     """
     Test UniformAugment with greater crop size
@@ -273,8 +260,8 @@ def test_cpp_uniform_augment_random_crop_badinput(num_ops=1):
 
 
 if __name__ == "__main__":
-    test_uniform_augment(num_ops=1)
-    test_cpp_uniform_augment(num_ops=1)
+    test_uniform_augment(num_ops=1, plot=True)
+    test_cpp_uniform_augment(num_ops=1, plot=True)
     test_cpp_uniform_augment_exception_pyops(num_ops=1)
     test_cpp_uniform_augment_exception_large_numops(num_ops=6)
     test_cpp_uniform_augment_exception_nonpositive_numops(num_ops=0)

tests/ut/python/dataset/util.py

@@ -25,6 +25,11 @@ from mindspore import log as logger
 # These are the column names defined in the testTFTestAllTypes dataset
 COLUMNS = ["col_1d", "col_2d", "col_3d", "col_binary", "col_float",
            "col_sint16", "col_sint32", "col_sint64"]
+# These are list of plot title in different visualize modes
+PLOT_TITLE_DICT = {
+    1: ["Original image", "Transformed image"],
+    2: ["c_transform image", "py_transform image"]
+}
 SAVE_JSON = False
 
 
@@ -206,19 +211,54 @@ def diff_me(in1, in2):
     return mse / 255 * 100
 
 
-def visualize(image_original, image_transformed):
+def visualize_list(image_list_1, image_list_2, visualize_mode=1):
     """
-    visualizes the image using DE op and Numpy op
+    visualizes a list of images using DE op
     """
-    num = len(image_transformed)
+    plot_title = PLOT_TITLE_DICT[visualize_mode]
+    num = len(image_list_1)
     for i in range(num):
         plt.subplot(2, num, i + 1)
-        plt.imshow(image_original[i])
-        plt.title("Original image")
+        plt.imshow(image_list_1[i])
+        plt.title(plot_title[0])
 
         plt.subplot(2, num, i + num + 1)
-        plt.imshow(image_transformed[i])
-        plt.title("Transformed image")
+        plt.imshow(image_list_2[i])
+        plt.title(plot_title[1])
+
+    plt.show()
+
+
+def visualize_image(image_original, image_de, mse=None, image_lib=None):
+    """
+    visualizes one example image with optional input: mse, image using 3rd party op.
+    If three images are passing in, different image is calculated by 2nd and 3rd images.
+    """
+    num = 2
+    if image_lib is not None:
+        num += 1
+    if mse is not None:
+        num += 1
+    plt.subplot(1, num, 1)
+    plt.imshow(image_original)
+    plt.title("Original image")
+
+    plt.subplot(1, num, 2)
+    plt.imshow(image_de)
+    plt.title("DE Op image")
+
+    if image_lib is not None:
+        plt.subplot(1, num, 3)
+        plt.imshow(image_lib)
+        plt.title("Lib Op image")
+        if mse is not None:
+            plt.subplot(1, num, 4)
+            plt.imshow(image_de - image_lib)
+            plt.title("Diff image,\n mse : {}".format(mse))
+    elif mse is not None:
+        plt.subplot(1, num, 3)
+        plt.imshow(image_original - image_de)
+        plt.title("Diff image,\n mse : {}".format(mse))
 
     plt.show()