未验证 提交 0e332fb5 编写于 作者: L littletomatodonkey 提交者: GitHub

add prob for ra (#2194)

* add prob for ra

* add prob for ra
上级 1a087990
...@@ -117,13 +117,16 @@ class RawRandAugment(object): ...@@ -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)
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册