diff --git a/ppocr/data/imaug/randaugment.py b/ppocr/data/imaug/randaugment.py index 0bfac353..56f114d2 100644 --- a/ppocr/data/imaug/randaugment.py +++ b/ppocr/data/imaug/randaugment.py @@ -117,13 +117,16 @@ class RawRandAugment(object): class RandAugment(RawRandAugment): """ RandAugment wrapper to auto fit different img types """ - def __init__(self, *args, **kwargs): + def __init__(self, prob=0.5, *args, **kwargs): + self.prob = prob if six.PY2: super(RandAugment, self).__init__(*args, **kwargs) else: super().__init__(*args, **kwargs) def __call__(self, data): + if np.random.rand() > self.prob: + return data img = data['image'] if not isinstance(img, Image.Image): img = np.ascontiguousarray(img)