diff --git a/ppocr/data/imaug/operators.py b/ppocr/data/imaug/operators.py index efae1f417c709db9ec5f6275273ceae28e2f8597..09736515e7a388e191a12826e1e9e348e2fcde86 100644 --- a/ppocr/data/imaug/operators.py +++ b/ppocr/data/imaug/operators.py @@ -167,20 +167,37 @@ class KeepKeys(object): class Pad(object): - def __init__(self, size_div=32, **kwargs): + def __init__(self, size=None, size_div=32, **kwargs): + if size is not None and not isinstance(size, (int, list, tuple)): + raise TypeError("Type of target_size is invalid. Now is {}".format( + type(size))) + if isinstance(size, int): + size = [size, size] + self.size = size self.size_div = size_div def __call__(self, data): img = data['image'] - resize_h2 = max(int(math.ceil(img.shape[0] / 32) * 32), 32) - resize_w2 = max(int(math.ceil(img.shape[1] / 32) * 32), 32) + img_h, img_w = img.shape[0], img.shape[1] + if self.size: + resize_h2, resize_w2 = self.size + assert ( + img_h < resize_h2 and img_w < resize_w2 + ), '(h, w) of target size should be greater than (img_h, img_w)' + else: + resize_h2 = max( + int(math.ceil(img.shape[0] / self.size_div) * self.size_div), + self.size_div) + resize_w2 = max( + int(math.ceil(img.shape[1] / self.size_div) * self.size_div), + self.size_div) img = cv2.copyMakeBorder( img, 0, - resize_h2 - img.shape[0], + resize_h2 - img_h, 0, - resize_w2 - img.shape[1], + resize_w2 - img_w, cv2.BORDER_CONSTANT, value=0) data['image'] = img