|
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), |
|
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') |
|
|
