This commit is contained in:
littletomatodonkey 2020-12-13 17:19:52 +08:00 committed by GitHub
parent 52671b7db2
commit e84ea2667f
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 4 additions and 5 deletions

View File

@ -35,12 +35,13 @@ from .text_image_aug import tia_perspective, tia_stretch, tia_distort
class RecAug(object):
def __init__(self, use_tia=True, **kwargsz):
def __init__(self, use_tia=True, aug_prob=0.4, **kwargs):
self.use_tia = use_tia
self.aug_prob = aug_prob
def __call__(self, data):
img = data['image']
img = warp(img, 10, self.use_tia)
img = warp(img, 10, self.use_tia, self.aug_prob)
data['image'] = img
return data
@ -329,7 +330,7 @@ def get_warpAffine(config):
return rz
def warp(img, ang, use_tia=True):
def warp(img, ang, use_tia=True, prob=0.4):
"""
warp
"""
@ -338,8 +339,6 @@ def warp(img, ang, use_tia=True):
config.make(w, h, ang)
new_img = img
prob = 0.4
if config.distort:
img_height, img_width = img.shape[0:2]
if random.random() <= prob and img_height >= 20 and img_width >= 20: