add prob for ra (#2194)

* add prob for ra

* add prob for ra
This commit is contained in:
littletomatodonkey 2021-03-09 15:08:59 +08:00 committed by GitHub
parent 1a087990d0
commit 0e332fb59a
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 4 additions and 1 deletions

View File

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