How to ensure augmentation applied to each class

Hi,

I am applying edge operation on 10% of images on the dataset. Suppose these are 6 classes.

How do I ensure that edge operation is applied to samples from each class?
(each class should have at least 1 images with edge operation applied to it)
Sample code below:

class cannyedge(object):
    def __init__(self, threshold1 = 100, threshold2 = 200):
        self.threshold1 = threshold1
        self.threshold2 = threshold2
        
    def __call__(self, sample):
        prob = np.random.random_sample()
        if prob < 0.1:
            sample = np.array(sample)
            edges = cv2.Canny(image=sample, threshold1=self.threshold1, threshold2=self.threshold2) 
            edge3c = np.zeros_like(sample)
            edge3c[:,:,0] = edges
            edge3c[:,:,1] = edges
            edge3c[:,:,2] = edges
            invedge = cv2.bitwise_not(edge3c)

            return invedge
        else:
            return sample

color_jitter = transforms.ColorJitter(0.8, 0.8, 0.8, 0.2)
data_transforms = transforms.Compose([transforms.RandomResizedCrop(size=IMAGE_SIZE),
                                              transforms.RandomHorizontalFlip(),
                                              transforms.RandomApply([color_jitter], p=0.5),
                                              transforms.RandomGrayscale(p=0.5),
                                              transforms.RandomRotation(30),
                                              cannyedge(),
                                              transforms.ToTensor(),
                                              
] )

Since you are randomly applying the transformations you won’t have a guarantee to apply it, but could estimate the likelihood of getting a transformed sample for each class.
Based on your description you could e.g. apply the transformations separately to at least one sample per class and let the other samples use the random process.