augmentation.py

if self.transform:
    batch_n = idx // self.bs
    if batch_n % 2 == 0:
        self.transform[1].p = 1
    else:
        self.transform[1].p = 0
        
    # albumentations requires bboxes to be (x,y,w,h,class_idx)
    augmentations = self.transform(image=img,
                                   bboxes=np.roll(labels, axis=1, shift=4)
                                   )
    img = augmentations["image"]
    # loss fx requires bboxes to be (class_idx,x,y,w,h)
    labels = np.array(augmentations["bboxes"])
    if len(labels):
        labels = np.roll(labels, axis=1, shift=1)