From 0e332fb59a52fb093ea747d79a4019c99fe15cbe Mon Sep 17 00:00:00 2001
From: littletomatodonkey <2120160898@bit.edu.cn>
Date: Tue, 9 Mar 2021 15:08:59 +0800
Subject: [PATCH] add prob for ra (#2194)

* add prob for ra

* add prob for ra
---
 ppocr/data/imaug/randaugment.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

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)
-- 
GitLab