Random Gamma Correction Data Augmentation Pytorch

RandGammaCorrection.py

import torchvision.transforms.functional as TF

class RandomGammaCorrection(object):

"""
Apply Gamma Correction to the images
"""

def __init__(self, gamma = None):
    self.gamma = gamma
    
    
def __call__(self,image):
    
    if self.gamma == None:
        # more chances of selecting 0 (original image)
        gammas = [0,0,0,0.5,1,1.5]
        self.gamma = random.choice(gammas)
    
    print(self.gamma)
    if self.gamma == 0:
        return image
    
    else:
        return TF.adjust_gamma(image, self.gamma, gain=1)