sscbench/augmentation.py

import torchvision.transforms.functional as F

def get_color_aug_fn(params):
    fn_idx, brightness_factor, contrast_factor, saturation_factor, hue_factor = params
    def color_aug_fn(img):
        for fn_id in fn_idx:
            if fn_id == 0 and brightness_factor is not None:
                img = F.adjust_brightness(img, brightness_factor)
            elif fn_id == 1 and contrast_factor is not None:
                img = F.adjust_contrast(img, contrast_factor)
            elif fn_id == 2 and saturation_factor is not None:
                img = F.adjust_saturation(img, saturation_factor)
            elif fn_id == 3 and hue_factor is not None:
                img = F.adjust_hue(img, hue_factor)

        return img

    return color_aug_fn