mindspore/tests/ut/python/dataset/test_rand_augment.py

# Copyright 2022 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""
Testing RandAugment in DE
"""
import numpy as np

import mindspore.dataset as ds
from mindspore.dataset.vision import Decode, RandAugment, Resize, Inter
from mindspore import log as logger
from util import visualize_image, visualize_list, diff_mse

IMAGE_FILE = "../data/dataset/testImageNetData/train/class1/1_1.jpg"
DATA_DIR = "../data/dataset/testImageNetData/train/"


def test_rand_augment_pipeline(plot=False):
    """
    Feature: RandAugment
    Description: Test RandAugment pipeline
    Expectation: Pass without error
    """
    logger.info("Test RandAugment pipeline")

    # Original Images
    data_set = ds.ImageFolderDataset(dataset_dir=DATA_DIR, shuffle=False)
    transforms_original = [Decode(), Resize(size=[224, 224])]
    ds_original = data_set.map(operations=transforms_original, input_columns="image")
    ds_original = ds_original.batch(512)

    for idx, (image, _) in enumerate(ds_original):
        if idx == 0:
            images_original = image.asnumpy()
        else:
            images_original = np.append(images_original, image.asnumpy(), axis=0)

    data_set1 = ds.ImageFolderDataset(dataset_dir=DATA_DIR, shuffle=False)
    rand_augment_op = RandAugment(3, 10, 15, Inter.BICUBIC, 20)
    transforms = [Decode(), Resize(size=[224, 224]), rand_augment_op]
    ds_rand_augment = data_set1.map(operations=transforms, input_columns="image")
    ds_rand_augment = ds_rand_augment.batch(512)
    for idx, (image, _) in enumerate(ds_rand_augment):
        if idx == 0:
            images_rand_augment = image.asnumpy()
        else:
            images_rand_augment = np.append(images_rand_augment, image.asnumpy(), axis=0)
    assert images_original.shape[0] == images_rand_augment.shape[0]
    if plot:
        visualize_list(images_original, images_rand_augment)

    num_samples = images_original.shape[0]
    mse = np.zeros(num_samples)
    for i in range(num_samples):
        mse[i] = diff_mse(images_rand_augment[i], images_original[i])
    logger.info("MSE= {}".format(str(np.mean(mse))))


def test_rand_augment_eager(plot=False):
    """
    Feature: RandAugment
    Description: Test RandAugment in eager mode
    Expectation: Pass without error
    """
    img = np.fromfile(IMAGE_FILE, dtype=np.uint8)
    logger.info("Image.type: {}, Image.shape: {}".format(type(img), img.shape))
    img = Decode()(img)
    img_rand_augmented = RandAugment()(img)
    if plot:
        visualize_image(img, img_rand_augmented)
    logger.info("Image.type: {}, Image.shape: {}".format(type(img_rand_augmented), img_rand_augmented.shape))
    mse = diff_mse(img_rand_augmented, img)
    logger.info("MSE= {}".format(str(mse)))


def test_rand_augment_invalid_params_int():
    """
    Feature: RandAugment
    Description: Test RandAugment with invalid first three parameters
    Expectation: Error is raised as expected
    """
    logger.info("test_rand_augment_invalid_params_int")
    dataset = ds.ImageFolderDataset(DATA_DIR, 1, shuffle=False, decode=True)
    try:
        rand_augment_op = RandAugment(num_ops=-1)
        dataset.map(operations=rand_augment_op, input_columns=['image'])
    except ValueError as e:
        logger.info("Got an exception in DE: {}".format(str(e)))
        assert "Input num_ops is not within the required interval of [0, 16777216]." in str(e)
    try:
        rand_augment_op = RandAugment(magnitude=-1)
        dataset.map(operations=rand_augment_op, input_columns=['image'])
    except ValueError as e:
        logger.info("Got an exception in DE: {}".format(str(e)))
        assert "Input magnitude is not within the required interval of [0, 31)." in str(e)
    try:
        rand_augment_op = RandAugment(magnitude=0, num_magnitude_bins=1)
        dataset.map(operations=rand_augment_op, input_columns=['image'])
    except ValueError as e:
        logger.info("Got an exception in DE: {}".format(str(e)))
        assert "Input num_magnitude_bins is not within the required interval of [2, 16777216]." in str(e)


def test_rand_augment_invalid_interpolation():
    """
    Feature: RandAugment
    Description: Test RandAugment with invalid interpolation
    Expectation: Error is raised as expected
    """
    logger.info("test_rand_augment_invalid_interpolation")
    dataset = ds.ImageFolderDataset(DATA_DIR, 1, shuffle=False, decode=True)
    try:
        rand_augment_op = RandAugment(interpolation="invalid")
        dataset.map(operations=rand_augment_op, input_columns=['image'])
    except TypeError as e:
        logger.info("Got an exception in DE: {}".format(str(e)))
        assert "Argument interpolation with value invalid is not of type [<enum 'Inter'>]" in str(e)


def test_rand_augment_invalid_fill_value():
    """
    Feature: RandAugment
    Description: Test RandAugment with invalid fill_value
    Expectation: Correct error is raised as expected
    """
    logger.info("test_rand_augment_invalid_fill_value")
    dataset = ds.ImageFolderDataset(DATA_DIR, 1, shuffle=False, decode=True)
    try:
        rand_augment_op = RandAugment(fill_value=(10, 10))
        dataset.map(operations=rand_augment_op, input_columns=['image'])
    except TypeError as e:
        logger.info("Got an exception in DE: {}".format(str(e)))
        assert "fill_value should be a single integer or a 3-tuple." in str(e)
    try:
        rand_augment_op = RandAugment(fill_value=-1)
        dataset.map(operations=rand_augment_op, input_columns=['image'])
    except ValueError as e:
        logger.info("Got an exception in DE: {}".format(str(e)))
        assert "is not within the required interval of [0, 255]." in str(e)


def test_rand_augment_invalid_magnitude_value():
    """
    Feature: RandAugment
    Description: Test RandAugment with invalid magnitude_value
    Expectation: Correct error is raised as expected
    """
    logger.info("test_rand_augment_invalid_magnitude_value")
    try:
        _ = RandAugment(3, 4, 3)
    except ValueError as e:
        logger.info("Got an exception in DE: {}".format(str(e)))
        assert "Input magnitude is not within the required interval of [0, 3)." in str(e)


if __name__ == "__main__":
    test_rand_augment_pipeline(plot=True)
    test_rand_augment_eager()
    test_rand_augment_invalid_params_int()
    test_rand_augment_invalid_interpolation()
    test_rand_augment_invalid_fill_value()
    test_rand_augment_invalid_magnitude_value()