from fastai.vision.all import *
import albumentations as Alb

class AlbTransform(Transform):
    def __init__(self, aug): self.aug = aug
    def encodes(self, img: PILImage):
        aug_img = self.aug(image=np.array(img))['image']
        return PILImage.create(aug_img)
    
def get_augs(): return Alb.Compose([
    Alb.Transpose(),
    Alb.Flip(),
    Alb.RandomRotate90(),
    Alb.HueSaturationValue(
      hue_shift_limit=5, 
      sat_shift_limit=5, 
      val_shift_limit=5 ),
])