diff --git a/ppocr/data/imaug/rec_img_aug.py b/ppocr/data/imaug/rec_img_aug.py index 3ab5a8f3238c3e00a143d8d8058ae65bd2bb8fd0..2ccb2d1d2b6780138098f08c78cce3be3e3b9ceb 100644 --- a/ppocr/data/imaug/rec_img_aug.py +++ b/ppocr/data/imaug/rec_img_aug.py @@ -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: