Just wanted to share my code for adding advanced data augmentation into fastai training pipeline. I really wonder if there is a more elegant way to do it?
Libraries like albumentations or imgaug require PIL like(or CV2 converted to RGB) image array. But fastai augmentations work with tensor. So we have to convert the data back and forth. Also, there is some wrapping required by fastai.
import albumentations as A def tensor2np(x): np_image = x.cpu().permute(1, 2, 0).numpy() np_image = (np_image * 255).astype(np.uint8) return np_image def alb_tfm2fastai(alb_tfm): def _alb_transformer(x): # tensor to numpy np_image = tensor2np(x) # apply albumentations transformed = alb_tfm(image=np_image)['image'] # back to tensor tensor_image = pil2tensor(transformed, np.float32) tensor_image.div_(255) return tensor_image transformer = TfmPixel(_alb_transformer) return transformer()
You can use this wrapper method with any augmentations set like this:
tfms = alb_tfm2fastai(A.HueSaturationValue())
And then pass it to
xtra_tfms or to
transform as the only set of transforms.
I hope someone will find it useful!