From e84ea2667f3da91c75e03dde86441ef3c292db3c Mon Sep 17 00:00:00 2001 From: littletomatodonkey <2120160898@bit.edu.cn> Date: Sun, 13 Dec 2020 17:19:52 +0800 Subject: [PATCH] fix prob (#1404) --- ppocr/data/imaug/rec_img_aug.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/ppocr/data/imaug/rec_img_aug.py b/ppocr/data/imaug/rec_img_aug.py index 3ab5a8f3..2ccb2d1d 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: