Spaces:
Running
Running
| import albumentations as A | |
| class DarkAug(object): | |
| """ | |
| Extreme dark augmentation aiming at Aachen Day-Night | |
| """ | |
| def __init__(self) -> None: | |
| self.augmentor = A.Compose([ | |
| A.RandomBrightnessContrast(p=0.75, brightness_limit=(-0.6, 0.0), contrast_limit=(-0.5, 0.3)), | |
| A.Blur(p=0.1, blur_limit=(3, 9)), | |
| A.MotionBlur(p=0.2, blur_limit=(3, 25)), | |
| A.RandomGamma(p=0.1, gamma_limit=(15, 65)), | |
| A.HueSaturationValue(p=0.1, val_shift_limit=(-100, -40)) | |
| ], p=0.75) | |
| def __call__(self, x): | |
| return self.augmentor(image=x)['image'] | |
| class MobileAug(object): | |
| """ | |
| Random augmentations aiming at images of mobile/handhold devices. | |
| """ | |
| def __init__(self): | |
| self.augmentor = A.Compose([ | |
| A.MotionBlur(p=0.25), | |
| A.ColorJitter(p=0.5), | |
| A.RandomRain(p=0.1), # random occlusion | |
| A.RandomSunFlare(p=0.1), | |
| A.JpegCompression(p=0.25), | |
| A.ISONoise(p=0.25) | |
| ], p=1.0) | |
| def __call__(self, x): | |
| return self.augmentor(image=x)['image'] | |
| def build_augmentor(method=None, **kwargs): | |
| if method is not None: | |
| raise NotImplementedError('Using of augmentation functions are not supported yet!') | |
| if method == 'dark': | |
| return DarkAug() | |
| elif method == 'mobile': | |
| return MobileAug() | |
| elif method is None: | |
| return None | |
| else: | |
| raise ValueError(f'Invalid augmentation method: {method}') | |
| if __name__ == '__main__': | |
| augmentor = build_augmentor('FDA') | |