Merge branch 'dygraph' into fix_mem
This commit is contained in:
commit
6783c2a4c1
|
@ -117,13 +117,16 @@ class RawRandAugment(object):
|
||||||
class RandAugment(RawRandAugment):
|
class RandAugment(RawRandAugment):
|
||||||
""" RandAugment wrapper to auto fit different img types """
|
""" 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:
|
if six.PY2:
|
||||||
super(RandAugment, self).__init__(*args, **kwargs)
|
super(RandAugment, self).__init__(*args, **kwargs)
|
||||||
else:
|
else:
|
||||||
super().__init__(*args, **kwargs)
|
super().__init__(*args, **kwargs)
|
||||||
|
|
||||||
def __call__(self, data):
|
def __call__(self, data):
|
||||||
|
if np.random.rand() > self.prob:
|
||||||
|
return data
|
||||||
img = data['image']
|
img = data['image']
|
||||||
if not isinstance(img, Image.Image):
|
if not isinstance(img, Image.Image):
|
||||||
img = np.ascontiguousarray(img)
|
img = np.ascontiguousarray(img)
|
||||||
|
|
Loading…
Reference in New Issue