This commit is contained in:
27
ppocr/data/imaug/ColorJitter.py
Normal file
27
ppocr/data/imaug/ColorJitter.py
Normal file
@@ -0,0 +1,27 @@
|
||||
# copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
from paddle.vision.transforms import ColorJitter as pp_ColorJitter
|
||||
|
||||
__all__ = ["ColorJitter"]
|
||||
|
||||
|
||||
class ColorJitter(object):
|
||||
def __init__(self, brightness=0, contrast=0, saturation=0, hue=0, **kwargs):
|
||||
self.aug = pp_ColorJitter(brightness, contrast, saturation, hue)
|
||||
|
||||
def __call__(self, data):
|
||||
image = data["image"]
|
||||
image = self.aug(image)
|
||||
data["image"] = image
|
||||
return data
|
||||
96
ppocr/data/imaug/__init__.py
Normal file
96
ppocr/data/imaug/__init__.py
Normal file
@@ -0,0 +1,96 @@
|
||||
# copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
from __future__ import absolute_import
|
||||
from __future__ import division
|
||||
from __future__ import print_function
|
||||
from __future__ import unicode_literals
|
||||
|
||||
from .iaa_augment import IaaAugment
|
||||
from .make_border_map import MakeBorderMap
|
||||
from .make_shrink_map import MakeShrinkMap
|
||||
from .random_crop_data import EastRandomCropData, RandomCropImgMask
|
||||
from .make_pse_gt import MakePseGt
|
||||
|
||||
|
||||
from .rec_img_aug import (
|
||||
BaseDataAugmentation,
|
||||
RecAug,
|
||||
RecConAug,
|
||||
RecResizeImg,
|
||||
ClsResizeImg,
|
||||
SRNRecResizeImg,
|
||||
GrayRecResizeImg,
|
||||
SARRecResizeImg,
|
||||
PRENResizeImg,
|
||||
ABINetRecResizeImg,
|
||||
SVTRRecResizeImg,
|
||||
ABINetRecAug,
|
||||
VLRecResizeImg,
|
||||
SPINRecResizeImg,
|
||||
RobustScannerRecResizeImg,
|
||||
RFLRecResizeImg,
|
||||
SVTRRecAug,
|
||||
ParseQRecAug,
|
||||
)
|
||||
from .ssl_img_aug import SSLRotateResize
|
||||
from .randaugment import RandAugment
|
||||
from .copy_paste import CopyPaste
|
||||
from .ColorJitter import ColorJitter
|
||||
from .operators import *
|
||||
from .label_ops import *
|
||||
|
||||
from .east_process import *
|
||||
from .sast_process import *
|
||||
from .pg_process import *
|
||||
from .table_ops import *
|
||||
|
||||
from .vqa import *
|
||||
|
||||
from .fce_aug import *
|
||||
from .fce_targets import FCENetTargets
|
||||
from .ct_process import *
|
||||
from .drrg_targets import DRRGTargets
|
||||
from .latex_ocr_aug import *
|
||||
from .unimernet_aug import *
|
||||
|
||||
|
||||
def transform(data, ops=None):
|
||||
"""transform"""
|
||||
if ops is None:
|
||||
ops = []
|
||||
for op in ops:
|
||||
data = op(data)
|
||||
if data is None:
|
||||
return None
|
||||
return data
|
||||
|
||||
|
||||
def create_operators(op_param_list, global_config=None):
|
||||
"""
|
||||
create operators based on the config
|
||||
|
||||
Args:
|
||||
params(list): a dict list, used to create some operators
|
||||
"""
|
||||
assert isinstance(op_param_list, list), "operator config should be a list"
|
||||
ops = []
|
||||
for operator in op_param_list:
|
||||
assert isinstance(operator, dict) and len(operator) == 1, "yaml format error"
|
||||
op_name = list(operator)[0]
|
||||
param = {} if operator[op_name] is None else operator[op_name]
|
||||
if global_config is not None:
|
||||
param.update(global_config)
|
||||
op = eval(op_name)(**param)
|
||||
ops.append(op)
|
||||
return ops
|
||||
512
ppocr/data/imaug/abinet_aug.py
Normal file
512
ppocr/data/imaug/abinet_aug.py
Normal file
@@ -0,0 +1,512 @@
|
||||
# copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
"""
|
||||
This code is refer from:
|
||||
https://github.com/FangShancheng/ABINet/blob/main/transforms.py
|
||||
"""
|
||||
import math
|
||||
import numbers
|
||||
import random
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
from paddle.vision.transforms import Compose, ColorJitter
|
||||
|
||||
|
||||
def sample_asym(magnitude, size=None):
|
||||
return np.random.beta(1, 4, size) * magnitude
|
||||
|
||||
|
||||
def sample_sym(magnitude, size=None):
|
||||
return (np.random.beta(4, 4, size=size) - 0.5) * 2 * magnitude
|
||||
|
||||
|
||||
def sample_uniform(low, high, size=None):
|
||||
return np.random.uniform(low, high, size=size)
|
||||
|
||||
|
||||
def get_interpolation(type="random"):
|
||||
if type == "random":
|
||||
choice = [cv2.INTER_NEAREST, cv2.INTER_LINEAR, cv2.INTER_CUBIC, cv2.INTER_AREA]
|
||||
interpolation = choice[random.randint(0, len(choice) - 1)]
|
||||
elif type == "nearest":
|
||||
interpolation = cv2.INTER_NEAREST
|
||||
elif type == "linear":
|
||||
interpolation = cv2.INTER_LINEAR
|
||||
elif type == "cubic":
|
||||
interpolation = cv2.INTER_CUBIC
|
||||
elif type == "area":
|
||||
interpolation = cv2.INTER_AREA
|
||||
else:
|
||||
raise TypeError(
|
||||
"Interpolation types only nearest, linear, cubic, area are supported!"
|
||||
)
|
||||
return interpolation
|
||||
|
||||
|
||||
class CVRandomRotation(object):
|
||||
def __init__(self, degrees=15):
|
||||
assert isinstance(degrees, numbers.Number), "degree should be a single number."
|
||||
assert degrees >= 0, "degree must be positive."
|
||||
self.degrees = degrees
|
||||
|
||||
@staticmethod
|
||||
def get_params(degrees):
|
||||
return sample_sym(degrees)
|
||||
|
||||
def __call__(self, img):
|
||||
angle = self.get_params(self.degrees)
|
||||
src_h, src_w = img.shape[:2]
|
||||
M = cv2.getRotationMatrix2D(
|
||||
center=(src_w / 2, src_h / 2), angle=angle, scale=1.0
|
||||
)
|
||||
abs_cos, abs_sin = abs(M[0, 0]), abs(M[0, 1])
|
||||
dst_w = int(src_h * abs_sin + src_w * abs_cos)
|
||||
dst_h = int(src_h * abs_cos + src_w * abs_sin)
|
||||
M[0, 2] += (dst_w - src_w) / 2
|
||||
M[1, 2] += (dst_h - src_h) / 2
|
||||
|
||||
flags = get_interpolation()
|
||||
return cv2.warpAffine(
|
||||
img, M, (dst_w, dst_h), flags=flags, borderMode=cv2.BORDER_REPLICATE
|
||||
)
|
||||
|
||||
|
||||
class CVRandomAffine(object):
|
||||
def __init__(self, degrees, translate=None, scale=None, shear=None):
|
||||
assert isinstance(degrees, numbers.Number), "degree should be a single number."
|
||||
assert degrees >= 0, "degree must be positive."
|
||||
self.degrees = degrees
|
||||
|
||||
if translate is not None:
|
||||
assert (
|
||||
isinstance(translate, (tuple, list)) and len(translate) == 2
|
||||
), "translate should be a list or tuple and it must be of length 2."
|
||||
for t in translate:
|
||||
if not (0.0 <= t <= 1.0):
|
||||
raise ValueError("translation values should be between 0 and 1")
|
||||
self.translate = translate
|
||||
|
||||
if scale is not None:
|
||||
assert (
|
||||
isinstance(scale, (tuple, list)) and len(scale) == 2
|
||||
), "scale should be a list or tuple and it must be of length 2."
|
||||
for s in scale:
|
||||
if s <= 0:
|
||||
raise ValueError("scale values should be positive")
|
||||
self.scale = scale
|
||||
|
||||
if shear is not None:
|
||||
if isinstance(shear, numbers.Number):
|
||||
if shear < 0:
|
||||
raise ValueError(
|
||||
"If shear is a single number, it must be positive."
|
||||
)
|
||||
self.shear = [shear]
|
||||
else:
|
||||
assert isinstance(shear, (tuple, list)) and (
|
||||
len(shear) == 2
|
||||
), "shear should be a list or tuple and it must be of length 2."
|
||||
self.shear = shear
|
||||
else:
|
||||
self.shear = shear
|
||||
|
||||
def _get_inverse_affine_matrix(self, center, angle, translate, scale, shear):
|
||||
# https://github.com/pytorch/vision/blob/v0.4.0/torchvision/transforms/functional.py#L717
|
||||
from numpy import sin, cos, tan
|
||||
|
||||
if isinstance(shear, numbers.Number):
|
||||
shear = [shear, 0]
|
||||
|
||||
if not isinstance(shear, (tuple, list)) and len(shear) == 2:
|
||||
raise ValueError(
|
||||
"Shear should be a single value or a tuple/list containing "
|
||||
+ "two values. Got {}".format(shear)
|
||||
)
|
||||
|
||||
rot = math.radians(angle)
|
||||
sx, sy = [math.radians(s) for s in shear]
|
||||
|
||||
cx, cy = center
|
||||
tx, ty = translate
|
||||
|
||||
# RSS without scaling
|
||||
a = cos(rot - sy) / cos(sy)
|
||||
b = -cos(rot - sy) * tan(sx) / cos(sy) - sin(rot)
|
||||
c = sin(rot - sy) / cos(sy)
|
||||
d = -sin(rot - sy) * tan(sx) / cos(sy) + cos(rot)
|
||||
|
||||
# Inverted rotation matrix with scale and shear
|
||||
# det([[a, b], [c, d]]) == 1, since det(rotation) = 1 and det(shear) = 1
|
||||
M = [d, -b, 0, -c, a, 0]
|
||||
M = [x / scale for x in M]
|
||||
|
||||
# Apply inverse of translation and of center translation: RSS^-1 * C^-1 * T^-1
|
||||
M[2] += M[0] * (-cx - tx) + M[1] * (-cy - ty)
|
||||
M[5] += M[3] * (-cx - tx) + M[4] * (-cy - ty)
|
||||
|
||||
# Apply center translation: C * RSS^-1 * C^-1 * T^-1
|
||||
M[2] += cx
|
||||
M[5] += cy
|
||||
return M
|
||||
|
||||
@staticmethod
|
||||
def get_params(degrees, translate, scale_ranges, shears, height):
|
||||
angle = sample_sym(degrees)
|
||||
if translate is not None:
|
||||
max_dx = translate[0] * height
|
||||
max_dy = translate[1] * height
|
||||
translations = (np.round(sample_sym(max_dx)), np.round(sample_sym(max_dy)))
|
||||
else:
|
||||
translations = (0, 0)
|
||||
|
||||
if scale_ranges is not None:
|
||||
scale = sample_uniform(scale_ranges[0], scale_ranges[1])
|
||||
else:
|
||||
scale = 1.0
|
||||
|
||||
if shears is not None:
|
||||
if len(shears) == 1:
|
||||
shear = [sample_sym(shears[0]), 0.0]
|
||||
elif len(shears) == 2:
|
||||
shear = [sample_sym(shears[0]), sample_sym(shears[1])]
|
||||
else:
|
||||
shear = 0.0
|
||||
|
||||
return angle, translations, scale, shear
|
||||
|
||||
def __call__(self, img):
|
||||
src_h, src_w = img.shape[:2]
|
||||
angle, translate, scale, shear = self.get_params(
|
||||
self.degrees, self.translate, self.scale, self.shear, src_h
|
||||
)
|
||||
|
||||
M = self._get_inverse_affine_matrix(
|
||||
(src_w / 2, src_h / 2), angle, (0, 0), scale, shear
|
||||
)
|
||||
M = np.array(M).reshape(2, 3)
|
||||
|
||||
startpoints = [(0, 0), (src_w - 1, 0), (src_w - 1, src_h - 1), (0, src_h - 1)]
|
||||
project = lambda x, y, a, b, c: int(a * x + b * y + c)
|
||||
endpoints = [
|
||||
(project(x, y, *M[0]), project(x, y, *M[1])) for x, y in startpoints
|
||||
]
|
||||
|
||||
rect = cv2.minAreaRect(np.array(endpoints))
|
||||
bbox = cv2.boxPoints(rect).astype(dtype=np.int32)
|
||||
max_x, max_y = bbox[:, 0].max(), bbox[:, 1].max()
|
||||
min_x, min_y = bbox[:, 0].min(), bbox[:, 1].min()
|
||||
|
||||
dst_w = int(max_x - min_x)
|
||||
dst_h = int(max_y - min_y)
|
||||
M[0, 2] += (dst_w - src_w) / 2
|
||||
M[1, 2] += (dst_h - src_h) / 2
|
||||
|
||||
# add translate
|
||||
dst_w += int(abs(translate[0]))
|
||||
dst_h += int(abs(translate[1]))
|
||||
if translate[0] < 0:
|
||||
M[0, 2] += abs(translate[0])
|
||||
if translate[1] < 0:
|
||||
M[1, 2] += abs(translate[1])
|
||||
|
||||
flags = get_interpolation()
|
||||
return cv2.warpAffine(
|
||||
img, M, (dst_w, dst_h), flags=flags, borderMode=cv2.BORDER_REPLICATE
|
||||
)
|
||||
|
||||
|
||||
class CVRandomPerspective(object):
|
||||
def __init__(self, distortion=0.5):
|
||||
self.distortion = distortion
|
||||
|
||||
def get_params(self, width, height, distortion):
|
||||
offset_h = sample_asym(distortion * height / 2, size=4).astype(dtype=np.int32)
|
||||
offset_w = sample_asym(distortion * width / 2, size=4).astype(dtype=np.int32)
|
||||
topleft = (offset_w[0], offset_h[0])
|
||||
topright = (width - 1 - offset_w[1], offset_h[1])
|
||||
botright = (width - 1 - offset_w[2], height - 1 - offset_h[2])
|
||||
botleft = (offset_w[3], height - 1 - offset_h[3])
|
||||
|
||||
startpoints = [(0, 0), (width - 1, 0), (width - 1, height - 1), (0, height - 1)]
|
||||
endpoints = [topleft, topright, botright, botleft]
|
||||
return np.array(startpoints, dtype=np.float32), np.array(
|
||||
endpoints, dtype=np.float32
|
||||
)
|
||||
|
||||
def __call__(self, img):
|
||||
height, width = img.shape[:2]
|
||||
startpoints, endpoints = self.get_params(width, height, self.distortion)
|
||||
M = cv2.getPerspectiveTransform(startpoints, endpoints)
|
||||
|
||||
# TODO: more robust way to crop image
|
||||
rect = cv2.minAreaRect(endpoints)
|
||||
bbox = cv2.boxPoints(rect).astype(dtype=np.int32)
|
||||
max_x, max_y = bbox[:, 0].max(), bbox[:, 1].max()
|
||||
min_x, min_y = bbox[:, 0].min(), bbox[:, 1].min()
|
||||
min_x, min_y = max(min_x, 0), max(min_y, 0)
|
||||
|
||||
flags = get_interpolation()
|
||||
img = cv2.warpPerspective(
|
||||
img, M, (max_x, max_y), flags=flags, borderMode=cv2.BORDER_REPLICATE
|
||||
)
|
||||
img = img[min_y:, min_x:]
|
||||
return img
|
||||
|
||||
|
||||
class CVRescale(object):
|
||||
def __init__(self, factor=4, base_size=(128, 512)):
|
||||
"""Define image scales using gaussian pyramid and rescale image to target scale.
|
||||
|
||||
Args:
|
||||
factor: the decayed factor from base size, factor=4 keeps target scale by default.
|
||||
base_size: base size the build the bottom layer of pyramid
|
||||
"""
|
||||
if isinstance(factor, numbers.Number):
|
||||
self.factor = round(sample_uniform(0, factor))
|
||||
elif isinstance(factor, (tuple, list)) and len(factor) == 2:
|
||||
self.factor = round(sample_uniform(factor[0], factor[1]))
|
||||
else:
|
||||
raise Exception("factor must be number or list with length 2")
|
||||
# assert factor is valid
|
||||
self.base_h, self.base_w = base_size[:2]
|
||||
|
||||
def __call__(self, img):
|
||||
if self.factor == 0:
|
||||
return img
|
||||
src_h, src_w = img.shape[:2]
|
||||
cur_w, cur_h = self.base_w, self.base_h
|
||||
scale_img = cv2.resize(img, (cur_w, cur_h), interpolation=get_interpolation())
|
||||
for _ in range(self.factor):
|
||||
scale_img = cv2.pyrDown(scale_img)
|
||||
scale_img = cv2.resize(
|
||||
scale_img, (src_w, src_h), interpolation=get_interpolation()
|
||||
)
|
||||
return scale_img
|
||||
|
||||
|
||||
class CVGaussianNoise(object):
|
||||
def __init__(self, mean=0, var=20):
|
||||
self.mean = mean
|
||||
if isinstance(var, numbers.Number):
|
||||
self.var = max(int(sample_asym(var)), 1)
|
||||
elif isinstance(var, (tuple, list)) and len(var) == 2:
|
||||
self.var = int(sample_uniform(var[0], var[1]))
|
||||
else:
|
||||
raise Exception("degree must be number or list with length 2")
|
||||
|
||||
def __call__(self, img):
|
||||
noise = np.random.normal(self.mean, self.var**0.5, img.shape)
|
||||
img = np.clip(img + noise, 0, 255).astype(np.uint8)
|
||||
return img
|
||||
|
||||
|
||||
class CVPossionNoise(object):
|
||||
def __init__(self, lam=20):
|
||||
self.lam = lam
|
||||
if isinstance(lam, numbers.Number):
|
||||
self.lam = max(int(sample_asym(lam)), 1)
|
||||
elif isinstance(lam, (tuple, list)) and len(lam) == 2:
|
||||
self.lam = int(sample_uniform(lam[0], lam[1]))
|
||||
else:
|
||||
raise Exception("lam must be number or list with length 2")
|
||||
|
||||
def __call__(self, img):
|
||||
noise = np.random.poisson(lam=self.lam, size=img.shape)
|
||||
img = np.clip(img + noise, 0, 255).astype(np.uint8)
|
||||
return img
|
||||
|
||||
|
||||
class CVGaussionBlur(object):
|
||||
def __init__(self, radius):
|
||||
self.radius = radius
|
||||
if isinstance(radius, numbers.Number):
|
||||
self.radius = max(int(sample_asym(radius)), 1)
|
||||
elif isinstance(radius, (tuple, list)) and len(radius) == 2:
|
||||
self.radius = int(sample_uniform(radius[0], radius[1]))
|
||||
else:
|
||||
raise Exception("radius must be number or list with length 2")
|
||||
|
||||
def __call__(self, img):
|
||||
fil = cv2.getGaussianKernel(ksize=self.radius, sigma=1, ktype=cv2.CV_32F)
|
||||
img = cv2.sepFilter2D(img, -1, fil, fil)
|
||||
return img
|
||||
|
||||
|
||||
class CVMotionBlur(object):
|
||||
def __init__(self, degrees=12, angle=90):
|
||||
if isinstance(degrees, numbers.Number):
|
||||
self.degree = max(int(sample_asym(degrees)), 1)
|
||||
elif isinstance(degrees, (tuple, list)) and len(degrees) == 2:
|
||||
self.degree = int(sample_uniform(degrees[0], degrees[1]))
|
||||
else:
|
||||
raise Exception("degree must be number or list with length 2")
|
||||
self.angle = sample_uniform(-angle, angle)
|
||||
|
||||
def __call__(self, img):
|
||||
M = cv2.getRotationMatrix2D((self.degree // 2, self.degree // 2), self.angle, 1)
|
||||
motion_blur_kernel = np.zeros((self.degree, self.degree))
|
||||
motion_blur_kernel[self.degree // 2, :] = 1
|
||||
motion_blur_kernel = cv2.warpAffine(
|
||||
motion_blur_kernel, M, (self.degree, self.degree)
|
||||
)
|
||||
motion_blur_kernel = motion_blur_kernel / self.degree
|
||||
img = cv2.filter2D(img, -1, motion_blur_kernel)
|
||||
img = np.clip(img, 0, 255).astype(np.uint8)
|
||||
return img
|
||||
|
||||
|
||||
class CVGeometry(object):
|
||||
def __init__(
|
||||
self,
|
||||
degrees=15,
|
||||
translate=(0.3, 0.3),
|
||||
scale=(0.5, 2.0),
|
||||
shear=(45, 15),
|
||||
distortion=0.5,
|
||||
p=0.5,
|
||||
):
|
||||
self.p = p
|
||||
type_p = random.random()
|
||||
if type_p < 0.33:
|
||||
self.transforms = CVRandomRotation(degrees=degrees)
|
||||
elif type_p < 0.66:
|
||||
self.transforms = CVRandomAffine(
|
||||
degrees=degrees, translate=translate, scale=scale, shear=shear
|
||||
)
|
||||
else:
|
||||
self.transforms = CVRandomPerspective(distortion=distortion)
|
||||
|
||||
def __call__(self, img):
|
||||
if random.random() < self.p:
|
||||
return self.transforms(img)
|
||||
else:
|
||||
return img
|
||||
|
||||
|
||||
class CVDeterioration(object):
|
||||
def __init__(self, var, degrees, factor, p=0.5):
|
||||
self.p = p
|
||||
transforms = []
|
||||
if var is not None:
|
||||
transforms.append(CVGaussianNoise(var=var))
|
||||
if degrees is not None:
|
||||
transforms.append(CVMotionBlur(degrees=degrees))
|
||||
if factor is not None:
|
||||
transforms.append(CVRescale(factor=factor))
|
||||
|
||||
random.shuffle(transforms)
|
||||
transforms = Compose(transforms)
|
||||
self.transforms = transforms
|
||||
|
||||
def __call__(self, img):
|
||||
if random.random() < self.p:
|
||||
return self.transforms(img)
|
||||
else:
|
||||
return img
|
||||
|
||||
|
||||
class CVColorJitter(object):
|
||||
def __init__(self, brightness=0.5, contrast=0.5, saturation=0.5, hue=0.1, p=0.5):
|
||||
self.p = p
|
||||
self.transforms = ColorJitter(
|
||||
brightness=brightness, contrast=contrast, saturation=saturation, hue=hue
|
||||
)
|
||||
|
||||
def __call__(self, img):
|
||||
if random.random() < self.p:
|
||||
return self.transforms(img)
|
||||
else:
|
||||
return img
|
||||
|
||||
|
||||
class SVTRDeterioration(object):
|
||||
def __init__(self, var, degrees, factor, p=0.5):
|
||||
self.p = p
|
||||
transforms = []
|
||||
if var is not None:
|
||||
transforms.append(CVGaussianNoise(var=var))
|
||||
if degrees is not None:
|
||||
transforms.append(CVMotionBlur(degrees=degrees))
|
||||
if factor is not None:
|
||||
transforms.append(CVRescale(factor=factor))
|
||||
self.transforms = transforms
|
||||
|
||||
def __call__(self, img):
|
||||
if random.random() < self.p:
|
||||
random.shuffle(self.transforms)
|
||||
transforms = Compose(self.transforms)
|
||||
return transforms(img)
|
||||
else:
|
||||
return img
|
||||
|
||||
|
||||
class ParseQDeterioration(object):
|
||||
def __init__(self, var, degrees, lam, radius, factor, p=0.5):
|
||||
self.p = p
|
||||
transforms = []
|
||||
if var is not None:
|
||||
transforms.append(CVGaussianNoise(var=var))
|
||||
if degrees is not None:
|
||||
transforms.append(CVMotionBlur(degrees=degrees))
|
||||
if lam is not None:
|
||||
transforms.append(CVPossionNoise(lam=lam))
|
||||
if radius is not None:
|
||||
transforms.append(CVGaussionBlur(radius=radius))
|
||||
if factor is not None:
|
||||
transforms.append(CVRescale(factor=factor))
|
||||
self.transforms = transforms
|
||||
|
||||
def __call__(self, img):
|
||||
if random.random() < self.p:
|
||||
random.shuffle(self.transforms)
|
||||
transforms = Compose(self.transforms)
|
||||
return transforms(img)
|
||||
else:
|
||||
return img
|
||||
|
||||
|
||||
class SVTRGeometry(object):
|
||||
def __init__(
|
||||
self,
|
||||
aug_type=0,
|
||||
degrees=15,
|
||||
translate=(0.3, 0.3),
|
||||
scale=(0.5, 2.0),
|
||||
shear=(45, 15),
|
||||
distortion=0.5,
|
||||
p=0.5,
|
||||
):
|
||||
self.aug_type = aug_type
|
||||
self.p = p
|
||||
self.transforms = []
|
||||
self.transforms.append(CVRandomRotation(degrees=degrees))
|
||||
self.transforms.append(
|
||||
CVRandomAffine(
|
||||
degrees=degrees, translate=translate, scale=scale, shear=shear
|
||||
)
|
||||
)
|
||||
self.transforms.append(CVRandomPerspective(distortion=distortion))
|
||||
|
||||
def __call__(self, img):
|
||||
if random.random() < self.p:
|
||||
if self.aug_type:
|
||||
random.shuffle(self.transforms)
|
||||
transforms = Compose(self.transforms[: random.randint(1, 3)])
|
||||
img = transforms(img)
|
||||
else:
|
||||
img = self.transforms[random.randint(0, 2)](img)
|
||||
return img
|
||||
else:
|
||||
return img
|
||||
179
ppocr/data/imaug/copy_paste.py
Normal file
179
ppocr/data/imaug/copy_paste.py
Normal file
@@ -0,0 +1,179 @@
|
||||
# copyright (c) 2021 PaddlePaddle Authors. All Rights Reserve.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
import copy
|
||||
import cv2
|
||||
import random
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
from shapely.geometry import Polygon
|
||||
|
||||
from ppocr.data.imaug.iaa_augment import IaaAugment
|
||||
from ppocr.data.imaug.random_crop_data import is_poly_outside_rect
|
||||
from tools.infer.utility import get_rotate_crop_image
|
||||
|
||||
|
||||
class CopyPaste(object):
|
||||
def __init__(self, objects_paste_ratio=0.2, limit_paste=True, **kwargs):
|
||||
self.ext_data_num = 1
|
||||
self.objects_paste_ratio = objects_paste_ratio
|
||||
self.limit_paste = limit_paste
|
||||
augmenter_args = [{"type": "Resize", "args": {"size": [0.5, 3]}}]
|
||||
self.aug = IaaAugment(augmenter_args)
|
||||
|
||||
def __call__(self, data):
|
||||
point_num = data["polys"].shape[1]
|
||||
src_img = data["image"]
|
||||
src_polys = data["polys"].tolist()
|
||||
src_texts = data["texts"]
|
||||
src_ignores = data["ignore_tags"].tolist()
|
||||
ext_data = data["ext_data"][0]
|
||||
ext_image = ext_data["image"]
|
||||
ext_polys = ext_data["polys"]
|
||||
ext_texts = ext_data["texts"]
|
||||
ext_ignores = ext_data["ignore_tags"]
|
||||
|
||||
indexes = [i for i in range(len(ext_ignores)) if not ext_ignores[i]]
|
||||
select_num = max(1, min(int(self.objects_paste_ratio * len(ext_polys)), 30))
|
||||
|
||||
random.shuffle(indexes)
|
||||
select_idxs = indexes[:select_num]
|
||||
select_polys = ext_polys[select_idxs]
|
||||
select_ignores = ext_ignores[select_idxs]
|
||||
|
||||
src_img = cv2.cvtColor(src_img, cv2.COLOR_BGR2RGB)
|
||||
ext_image = cv2.cvtColor(ext_image, cv2.COLOR_BGR2RGB)
|
||||
src_img = Image.fromarray(src_img).convert("RGBA")
|
||||
for idx, poly, tag in zip(select_idxs, select_polys, select_ignores):
|
||||
box_img = get_rotate_crop_image(ext_image, poly)
|
||||
|
||||
src_img, box = self.paste_img(src_img, box_img, src_polys)
|
||||
if box is not None:
|
||||
box = box.tolist()
|
||||
for _ in range(len(box), point_num):
|
||||
box.append(box[-1])
|
||||
src_polys.append(box)
|
||||
src_texts.append(ext_texts[idx])
|
||||
src_ignores.append(tag)
|
||||
src_img = cv2.cvtColor(np.array(src_img), cv2.COLOR_RGB2BGR)
|
||||
h, w = src_img.shape[:2]
|
||||
src_polys = np.array(src_polys)
|
||||
src_polys[:, :, 0] = np.clip(src_polys[:, :, 0], 0, w)
|
||||
src_polys[:, :, 1] = np.clip(src_polys[:, :, 1], 0, h)
|
||||
data["image"] = src_img
|
||||
data["polys"] = src_polys
|
||||
data["texts"] = src_texts
|
||||
data["ignore_tags"] = np.array(src_ignores)
|
||||
return data
|
||||
|
||||
def paste_img(self, src_img, box_img, src_polys):
|
||||
box_img_pil = Image.fromarray(box_img).convert("RGBA")
|
||||
src_w, src_h = src_img.size
|
||||
box_w, box_h = box_img_pil.size
|
||||
|
||||
angle = np.random.randint(0, 360)
|
||||
box = np.array([[[0, 0], [box_w, 0], [box_w, box_h], [0, box_h]]])
|
||||
box = rotate_bbox(box_img, box, angle)[0]
|
||||
box_img_pil = box_img_pil.rotate(angle, expand=1)
|
||||
box_w, box_h = box_img_pil.width, box_img_pil.height
|
||||
if src_w - box_w < 0 or src_h - box_h < 0:
|
||||
return src_img, None
|
||||
|
||||
paste_x, paste_y = self.select_coord(
|
||||
src_polys, box, src_w - box_w, src_h - box_h
|
||||
)
|
||||
if paste_x is None:
|
||||
return src_img, None
|
||||
box[:, 0] += paste_x
|
||||
box[:, 1] += paste_y
|
||||
r, g, b, A = box_img_pil.split()
|
||||
src_img.paste(box_img_pil, (paste_x, paste_y), mask=A)
|
||||
|
||||
return src_img, box
|
||||
|
||||
def select_coord(self, src_polys, box, endx, endy):
|
||||
if self.limit_paste:
|
||||
xmin, ymin, xmax, ymax = (
|
||||
box[:, 0].min(),
|
||||
box[:, 1].min(),
|
||||
box[:, 0].max(),
|
||||
box[:, 1].max(),
|
||||
)
|
||||
for _ in range(50):
|
||||
paste_x = random.randint(0, endx)
|
||||
paste_y = random.randint(0, endy)
|
||||
xmin1 = xmin + paste_x
|
||||
xmax1 = xmax + paste_x
|
||||
ymin1 = ymin + paste_y
|
||||
ymax1 = ymax + paste_y
|
||||
|
||||
num_poly_in_rect = 0
|
||||
for poly in src_polys:
|
||||
if not is_poly_outside_rect(
|
||||
poly, xmin1, ymin1, xmax1 - xmin1, ymax1 - ymin1
|
||||
):
|
||||
num_poly_in_rect += 1
|
||||
break
|
||||
if num_poly_in_rect == 0:
|
||||
return paste_x, paste_y
|
||||
return None, None
|
||||
else:
|
||||
paste_x = random.randint(0, endx)
|
||||
paste_y = random.randint(0, endy)
|
||||
return paste_x, paste_y
|
||||
|
||||
|
||||
def get_union(pD, pG):
|
||||
return Polygon(pD).union(Polygon(pG)).area
|
||||
|
||||
|
||||
def get_intersection_over_union(pD, pG):
|
||||
return get_intersection(pD, pG) / get_union(pD, pG)
|
||||
|
||||
|
||||
def get_intersection(pD, pG):
|
||||
return Polygon(pD).intersection(Polygon(pG)).area
|
||||
|
||||
|
||||
def rotate_bbox(img, text_polys, angle, scale=1):
|
||||
"""
|
||||
from https://github.com/WenmuZhou/DBNet.pytorch/blob/master/data_loader/modules/augment.py
|
||||
Args:
|
||||
img: np.ndarray
|
||||
text_polys: np.ndarray N*4*2
|
||||
angle: int
|
||||
scale: int
|
||||
|
||||
Returns:
|
||||
|
||||
"""
|
||||
w = img.shape[1]
|
||||
h = img.shape[0]
|
||||
|
||||
rangle = np.deg2rad(angle)
|
||||
nw = abs(np.sin(rangle) * h) + abs(np.cos(rangle) * w)
|
||||
nh = abs(np.cos(rangle) * h) + abs(np.sin(rangle) * w)
|
||||
rot_mat = cv2.getRotationMatrix2D((nw * 0.5, nh * 0.5), angle, scale)
|
||||
rot_move = np.dot(rot_mat, np.array([(nw - w) * 0.5, (nh - h) * 0.5, 0]))
|
||||
rot_mat[0, 2] += rot_move[0]
|
||||
rot_mat[1, 2] += rot_move[1]
|
||||
|
||||
# ---------------------- rotate box ----------------------
|
||||
rot_text_polys = list()
|
||||
for bbox in text_polys:
|
||||
point1 = np.dot(rot_mat, np.array([bbox[0, 0], bbox[0, 1], 1]))
|
||||
point2 = np.dot(rot_mat, np.array([bbox[1, 0], bbox[1, 1], 1]))
|
||||
point3 = np.dot(rot_mat, np.array([bbox[2, 0], bbox[2, 1], 1]))
|
||||
point4 = np.dot(rot_mat, np.array([bbox[3, 0], bbox[3, 1], 1]))
|
||||
rot_text_polys.append([point1, point2, point3, point4])
|
||||
return np.array(rot_text_polys, dtype=np.float32)
|
||||
376
ppocr/data/imaug/ct_process.py
Normal file
376
ppocr/data/imaug/ct_process.py
Normal file
@@ -0,0 +1,376 @@
|
||||
# copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import os
|
||||
import cv2
|
||||
import paddle
|
||||
import random
|
||||
import pyclipper
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
import paddle.vision.transforms as transforms
|
||||
|
||||
from ppocr.utils.utility import check_install
|
||||
|
||||
|
||||
class RandomScale:
|
||||
def __init__(self, short_size=640, **kwargs):
|
||||
self.short_size = short_size
|
||||
|
||||
def scale_aligned(self, img, scale):
|
||||
oh, ow = img.shape[0:2]
|
||||
h = int(oh * scale + 0.5)
|
||||
w = int(ow * scale + 0.5)
|
||||
if h % 32 != 0:
|
||||
h = h + (32 - h % 32)
|
||||
if w % 32 != 0:
|
||||
w = w + (32 - w % 32)
|
||||
img = cv2.resize(img, dsize=(w, h))
|
||||
factor_h = h / oh
|
||||
factor_w = w / ow
|
||||
return img, factor_h, factor_w
|
||||
|
||||
def __call__(self, data):
|
||||
img = data["image"]
|
||||
|
||||
h, w = img.shape[0:2]
|
||||
random_scale = np.array([0.7, 0.8, 0.9, 1.0, 1.1, 1.2, 1.3])
|
||||
scale = (np.random.choice(random_scale) * self.short_size) / min(h, w)
|
||||
img, factor_h, factor_w = self.scale_aligned(img, scale)
|
||||
|
||||
data["scale_factor"] = (factor_w, factor_h)
|
||||
data["image"] = img
|
||||
return data
|
||||
|
||||
|
||||
class MakeShrink:
|
||||
def __init__(self, kernel_scale=0.7, **kwargs):
|
||||
self.kernel_scale = kernel_scale
|
||||
|
||||
def dist(self, a, b):
|
||||
return np.linalg.norm((a - b), ord=2, axis=0)
|
||||
|
||||
def perimeter(self, bbox):
|
||||
peri = 0.0
|
||||
for i in range(bbox.shape[0]):
|
||||
peri += self.dist(bbox[i], bbox[(i + 1) % bbox.shape[0]])
|
||||
return peri
|
||||
|
||||
def shrink(self, bboxes, rate, max_shr=20):
|
||||
check_install("Polygon", "Polygon3")
|
||||
import Polygon as plg
|
||||
|
||||
rate = rate * rate
|
||||
shrinked_bboxes = []
|
||||
for bbox in bboxes:
|
||||
area = plg.Polygon(bbox).area()
|
||||
peri = self.perimeter(bbox)
|
||||
|
||||
try:
|
||||
pco = pyclipper.PyclipperOffset()
|
||||
pco.AddPath(bbox, pyclipper.JT_ROUND, pyclipper.ET_CLOSEDPOLYGON)
|
||||
offset = min(int(area * (1 - rate) / (peri + 0.001) + 0.5), max_shr)
|
||||
|
||||
shrinked_bbox = pco.Execute(-offset)
|
||||
if len(shrinked_bbox) == 0:
|
||||
shrinked_bboxes.append(bbox)
|
||||
continue
|
||||
|
||||
shrinked_bbox = np.array(shrinked_bbox[0])
|
||||
if shrinked_bbox.shape[0] <= 2:
|
||||
shrinked_bboxes.append(bbox)
|
||||
continue
|
||||
|
||||
shrinked_bboxes.append(shrinked_bbox)
|
||||
except Exception as e:
|
||||
shrinked_bboxes.append(bbox)
|
||||
|
||||
return shrinked_bboxes
|
||||
|
||||
def __call__(self, data):
|
||||
img = data["image"]
|
||||
bboxes = data["polys"]
|
||||
words = data["texts"]
|
||||
scale_factor = data["scale_factor"]
|
||||
|
||||
gt_instance = np.zeros(img.shape[0:2], dtype="uint8") # h,w
|
||||
training_mask = np.ones(img.shape[0:2], dtype="uint8")
|
||||
training_mask_distance = np.ones(img.shape[0:2], dtype="uint8")
|
||||
|
||||
for i in range(len(bboxes)):
|
||||
bboxes[i] = np.reshape(
|
||||
bboxes[i]
|
||||
* ([scale_factor[0], scale_factor[1]] * (bboxes[i].shape[0] // 2)),
|
||||
(bboxes[i].shape[0] // 2, 2),
|
||||
).astype("int32")
|
||||
|
||||
for i in range(len(bboxes)):
|
||||
# different value for different bbox
|
||||
cv2.drawContours(gt_instance, [bboxes[i]], -1, i + 1, -1)
|
||||
|
||||
# set training mask to 0
|
||||
cv2.drawContours(training_mask, [bboxes[i]], -1, 0, -1)
|
||||
|
||||
# for not accurate annotation, use training_mask_distance
|
||||
if words[i] == "###" or words[i] == "???":
|
||||
cv2.drawContours(training_mask_distance, [bboxes[i]], -1, 0, -1)
|
||||
|
||||
# make shrink
|
||||
gt_kernel_instance = np.zeros(img.shape[0:2], dtype="uint8")
|
||||
kernel_bboxes = self.shrink(bboxes, self.kernel_scale)
|
||||
for i in range(len(bboxes)):
|
||||
cv2.drawContours(gt_kernel_instance, [kernel_bboxes[i]], -1, i + 1, -1)
|
||||
|
||||
# for training mask, kernel and background= 1, box region=0
|
||||
if words[i] != "###" and words[i] != "???":
|
||||
cv2.drawContours(training_mask, [kernel_bboxes[i]], -1, 1, -1)
|
||||
|
||||
gt_kernel = gt_kernel_instance.copy()
|
||||
# for gt_kernel, kernel = 1
|
||||
gt_kernel[gt_kernel > 0] = 1
|
||||
|
||||
# shrink 2 times
|
||||
tmp1 = gt_kernel_instance.copy()
|
||||
erode_kernel = np.ones((3, 3), np.uint8)
|
||||
tmp1 = cv2.erode(tmp1, erode_kernel, iterations=1)
|
||||
tmp2 = tmp1.copy()
|
||||
tmp2 = cv2.erode(tmp2, erode_kernel, iterations=1)
|
||||
|
||||
# compute text region
|
||||
gt_kernel_inner = tmp1 - tmp2
|
||||
|
||||
# gt_instance: text instance, bg=0, diff word use diff value
|
||||
# training_mask: text instance mask, word=0，kernel and bg=1
|
||||
# gt_kernel_instance: text kernel instance, bg=0, diff word use diff value
|
||||
# gt_kernel: text_kernel, bg=0，diff word use same value
|
||||
# gt_kernel_inner: text kernel reference
|
||||
# training_mask_distance: word without anno = 0, else 1
|
||||
|
||||
data["image"] = [
|
||||
img,
|
||||
gt_instance,
|
||||
training_mask,
|
||||
gt_kernel_instance,
|
||||
gt_kernel,
|
||||
gt_kernel_inner,
|
||||
training_mask_distance,
|
||||
]
|
||||
return data
|
||||
|
||||
|
||||
class GroupRandomHorizontalFlip:
|
||||
def __init__(self, p=0.5, **kwargs):
|
||||
self.p = p
|
||||
|
||||
def __call__(self, data):
|
||||
imgs = data["image"]
|
||||
|
||||
if random.random() < self.p:
|
||||
for i in range(len(imgs)):
|
||||
imgs[i] = np.flip(imgs[i], axis=1).copy()
|
||||
data["image"] = imgs
|
||||
return data
|
||||
|
||||
|
||||
class GroupRandomRotate:
|
||||
def __init__(self, **kwargs):
|
||||
pass
|
||||
|
||||
def __call__(self, data):
|
||||
imgs = data["image"]
|
||||
|
||||
max_angle = 10
|
||||
angle = random.random() * 2 * max_angle - max_angle
|
||||
for i in range(len(imgs)):
|
||||
img = imgs[i]
|
||||
w, h = img.shape[:2]
|
||||
rotation_matrix = cv2.getRotationMatrix2D((h / 2, w / 2), angle, 1)
|
||||
img_rotation = cv2.warpAffine(
|
||||
img, rotation_matrix, (h, w), flags=cv2.INTER_NEAREST
|
||||
)
|
||||
imgs[i] = img_rotation
|
||||
|
||||
data["image"] = imgs
|
||||
return data
|
||||
|
||||
|
||||
class GroupRandomCropPadding:
|
||||
def __init__(self, target_size=(640, 640), **kwargs):
|
||||
self.target_size = target_size
|
||||
|
||||
def __call__(self, data):
|
||||
imgs = data["image"]
|
||||
|
||||
h, w = imgs[0].shape[0:2]
|
||||
t_w, t_h = self.target_size
|
||||
p_w, p_h = self.target_size
|
||||
if w == t_w and h == t_h:
|
||||
return data
|
||||
|
||||
t_h = t_h if t_h < h else h
|
||||
t_w = t_w if t_w < w else w
|
||||
|
||||
if random.random() > 3.0 / 8.0 and np.max(imgs[1]) > 0:
|
||||
# make sure to crop the text region
|
||||
tl = np.min(np.where(imgs[1] > 0), axis=1) - (t_h, t_w)
|
||||
tl[tl < 0] = 0
|
||||
br = np.max(np.where(imgs[1] > 0), axis=1) - (t_h, t_w)
|
||||
br[br < 0] = 0
|
||||
br[0] = min(br[0], h - t_h)
|
||||
br[1] = min(br[1], w - t_w)
|
||||
|
||||
i = random.randint(tl[0], br[0]) if tl[0] < br[0] else 0
|
||||
j = random.randint(tl[1], br[1]) if tl[1] < br[1] else 0
|
||||
else:
|
||||
i = random.randint(0, h - t_h) if h - t_h > 0 else 0
|
||||
j = random.randint(0, w - t_w) if w - t_w > 0 else 0
|
||||
|
||||
n_imgs = []
|
||||
for idx in range(len(imgs)):
|
||||
if len(imgs[idx].shape) == 3:
|
||||
s3_length = int(imgs[idx].shape[-1])
|
||||
img = imgs[idx][i : i + t_h, j : j + t_w, :]
|
||||
img_p = cv2.copyMakeBorder(
|
||||
img,
|
||||
0,
|
||||
p_h - t_h,
|
||||
0,
|
||||
p_w - t_w,
|
||||
borderType=cv2.BORDER_CONSTANT,
|
||||
value=tuple(0 for i in range(s3_length)),
|
||||
)
|
||||
else:
|
||||
img = imgs[idx][i : i + t_h, j : j + t_w]
|
||||
img_p = cv2.copyMakeBorder(
|
||||
img,
|
||||
0,
|
||||
p_h - t_h,
|
||||
0,
|
||||
p_w - t_w,
|
||||
borderType=cv2.BORDER_CONSTANT,
|
||||
value=(0,),
|
||||
)
|
||||
n_imgs.append(img_p)
|
||||
|
||||
data["image"] = n_imgs
|
||||
return data
|
||||
|
||||
|
||||
class MakeCentripetalShift:
|
||||
def __init__(self, **kwargs):
|
||||
pass
|
||||
|
||||
def jaccard(self, As, Bs):
|
||||
A = As.shape[0] # small
|
||||
B = Bs.shape[0] # large
|
||||
|
||||
dis = np.sqrt(
|
||||
np.sum(
|
||||
(
|
||||
As[:, np.newaxis, :].repeat(B, axis=1)
|
||||
- Bs[np.newaxis, :, :].repeat(A, axis=0)
|
||||
)
|
||||
** 2,
|
||||
axis=-1,
|
||||
)
|
||||
)
|
||||
|
||||
ind = np.argmin(dis, axis=-1)
|
||||
|
||||
return ind
|
||||
|
||||
def __call__(self, data):
|
||||
imgs = data["image"]
|
||||
|
||||
(
|
||||
img,
|
||||
gt_instance,
|
||||
training_mask,
|
||||
gt_kernel_instance,
|
||||
gt_kernel,
|
||||
gt_kernel_inner,
|
||||
training_mask_distance,
|
||||
) = (imgs[0], imgs[1], imgs[2], imgs[3], imgs[4], imgs[5], imgs[6])
|
||||
|
||||
max_instance = np.max(gt_instance) # num bbox
|
||||
|
||||
# make centripetal shift
|
||||
gt_distance = np.zeros((2, *img.shape[0:2]), dtype=np.float32)
|
||||
for i in range(1, max_instance + 1):
|
||||
# kernel_reference
|
||||
ind = gt_kernel_inner == i
|
||||
|
||||
if np.sum(ind) == 0:
|
||||
training_mask[gt_instance == i] = 0
|
||||
training_mask_distance[gt_instance == i] = 0
|
||||
continue
|
||||
|
||||
kpoints = (
|
||||
np.array(np.where(ind)).transpose((1, 0))[:, ::-1].astype("float32")
|
||||
)
|
||||
|
||||
ind = (gt_instance == i) * (gt_kernel_instance == 0)
|
||||
if np.sum(ind) == 0:
|
||||
continue
|
||||
pixels = np.where(ind)
|
||||
|
||||
points = np.array(pixels).transpose((1, 0))[:, ::-1].astype("float32")
|
||||
|
||||
bbox_ind = self.jaccard(points, kpoints)
|
||||
|
||||
offset_gt = kpoints[bbox_ind] - points
|
||||
|
||||
gt_distance[:, pixels[0], pixels[1]] = offset_gt.T * 0.1
|
||||
|
||||
img = Image.fromarray(img)
|
||||
img = img.convert("RGB")
|
||||
|
||||
data["image"] = img
|
||||
data["gt_kernel"] = gt_kernel.astype("int64")
|
||||
data["training_mask"] = training_mask.astype("int64")
|
||||
data["gt_instance"] = gt_instance.astype("int64")
|
||||
data["gt_kernel_instance"] = gt_kernel_instance.astype("int64")
|
||||
data["training_mask_distance"] = training_mask_distance.astype("int64")
|
||||
data["gt_distance"] = gt_distance.astype("float32")
|
||||
|
||||
return data
|
||||
|
||||
|
||||
class ScaleAlignedShort:
|
||||
def __init__(self, short_size=640, **kwargs):
|
||||
self.short_size = short_size
|
||||
|
||||
def __call__(self, data):
|
||||
img = data["image"]
|
||||
|
||||
org_img_shape = img.shape
|
||||
|
||||
h, w = img.shape[0:2]
|
||||
scale = self.short_size * 1.0 / min(h, w)
|
||||
h = int(h * scale + 0.5)
|
||||
w = int(w * scale + 0.5)
|
||||
if h % 32 != 0:
|
||||
h = h + (32 - h % 32)
|
||||
if w % 32 != 0:
|
||||
w = w + (32 - w % 32)
|
||||
img = cv2.resize(img, dsize=(w, h))
|
||||
|
||||
new_img_shape = img.shape
|
||||
img_shape = np.array(org_img_shape + new_img_shape)
|
||||
|
||||
data["shape"] = img_shape
|
||||
data["image"] = img
|
||||
|
||||
return data
|
||||
770
ppocr/data/imaug/drrg_targets.py
Normal file
770
ppocr/data/imaug/drrg_targets.py
Normal file
@@ -0,0 +1,770 @@
|
||||
# copyright (c) 2022 PaddlePaddle Authors. All Rights Reserve.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
"""
|
||||
This code is refer from:
|
||||
https://github.com/open-mmlab/mmocr/blob/main/mmocr/datasets/pipelines/textdet_targets/drrg_targets.py
|
||||
"""
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
from ppocr.utils.utility import check_install
|
||||
from numpy.linalg import norm
|
||||
|
||||
|
||||
class DRRGTargets(object):
|
||||
def __init__(
|
||||
self,
|
||||
orientation_thr=2.0,
|
||||
resample_step=8.0,
|
||||
num_min_comps=9,
|
||||
num_max_comps=600,
|
||||
min_width=8.0,
|
||||
max_width=24.0,
|
||||
center_region_shrink_ratio=0.3,
|
||||
comp_shrink_ratio=1.0,
|
||||
comp_w_h_ratio=0.3,
|
||||
text_comp_nms_thr=0.25,
|
||||
min_rand_half_height=8.0,
|
||||
max_rand_half_height=24.0,
|
||||
jitter_level=0.2,
|
||||
**kwargs,
|
||||
):
|
||||
super().__init__()
|
||||
self.orientation_thr = orientation_thr
|
||||
self.resample_step = resample_step
|
||||
self.num_max_comps = num_max_comps
|
||||
self.num_min_comps = num_min_comps
|
||||
self.min_width = min_width
|
||||
self.max_width = max_width
|
||||
self.center_region_shrink_ratio = center_region_shrink_ratio
|
||||
self.comp_shrink_ratio = comp_shrink_ratio
|
||||
self.comp_w_h_ratio = comp_w_h_ratio
|
||||
self.text_comp_nms_thr = text_comp_nms_thr
|
||||
self.min_rand_half_height = min_rand_half_height
|
||||
self.max_rand_half_height = max_rand_half_height
|
||||
self.jitter_level = jitter_level
|
||||
self.eps = 1e-8
|
||||
|
||||
def vector_angle(self, vec1, vec2):
|
||||
if vec1.ndim > 1:
|
||||
unit_vec1 = vec1 / (norm(vec1, axis=-1) + self.eps).reshape((-1, 1))
|
||||
else:
|
||||
unit_vec1 = vec1 / (norm(vec1, axis=-1) + self.eps)
|
||||
if vec2.ndim > 1:
|
||||
unit_vec2 = vec2 / (norm(vec2, axis=-1) + self.eps).reshape((-1, 1))
|
||||
else:
|
||||
unit_vec2 = vec2 / (norm(vec2, axis=-1) + self.eps)
|
||||
return np.arccos(np.clip(np.sum(unit_vec1 * unit_vec2, axis=-1), -1.0, 1.0))
|
||||
|
||||
def vector_slope(self, vec):
|
||||
assert len(vec) == 2
|
||||
return abs(vec[1] / (vec[0] + self.eps))
|
||||
|
||||
def vector_sin(self, vec):
|
||||
assert len(vec) == 2
|
||||
return vec[1] / (norm(vec) + self.eps)
|
||||
|
||||
def vector_cos(self, vec):
|
||||
assert len(vec) == 2
|
||||
return vec[0] / (norm(vec) + self.eps)
|
||||
|
||||
def find_head_tail(self, points, orientation_thr):
|
||||
assert points.ndim == 2
|
||||
assert points.shape[0] >= 4
|
||||
assert points.shape[1] == 2
|
||||
assert isinstance(orientation_thr, float)
|
||||
|
||||
if len(points) > 4:
|
||||
pad_points = np.vstack([points, points[0]])
|
||||
edge_vec = pad_points[1:] - pad_points[:-1]
|
||||
|
||||
theta_sum = []
|
||||
adjacent_vec_theta = []
|
||||
for i, edge_vec1 in enumerate(edge_vec):
|
||||
adjacent_ind = [x % len(edge_vec) for x in [i - 1, i + 1]]
|
||||
adjacent_edge_vec = edge_vec[adjacent_ind]
|
||||
temp_theta_sum = np.sum(self.vector_angle(edge_vec1, adjacent_edge_vec))
|
||||
temp_adjacent_theta = self.vector_angle(
|
||||
adjacent_edge_vec[0], adjacent_edge_vec[1]
|
||||
)
|
||||
theta_sum.append(temp_theta_sum)
|
||||
adjacent_vec_theta.append(temp_adjacent_theta)
|
||||
theta_sum_score = np.array(theta_sum) / np.pi
|
||||
adjacent_theta_score = np.array(adjacent_vec_theta) / np.pi
|
||||
poly_center = np.mean(points, axis=0)
|
||||
edge_dist = np.maximum(
|
||||
norm(pad_points[1:] - poly_center, axis=-1),
|
||||
norm(pad_points[:-1] - poly_center, axis=-1),
|
||||
)
|
||||
dist_score = edge_dist / (np.max(edge_dist) + self.eps)
|
||||
position_score = np.zeros(len(edge_vec))
|
||||
score = 0.5 * theta_sum_score + 0.15 * adjacent_theta_score
|
||||
score += 0.35 * dist_score
|
||||
if len(points) % 2 == 0:
|
||||
position_score[(len(score) // 2 - 1)] += 1
|
||||
position_score[-1] += 1
|
||||
score += 0.1 * position_score
|
||||
pad_score = np.concatenate([score, score])
|
||||
score_matrix = np.zeros((len(score), len(score) - 3))
|
||||
x = np.arange(len(score) - 3) / float(len(score) - 4)
|
||||
gaussian = (
|
||||
1.0
|
||||
/ (np.sqrt(2.0 * np.pi) * 0.5)
|
||||
* np.exp(-np.power((x - 0.5) / 0.5, 2.0) / 2)
|
||||
)
|
||||
gaussian = gaussian / np.max(gaussian)
|
||||
for i in range(len(score)):
|
||||
score_matrix[i, :] = (
|
||||
score[i]
|
||||
+ pad_score[(i + 2) : (i + len(score) - 1)] * gaussian * 0.3
|
||||
)
|
||||
|
||||
head_start, tail_increment = np.unravel_index(
|
||||
score_matrix.argmax(), score_matrix.shape
|
||||
)
|
||||
tail_start = (head_start + tail_increment + 2) % len(points)
|
||||
head_end = (head_start + 1) % len(points)
|
||||
tail_end = (tail_start + 1) % len(points)
|
||||
|
||||
if head_end > tail_end:
|
||||
head_start, tail_start = tail_start, head_start
|
||||
head_end, tail_end = tail_end, head_end
|
||||
head_inds = [head_start, head_end]
|
||||
tail_inds = [tail_start, tail_end]
|
||||
else:
|
||||
if self.vector_slope(points[1] - points[0]) + self.vector_slope(
|
||||
points[3] - points[2]
|
||||
) < self.vector_slope(points[2] - points[1]) + self.vector_slope(
|
||||
points[0] - points[3]
|
||||
):
|
||||
horizontal_edge_inds = [[0, 1], [2, 3]]
|
||||
vertical_edge_inds = [[3, 0], [1, 2]]
|
||||
else:
|
||||
horizontal_edge_inds = [[3, 0], [1, 2]]
|
||||
vertical_edge_inds = [[0, 1], [2, 3]]
|
||||
|
||||
vertical_len_sum = norm(
|
||||
points[vertical_edge_inds[0][0]] - points[vertical_edge_inds[0][1]]
|
||||
) + norm(
|
||||
points[vertical_edge_inds[1][0]] - points[vertical_edge_inds[1][1]]
|
||||
)
|
||||
horizontal_len_sum = norm(
|
||||
points[horizontal_edge_inds[0][0]] - points[horizontal_edge_inds[0][1]]
|
||||
) + norm(
|
||||
points[horizontal_edge_inds[1][0]] - points[horizontal_edge_inds[1][1]]
|
||||
)
|
||||
|
||||
if vertical_len_sum > horizontal_len_sum * orientation_thr:
|
||||
head_inds = horizontal_edge_inds[0]
|
||||
tail_inds = horizontal_edge_inds[1]
|
||||
else:
|
||||
head_inds = vertical_edge_inds[0]
|
||||
tail_inds = vertical_edge_inds[1]
|
||||
|
||||
return head_inds, tail_inds
|
||||
|
||||
def reorder_poly_edge(self, points):
|
||||
assert points.ndim == 2
|
||||
assert points.shape[0] >= 4
|
||||
assert points.shape[1] == 2
|
||||
|
||||
head_inds, tail_inds = self.find_head_tail(points, self.orientation_thr)
|
||||
head_edge, tail_edge = points[head_inds], points[tail_inds]
|
||||
|
||||
pad_points = np.vstack([points, points])
|
||||
if tail_inds[1] < 1:
|
||||
tail_inds[1] = len(points)
|
||||
sideline1 = pad_points[head_inds[1] : tail_inds[1]]
|
||||
sideline2 = pad_points[tail_inds[1] : (head_inds[1] + len(points))]
|
||||
sideline_mean_shift = np.mean(sideline1, axis=0) - np.mean(sideline2, axis=0)
|
||||
|
||||
if sideline_mean_shift[1] > 0:
|
||||
top_sideline, bot_sideline = sideline2, sideline1
|
||||
else:
|
||||
top_sideline, bot_sideline = sideline1, sideline2
|
||||
|
||||
return head_edge, tail_edge, top_sideline, bot_sideline
|
||||
|
||||
def cal_curve_length(self, line):
|
||||
assert line.ndim == 2
|
||||
assert len(line) >= 2
|
||||
|
||||
edges_length = np.sqrt(
|
||||
(line[1:, 0] - line[:-1, 0]) ** 2 + (line[1:, 1] - line[:-1, 1]) ** 2
|
||||
)
|
||||
total_length = np.sum(edges_length)
|
||||
return edges_length, total_length
|
||||
|
||||
def resample_line(self, line, n):
|
||||
assert line.ndim == 2
|
||||
assert line.shape[0] >= 2
|
||||
assert line.shape[1] == 2
|
||||
assert isinstance(n, int)
|
||||
assert n > 2
|
||||
|
||||
edges_length, total_length = self.cal_curve_length(line)
|
||||
t_org = np.insert(np.cumsum(edges_length), 0, 0)
|
||||
unit_t = total_length / (n - 1)
|
||||
t_equidistant = np.arange(1, n - 1, dtype=np.float32) * unit_t
|
||||
edge_ind = 0
|
||||
points = [line[0]]
|
||||
for t in t_equidistant:
|
||||
while edge_ind < len(edges_length) - 1 and t > t_org[edge_ind + 1]:
|
||||
edge_ind += 1
|
||||
t_l, t_r = t_org[edge_ind], t_org[edge_ind + 1]
|
||||
weight = np.array([t_r - t, t - t_l], dtype=np.float32) / (
|
||||
t_r - t_l + self.eps
|
||||
)
|
||||
p_coords = np.dot(weight, line[[edge_ind, edge_ind + 1]])
|
||||
points.append(p_coords)
|
||||
points.append(line[-1])
|
||||
resampled_line = np.vstack(points)
|
||||
|
||||
return resampled_line
|
||||
|
||||
def resample_sidelines(self, sideline1, sideline2, resample_step):
|
||||
assert sideline1.ndim == sideline2.ndim == 2
|
||||
assert sideline1.shape[1] == sideline2.shape[1] == 2
|
||||
assert sideline1.shape[0] >= 2
|
||||
assert sideline2.shape[0] >= 2
|
||||
assert isinstance(resample_step, float)
|
||||
|
||||
_, length1 = self.cal_curve_length(sideline1)
|
||||
_, length2 = self.cal_curve_length(sideline2)
|
||||
|
||||
avg_length = (length1 + length2) / 2
|
||||
resample_point_num = max(int(float(avg_length) / resample_step) + 1, 3)
|
||||
|
||||
resampled_line1 = self.resample_line(sideline1, resample_point_num)
|
||||
resampled_line2 = self.resample_line(sideline2, resample_point_num)
|
||||
|
||||
return resampled_line1, resampled_line2
|
||||
|
||||
def dist_point2line(self, point, line):
|
||||
assert isinstance(line, tuple)
|
||||
point1, point2 = line
|
||||
d = abs(np.cross(point2 - point1, point - point1)) / (
|
||||
norm(point2 - point1) + 1e-8
|
||||
)
|
||||
return d
|
||||
|
||||
def draw_center_region_maps(
|
||||
self,
|
||||
top_line,
|
||||
bot_line,
|
||||
center_line,
|
||||
center_region_mask,
|
||||
top_height_map,
|
||||
bot_height_map,
|
||||
sin_map,
|
||||
cos_map,
|
||||
region_shrink_ratio,
|
||||
):
|
||||
assert top_line.shape == bot_line.shape == center_line.shape
|
||||
assert (
|
||||
center_region_mask.shape
|
||||
== top_height_map.shape
|
||||
== bot_height_map.shape
|
||||
== sin_map.shape
|
||||
== cos_map.shape
|
||||
)
|
||||
assert isinstance(region_shrink_ratio, float)
|
||||
|
||||
h, w = center_region_mask.shape
|
||||
for i in range(0, len(center_line) - 1):
|
||||
top_mid_point = (top_line[i] + top_line[i + 1]) / 2
|
||||
bot_mid_point = (bot_line[i] + bot_line[i + 1]) / 2
|
||||
|
||||
sin_theta = self.vector_sin(top_mid_point - bot_mid_point)
|
||||
cos_theta = self.vector_cos(top_mid_point - bot_mid_point)
|
||||
|
||||
tl = center_line[i] + (top_line[i] - center_line[i]) * region_shrink_ratio
|
||||
tr = (
|
||||
center_line[i + 1]
|
||||
+ (top_line[i + 1] - center_line[i + 1]) * region_shrink_ratio
|
||||
)
|
||||
br = (
|
||||
center_line[i + 1]
|
||||
+ (bot_line[i + 1] - center_line[i + 1]) * region_shrink_ratio
|
||||
)
|
||||
bl = center_line[i] + (bot_line[i] - center_line[i]) * region_shrink_ratio
|
||||
current_center_box = np.vstack([tl, tr, br, bl]).astype(np.int32)
|
||||
|
||||
cv2.fillPoly(center_region_mask, [current_center_box], color=1)
|
||||
cv2.fillPoly(sin_map, [current_center_box], color=sin_theta)
|
||||
cv2.fillPoly(cos_map, [current_center_box], color=cos_theta)
|
||||
|
||||
current_center_box[:, 0] = np.clip(current_center_box[:, 0], 0, w - 1)
|
||||
current_center_box[:, 1] = np.clip(current_center_box[:, 1], 0, h - 1)
|
||||
min_coord = np.min(current_center_box, axis=0).astype(np.int32)
|
||||
max_coord = np.max(current_center_box, axis=0).astype(np.int32)
|
||||
current_center_box = current_center_box - min_coord
|
||||
box_sz = max_coord - min_coord + 1
|
||||
|
||||
center_box_mask = np.zeros((box_sz[1], box_sz[0]), dtype=np.uint8)
|
||||
cv2.fillPoly(center_box_mask, [current_center_box], color=1)
|
||||
|
||||
inds = np.argwhere(center_box_mask > 0)
|
||||
inds = inds + (min_coord[1], min_coord[0])
|
||||
inds_xy = np.fliplr(inds)
|
||||
top_height_map[(inds[:, 0], inds[:, 1])] = self.dist_point2line(
|
||||
inds_xy, (top_line[i], top_line[i + 1])
|
||||
)
|
||||
bot_height_map[(inds[:, 0], inds[:, 1])] = self.dist_point2line(
|
||||
inds_xy, (bot_line[i], bot_line[i + 1])
|
||||
)
|
||||
|
||||
def generate_center_mask_attrib_maps(self, img_size, text_polys):
|
||||
assert isinstance(img_size, tuple)
|
||||
|
||||
h, w = img_size
|
||||
|
||||
center_lines = []
|
||||
center_region_mask = np.zeros((h, w), np.uint8)
|
||||
top_height_map = np.zeros((h, w), dtype=np.float32)
|
||||
bot_height_map = np.zeros((h, w), dtype=np.float32)
|
||||
sin_map = np.zeros((h, w), dtype=np.float32)
|
||||
cos_map = np.zeros((h, w), dtype=np.float32)
|
||||
|
||||
for poly in text_polys:
|
||||
polygon_points = poly
|
||||
_, _, top_line, bot_line = self.reorder_poly_edge(polygon_points)
|
||||
resampled_top_line, resampled_bot_line = self.resample_sidelines(
|
||||
top_line, bot_line, self.resample_step
|
||||
)
|
||||
resampled_bot_line = resampled_bot_line[::-1]
|
||||
center_line = (resampled_top_line + resampled_bot_line) / 2
|
||||
|
||||
if self.vector_slope(center_line[-1] - center_line[0]) > 2:
|
||||
if (center_line[-1] - center_line[0])[1] < 0:
|
||||
center_line = center_line[::-1]
|
||||
resampled_top_line = resampled_top_line[::-1]
|
||||
resampled_bot_line = resampled_bot_line[::-1]
|
||||
else:
|
||||
if (center_line[-1] - center_line[0])[0] < 0:
|
||||
center_line = center_line[::-1]
|
||||
resampled_top_line = resampled_top_line[::-1]
|
||||
resampled_bot_line = resampled_bot_line[::-1]
|
||||
|
||||
line_head_shrink_len = (
|
||||
np.clip(
|
||||
(norm(top_line[0] - bot_line[0]) * self.comp_w_h_ratio),
|
||||
self.min_width,
|
||||
self.max_width,
|
||||
)
|
||||
/ 2
|
||||
)
|
||||
line_tail_shrink_len = (
|
||||
np.clip(
|
||||
(norm(top_line[-1] - bot_line[-1]) * self.comp_w_h_ratio),
|
||||
self.min_width,
|
||||
self.max_width,
|
||||
)
|
||||
/ 2
|
||||
)
|
||||
num_head_shrink = int(line_head_shrink_len // self.resample_step)
|
||||
num_tail_shrink = int(line_tail_shrink_len // self.resample_step)
|
||||
if len(center_line) > num_head_shrink + num_tail_shrink + 2:
|
||||
center_line = center_line[
|
||||
num_head_shrink : len(center_line) - num_tail_shrink
|
||||
]
|
||||
resampled_top_line = resampled_top_line[
|
||||
num_head_shrink : len(resampled_top_line) - num_tail_shrink
|
||||
]
|
||||
resampled_bot_line = resampled_bot_line[
|
||||
num_head_shrink : len(resampled_bot_line) - num_tail_shrink
|
||||
]
|
||||
center_lines.append(center_line.astype(np.int32))
|
||||
|
||||
self.draw_center_region_maps(
|
||||
resampled_top_line,
|
||||
resampled_bot_line,
|
||||
center_line,
|
||||
center_region_mask,
|
||||
top_height_map,
|
||||
bot_height_map,
|
||||
sin_map,
|
||||
cos_map,
|
||||
self.center_region_shrink_ratio,
|
||||
)
|
||||
|
||||
return (
|
||||
center_lines,
|
||||
center_region_mask,
|
||||
top_height_map,
|
||||
bot_height_map,
|
||||
sin_map,
|
||||
cos_map,
|
||||
)
|
||||
|
||||
def generate_rand_comp_attribs(self, num_rand_comps, center_sample_mask):
|
||||
assert isinstance(num_rand_comps, int)
|
||||
assert num_rand_comps > 0
|
||||
assert center_sample_mask.ndim == 2
|
||||
|
||||
h, w = center_sample_mask.shape
|
||||
|
||||
max_rand_half_height = self.max_rand_half_height
|
||||
min_rand_half_height = self.min_rand_half_height
|
||||
max_rand_height = max_rand_half_height * 2
|
||||
max_rand_width = np.clip(
|
||||
max_rand_height * self.comp_w_h_ratio, self.min_width, self.max_width
|
||||
)
|
||||
margin = (
|
||||
int(np.sqrt((max_rand_height / 2) ** 2 + (max_rand_width / 2) ** 2)) + 1
|
||||
)
|
||||
|
||||
if 2 * margin + 1 > min(h, w):
|
||||
assert min(h, w) > (np.sqrt(2) * (self.min_width + 1))
|
||||
max_rand_half_height = max(min(h, w) / 4, self.min_width / 2 + 1)
|
||||
min_rand_half_height = max(max_rand_half_height / 4, self.min_width / 2)
|
||||
|
||||
max_rand_height = max_rand_half_height * 2
|
||||
max_rand_width = np.clip(
|
||||
max_rand_height * self.comp_w_h_ratio, self.min_width, self.max_width
|
||||
)
|
||||
margin = (
|
||||
int(np.sqrt((max_rand_height / 2) ** 2 + (max_rand_width / 2) ** 2)) + 1
|
||||
)
|
||||
|
||||
inner_center_sample_mask = np.zeros_like(center_sample_mask)
|
||||
inner_center_sample_mask[margin : h - margin, margin : w - margin] = (
|
||||
center_sample_mask[margin : h - margin, margin : w - margin]
|
||||
)
|
||||
kernel_size = int(np.clip(max_rand_half_height, 7, 21))
|
||||
inner_center_sample_mask = cv2.erode(
|
||||
inner_center_sample_mask, np.ones((kernel_size, kernel_size), np.uint8)
|
||||
)
|
||||
|
||||
center_candidates = np.argwhere(inner_center_sample_mask > 0)
|
||||
num_center_candidates = len(center_candidates)
|
||||
sample_inds = np.random.choice(num_center_candidates, num_rand_comps)
|
||||
rand_centers = center_candidates[sample_inds]
|
||||
|
||||
rand_top_height = np.random.randint(
|
||||
min_rand_half_height, max_rand_half_height, size=(len(rand_centers), 1)
|
||||
)
|
||||
rand_bot_height = np.random.randint(
|
||||
min_rand_half_height, max_rand_half_height, size=(len(rand_centers), 1)
|
||||
)
|
||||
|
||||
rand_cos = 2 * np.random.random(size=(len(rand_centers), 1)) - 1
|
||||
rand_sin = 2 * np.random.random(size=(len(rand_centers), 1)) - 1
|
||||
scale = np.sqrt(1.0 / (rand_cos**2 + rand_sin**2 + 1e-8))
|
||||
rand_cos = rand_cos * scale
|
||||
rand_sin = rand_sin * scale
|
||||
|
||||
height = rand_top_height + rand_bot_height
|
||||
width = np.clip(height * self.comp_w_h_ratio, self.min_width, self.max_width)
|
||||
|
||||
rand_comp_attribs = np.hstack(
|
||||
[
|
||||
rand_centers[:, ::-1],
|
||||
height,
|
||||
width,
|
||||
rand_cos,
|
||||
rand_sin,
|
||||
np.zeros_like(rand_sin),
|
||||
]
|
||||
).astype(np.float32)
|
||||
|
||||
return rand_comp_attribs
|
||||
|
||||
def jitter_comp_attribs(self, comp_attribs, jitter_level):
|
||||
"""Jitter text components attributes.
|
||||
|
||||
Args:
|
||||
comp_attribs (ndarray): The text component attributes.
|
||||
jitter_level (float): The jitter level of text components
|
||||
attributes.
|
||||
|
||||
Returns:
|
||||
jittered_comp_attribs (ndarray): The jittered text component
|
||||
attributes (x, y, h, w, cos, sin, comp_label).
|
||||
"""
|
||||
|
||||
assert comp_attribs.shape[1] == 7
|
||||
assert comp_attribs.shape[0] > 0
|
||||
assert isinstance(jitter_level, float)
|
||||
|
||||
x = comp_attribs[:, 0].reshape((-1, 1))
|
||||
y = comp_attribs[:, 1].reshape((-1, 1))
|
||||
h = comp_attribs[:, 2].reshape((-1, 1))
|
||||
w = comp_attribs[:, 3].reshape((-1, 1))
|
||||
cos = comp_attribs[:, 4].reshape((-1, 1))
|
||||
sin = comp_attribs[:, 5].reshape((-1, 1))
|
||||
comp_labels = comp_attribs[:, 6].reshape((-1, 1))
|
||||
|
||||
x += (
|
||||
(np.random.random(size=(len(comp_attribs), 1)) - 0.5)
|
||||
* (h * np.abs(cos) + w * np.abs(sin))
|
||||
* jitter_level
|
||||
)
|
||||
y += (
|
||||
(np.random.random(size=(len(comp_attribs), 1)) - 0.5)
|
||||
* (h * np.abs(sin) + w * np.abs(cos))
|
||||
* jitter_level
|
||||
)
|
||||
|
||||
h += (np.random.random(size=(len(comp_attribs), 1)) - 0.5) * h * jitter_level
|
||||
w += (np.random.random(size=(len(comp_attribs), 1)) - 0.5) * w * jitter_level
|
||||
|
||||
cos += (np.random.random(size=(len(comp_attribs), 1)) - 0.5) * 2 * jitter_level
|
||||
sin += (np.random.random(size=(len(comp_attribs), 1)) - 0.5) * 2 * jitter_level
|
||||
|
||||
scale = np.sqrt(1.0 / (cos**2 + sin**2 + 1e-8))
|
||||
cos = cos * scale
|
||||
sin = sin * scale
|
||||
|
||||
jittered_comp_attribs = np.hstack([x, y, h, w, cos, sin, comp_labels])
|
||||
|
||||
return jittered_comp_attribs
|
||||
|
||||
def generate_comp_attribs(
|
||||
self,
|
||||
center_lines,
|
||||
text_mask,
|
||||
center_region_mask,
|
||||
top_height_map,
|
||||
bot_height_map,
|
||||
sin_map,
|
||||
cos_map,
|
||||
):
|
||||
"""Generate text component attributes.
|
||||
|
||||
Args:
|
||||
center_lines (list[ndarray]): The list of text center lines .
|
||||
text_mask (ndarray): The text region mask.
|
||||
center_region_mask (ndarray): The text center region mask.
|
||||
top_height_map (ndarray): The map on which the distance from points
|
||||
to top side lines will be drawn for each pixel in text center
|
||||
regions.
|
||||
bot_height_map (ndarray): The map on which the distance from points
|
||||
to bottom side lines will be drawn for each pixel in text
|
||||
center regions.
|
||||
sin_map (ndarray): The sin(theta) map where theta is the angle
|
||||
between vector (top point - bottom point) and vector (1, 0).
|
||||
cos_map (ndarray): The cos(theta) map where theta is the angle
|
||||
between vector (top point - bottom point) and vector (1, 0).
|
||||
|
||||
Returns:
|
||||
pad_comp_attribs (ndarray): The padded text component attributes
|
||||
of a fixed size.
|
||||
"""
|
||||
|
||||
assert isinstance(center_lines, list)
|
||||
assert (
|
||||
text_mask.shape
|
||||
== center_region_mask.shape
|
||||
== top_height_map.shape
|
||||
== bot_height_map.shape
|
||||
== sin_map.shape
|
||||
== cos_map.shape
|
||||
)
|
||||
|
||||
center_lines_mask = np.zeros_like(center_region_mask)
|
||||
cv2.polylines(center_lines_mask, center_lines, 0, 1, 1)
|
||||
center_lines_mask = center_lines_mask * center_region_mask
|
||||
comp_centers = np.argwhere(center_lines_mask > 0)
|
||||
|
||||
y = comp_centers[:, 0]
|
||||
x = comp_centers[:, 1]
|
||||
|
||||
top_height = top_height_map[y, x].reshape((-1, 1)) * self.comp_shrink_ratio
|
||||
bot_height = bot_height_map[y, x].reshape((-1, 1)) * self.comp_shrink_ratio
|
||||
sin = sin_map[y, x].reshape((-1, 1))
|
||||
cos = cos_map[y, x].reshape((-1, 1))
|
||||
|
||||
top_mid_points = comp_centers + np.hstack([top_height * sin, top_height * cos])
|
||||
bot_mid_points = comp_centers - np.hstack([bot_height * sin, bot_height * cos])
|
||||
|
||||
width = (top_height + bot_height) * self.comp_w_h_ratio
|
||||
width = np.clip(width, self.min_width, self.max_width)
|
||||
r = width / 2
|
||||
|
||||
tl = top_mid_points[:, ::-1] - np.hstack([-r * sin, r * cos])
|
||||
tr = top_mid_points[:, ::-1] + np.hstack([-r * sin, r * cos])
|
||||
br = bot_mid_points[:, ::-1] + np.hstack([-r * sin, r * cos])
|
||||
bl = bot_mid_points[:, ::-1] - np.hstack([-r * sin, r * cos])
|
||||
text_comps = np.hstack([tl, tr, br, bl]).astype(np.float32)
|
||||
|
||||
score = np.ones((text_comps.shape[0], 1), dtype=np.float32)
|
||||
text_comps = np.hstack([text_comps, score])
|
||||
check_install("lanms", "lanms-neo")
|
||||
from lanms import merge_quadrangle_n9 as la_nms
|
||||
|
||||
text_comps = la_nms(text_comps, self.text_comp_nms_thr)
|
||||
|
||||
if text_comps.shape[0] >= 1:
|
||||
img_h, img_w = center_region_mask.shape
|
||||
text_comps[:, 0:8:2] = np.clip(text_comps[:, 0:8:2], 0, img_w - 1)
|
||||
text_comps[:, 1:8:2] = np.clip(text_comps[:, 1:8:2], 0, img_h - 1)
|
||||
|
||||
comp_centers = np.mean(
|
||||
text_comps[:, 0:8].reshape((-1, 4, 2)), axis=1
|
||||
).astype(np.int32)
|
||||
x = comp_centers[:, 0]
|
||||
y = comp_centers[:, 1]
|
||||
|
||||
height = (top_height_map[y, x] + bot_height_map[y, x]).reshape((-1, 1))
|
||||
width = np.clip(
|
||||
height * self.comp_w_h_ratio, self.min_width, self.max_width
|
||||
)
|
||||
|
||||
cos = cos_map[y, x].reshape((-1, 1))
|
||||
sin = sin_map[y, x].reshape((-1, 1))
|
||||
|
||||
_, comp_label_mask = cv2.connectedComponents(
|
||||
center_region_mask, connectivity=8
|
||||
)
|
||||
comp_labels = comp_label_mask[y, x].reshape((-1, 1)).astype(np.float32)
|
||||
|
||||
x = x.reshape((-1, 1)).astype(np.float32)
|
||||
y = y.reshape((-1, 1)).astype(np.float32)
|
||||
comp_attribs = np.hstack([x, y, height, width, cos, sin, comp_labels])
|
||||
comp_attribs = self.jitter_comp_attribs(comp_attribs, self.jitter_level)
|
||||
|
||||
if comp_attribs.shape[0] < self.num_min_comps:
|
||||
num_rand_comps = self.num_min_comps - comp_attribs.shape[0]
|
||||
rand_comp_attribs = self.generate_rand_comp_attribs(
|
||||
num_rand_comps, 1 - text_mask
|
||||
)
|
||||
comp_attribs = np.vstack([comp_attribs, rand_comp_attribs])
|
||||
else:
|
||||
comp_attribs = self.generate_rand_comp_attribs(
|
||||
self.num_min_comps, 1 - text_mask
|
||||
)
|
||||
|
||||
num_comps = (
|
||||
np.ones((comp_attribs.shape[0], 1), dtype=np.float32)
|
||||
* comp_attribs.shape[0]
|
||||
)
|
||||
comp_attribs = np.hstack([num_comps, comp_attribs])
|
||||
|
||||
if comp_attribs.shape[0] > self.num_max_comps:
|
||||
comp_attribs = comp_attribs[: self.num_max_comps, :]
|
||||
comp_attribs[:, 0] = self.num_max_comps
|
||||
|
||||
pad_comp_attribs = np.zeros(
|
||||
(self.num_max_comps, comp_attribs.shape[1]), dtype=np.float32
|
||||
)
|
||||
pad_comp_attribs[: comp_attribs.shape[0], :] = comp_attribs
|
||||
|
||||
return pad_comp_attribs
|
||||
|
||||
def generate_text_region_mask(self, img_size, text_polys):
|
||||
"""Generate text center region mask and geometry attribute maps.
|
||||
|
||||
Args:
|
||||
img_size (tuple): The image size (height, width).
|
||||
text_polys (list[list[ndarray]]): The list of text polygons.
|
||||
|
||||
Returns:
|
||||
text_region_mask (ndarray): The text region mask.
|
||||
"""
|
||||
|
||||
assert isinstance(img_size, tuple)
|
||||
|
||||
h, w = img_size
|
||||
text_region_mask = np.zeros((h, w), dtype=np.uint8)
|
||||
|
||||
for poly in text_polys:
|
||||
polygon = np.array(poly, dtype=np.int32).reshape((1, -1, 2))
|
||||
cv2.fillPoly(text_region_mask, polygon, 1)
|
||||
|
||||
return text_region_mask
|
||||
|
||||
def generate_effective_mask(self, mask_size: tuple, polygons_ignore):
|
||||
"""Generate effective mask by setting the ineffective regions to 0 and
|
||||
effective regions to 1.
|
||||
|
||||
Args:
|
||||
mask_size (tuple): The mask size.
|
||||
polygons_ignore (list[[ndarray]]: The list of ignored text
|
||||
polygons.
|
||||
|
||||
Returns:
|
||||
mask (ndarray): The effective mask of (height, width).
|
||||
"""
|
||||
mask = np.ones(mask_size, dtype=np.uint8)
|
||||
|
||||
for poly in polygons_ignore:
|
||||
instance = poly.astype(np.int32).reshape(1, -1, 2)
|
||||
cv2.fillPoly(mask, instance, 0)
|
||||
|
||||
return mask
|
||||
|
||||
def generate_targets(self, data):
|
||||
"""Generate the gt targets for DRRG.
|
||||
|
||||
Args:
|
||||
data (dict): The input result dictionary.
|
||||
|
||||
Returns:
|
||||
data (dict): The output result dictionary.
|
||||
"""
|
||||
|
||||
assert isinstance(data, dict)
|
||||
|
||||
image = data["image"]
|
||||
polygons = data["polys"]
|
||||
ignore_tags = data["ignore_tags"]
|
||||
h, w, _ = image.shape
|
||||
|
||||
polygon_masks = []
|
||||
polygon_masks_ignore = []
|
||||
for tag, polygon in zip(ignore_tags, polygons):
|
||||
if tag is True:
|
||||
polygon_masks_ignore.append(polygon)
|
||||
else:
|
||||
polygon_masks.append(polygon)
|
||||
|
||||
gt_text_mask = self.generate_text_region_mask((h, w), polygon_masks)
|
||||
gt_mask = self.generate_effective_mask((h, w), polygon_masks_ignore)
|
||||
(
|
||||
center_lines,
|
||||
gt_center_region_mask,
|
||||
gt_top_height_map,
|
||||
gt_bot_height_map,
|
||||
gt_sin_map,
|
||||
gt_cos_map,
|
||||
) = self.generate_center_mask_attrib_maps((h, w), polygon_masks)
|
||||
|
||||
gt_comp_attribs = self.generate_comp_attribs(
|
||||
center_lines,
|
||||
gt_text_mask,
|
||||
gt_center_region_mask,
|
||||
gt_top_height_map,
|
||||
gt_bot_height_map,
|
||||
gt_sin_map,
|
||||
gt_cos_map,
|
||||
)
|
||||
|
||||
mapping = {
|
||||
"gt_text_mask": gt_text_mask,
|
||||
"gt_center_region_mask": gt_center_region_mask,
|
||||
"gt_mask": gt_mask,
|
||||
"gt_top_height_map": gt_top_height_map,
|
||||
"gt_bot_height_map": gt_bot_height_map,
|
||||
"gt_sin_map": gt_sin_map,
|
||||
"gt_cos_map": gt_cos_map,
|
||||
}
|
||||
|
||||
data.update(mapping)
|
||||
data["gt_comp_attribs"] = gt_comp_attribs
|
||||
return data
|
||||
|
||||
def __call__(self, data):
|
||||
data = self.generate_targets(data)
|
||||
return data
|
||||
446
ppocr/data/imaug/east_process.py
Normal file
446
ppocr/data/imaug/east_process.py
Normal file
@@ -0,0 +1,446 @@
|
||||
# copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
"""
|
||||
This code is referred from:
|
||||
https://github.com/songdejia/EAST/blob/master/data_utils.py
|
||||
"""
|
||||
import math
|
||||
import cv2
|
||||
import numpy as np
|
||||
import json
|
||||
import sys
|
||||
import os
|
||||
|
||||
__all__ = ["EASTProcessTrain"]
|
||||
|
||||
|
||||
class EASTProcessTrain(object):
|
||||
def __init__(
|
||||
self,
|
||||
image_shape=[512, 512],
|
||||
background_ratio=0.125,
|
||||
min_crop_side_ratio=0.1,
|
||||
min_text_size=10,
|
||||
**kwargs,
|
||||
):
|
||||
self.input_size = image_shape[1]
|
||||
self.random_scale = np.array([0.5, 1, 2.0, 3.0])
|
||||
self.background_ratio = background_ratio
|
||||
self.min_crop_side_ratio = min_crop_side_ratio
|
||||
self.min_text_size = min_text_size
|
||||
|
||||
def preprocess(self, im):
|
||||
input_size = self.input_size
|
||||
im_shape = im.shape
|
||||
im_size_min = np.min(im_shape[0:2])
|
||||
im_size_max = np.max(im_shape[0:2])
|
||||
im_scale = float(input_size) / float(im_size_max)
|
||||
im = cv2.resize(im, None, None, fx=im_scale, fy=im_scale)
|
||||
img_mean = [0.485, 0.456, 0.406]
|
||||
img_std = [0.229, 0.224, 0.225]
|
||||
# im = im[:, :, ::-1].astype(np.float32)
|
||||
im = im / 255
|
||||
im -= img_mean
|
||||
im /= img_std
|
||||
new_h, new_w, _ = im.shape
|
||||
im_padded = np.zeros((input_size, input_size, 3), dtype=np.float32)
|
||||
im_padded[:new_h, :new_w, :] = im
|
||||
im_padded = im_padded.transpose((2, 0, 1))
|
||||
im_padded = im_padded[np.newaxis, :]
|
||||
return im_padded, im_scale
|
||||
|
||||
def rotate_im_poly(self, im, text_polys):
|
||||
"""
|
||||
rotate image with 90 / 180 / 270 degre
|
||||
"""
|
||||
im_w, im_h = im.shape[1], im.shape[0]
|
||||
dst_im = im.copy()
|
||||
dst_polys = []
|
||||
rand_degree_ratio = np.random.rand()
|
||||
rand_degree_cnt = 1
|
||||
if 0.333 < rand_degree_ratio < 0.666:
|
||||
rand_degree_cnt = 2
|
||||
elif rand_degree_ratio > 0.666:
|
||||
rand_degree_cnt = 3
|
||||
for i in range(rand_degree_cnt):
|
||||
dst_im = np.rot90(dst_im)
|
||||
rot_degree = -90 * rand_degree_cnt
|
||||
rot_angle = rot_degree * math.pi / 180.0
|
||||
n_poly = text_polys.shape[0]
|
||||
cx, cy = 0.5 * im_w, 0.5 * im_h
|
||||
ncx, ncy = 0.5 * dst_im.shape[1], 0.5 * dst_im.shape[0]
|
||||
for i in range(n_poly):
|
||||
wordBB = text_polys[i]
|
||||
poly = []
|
||||
for j in range(4):
|
||||
sx, sy = wordBB[j][0], wordBB[j][1]
|
||||
dx = (
|
||||
math.cos(rot_angle) * (sx - cx)
|
||||
- math.sin(rot_angle) * (sy - cy)
|
||||
+ ncx
|
||||
)
|
||||
dy = (
|
||||
math.sin(rot_angle) * (sx - cx)
|
||||
+ math.cos(rot_angle) * (sy - cy)
|
||||
+ ncy
|
||||
)
|
||||
poly.append([dx, dy])
|
||||
dst_polys.append(poly)
|
||||
dst_polys = np.array(dst_polys, dtype=np.float32)
|
||||
return dst_im, dst_polys
|
||||
|
||||
def polygon_area(self, poly):
|
||||
"""
|
||||
compute area of a polygon
|
||||
:param poly:
|
||||
:return:
|
||||
"""
|
||||
edge = [
|
||||
(poly[1][0] - poly[0][0]) * (poly[1][1] + poly[0][1]),
|
||||
(poly[2][0] - poly[1][0]) * (poly[2][1] + poly[1][1]),
|
||||
(poly[3][0] - poly[2][0]) * (poly[3][1] + poly[2][1]),
|
||||
(poly[0][0] - poly[3][0]) * (poly[0][1] + poly[3][1]),
|
||||
]
|
||||
return np.sum(edge) / 2.0
|
||||
|
||||
def check_and_validate_polys(self, polys, tags, img_height, img_width):
|
||||
"""
|
||||
check so that the text poly is in the same direction,
|
||||
and also filter some invalid polygons
|
||||
:param polys:
|
||||
:param tags:
|
||||
:return:
|
||||
"""
|
||||
h, w = img_height, img_width
|
||||
if polys.shape[0] == 0:
|
||||
return polys
|
||||
polys[:, :, 0] = np.clip(polys[:, :, 0], 0, w - 1)
|
||||
polys[:, :, 1] = np.clip(polys[:, :, 1], 0, h - 1)
|
||||
|
||||
validated_polys = []
|
||||
validated_tags = []
|
||||
for poly, tag in zip(polys, tags):
|
||||
p_area = self.polygon_area(poly)
|
||||
# invalid poly
|
||||
if abs(p_area) < 1:
|
||||
continue
|
||||
if p_area > 0:
|
||||
#'poly in wrong direction'
|
||||
if not tag:
|
||||
tag = True # reversed cases should be ignore
|
||||
poly = poly[(0, 3, 2, 1), :]
|
||||
validated_polys.append(poly)
|
||||
validated_tags.append(tag)
|
||||
return np.array(validated_polys), np.array(validated_tags)
|
||||
|
||||
def draw_img_polys(self, img, polys):
|
||||
if len(img.shape) == 4:
|
||||
img = np.squeeze(img, axis=0)
|
||||
if img.shape[0] == 3:
|
||||
img = img.transpose((1, 2, 0))
|
||||
img[:, :, 2] += 123.68
|
||||
img[:, :, 1] += 116.78
|
||||
img[:, :, 0] += 103.94
|
||||
cv2.imwrite("tmp.jpg", img)
|
||||
img = cv2.imread("tmp.jpg")
|
||||
for box in polys:
|
||||
box = box.astype(np.int32).reshape((-1, 1, 2))
|
||||
cv2.polylines(img, [box], True, color=(255, 255, 0), thickness=2)
|
||||
import random
|
||||
|
||||
ino = random.randint(0, 100)
|
||||
cv2.imwrite("tmp_%d.jpg" % ino, img)
|
||||
return
|
||||
|
||||
def shrink_poly(self, poly, r):
|
||||
"""
|
||||
fit a poly inside the origin poly, maybe bugs here...
|
||||
used for generate the score map
|
||||
:param poly: the text poly
|
||||
:param r: r in the paper
|
||||
:return: the shrunk poly
|
||||
"""
|
||||
# shrink ratio
|
||||
R = 0.3
|
||||
# find the longer pair
|
||||
dist0 = np.linalg.norm(poly[0] - poly[1])
|
||||
dist1 = np.linalg.norm(poly[2] - poly[3])
|
||||
dist2 = np.linalg.norm(poly[0] - poly[3])
|
||||
dist3 = np.linalg.norm(poly[1] - poly[2])
|
||||
if dist0 + dist1 > dist2 + dist3:
|
||||
# first move (p0, p1), (p2, p3), then (p0, p3), (p1, p2)
|
||||
## p0, p1
|
||||
theta = np.arctan2((poly[1][1] - poly[0][1]), (poly[1][0] - poly[0][0]))
|
||||
poly[0][0] += R * r[0] * np.cos(theta)
|
||||
poly[0][1] += R * r[0] * np.sin(theta)
|
||||
poly[1][0] -= R * r[1] * np.cos(theta)
|
||||
poly[1][1] -= R * r[1] * np.sin(theta)
|
||||
## p2, p3
|
||||
theta = np.arctan2((poly[2][1] - poly[3][1]), (poly[2][0] - poly[3][0]))
|
||||
poly[3][0] += R * r[3] * np.cos(theta)
|
||||
poly[3][1] += R * r[3] * np.sin(theta)
|
||||
poly[2][0] -= R * r[2] * np.cos(theta)
|
||||
poly[2][1] -= R * r[2] * np.sin(theta)
|
||||
## p0, p3
|
||||
theta = np.arctan2((poly[3][0] - poly[0][0]), (poly[3][1] - poly[0][1]))
|
||||
poly[0][0] += R * r[0] * np.sin(theta)
|
||||
poly[0][1] += R * r[0] * np.cos(theta)
|
||||
poly[3][0] -= R * r[3] * np.sin(theta)
|
||||
poly[3][1] -= R * r[3] * np.cos(theta)
|
||||
## p1, p2
|
||||
theta = np.arctan2((poly[2][0] - poly[1][0]), (poly[2][1] - poly[1][1]))
|
||||
poly[1][0] += R * r[1] * np.sin(theta)
|
||||
poly[1][1] += R * r[1] * np.cos(theta)
|
||||
poly[2][0] -= R * r[2] * np.sin(theta)
|
||||
poly[2][1] -= R * r[2] * np.cos(theta)
|
||||
else:
|
||||
## p0, p3
|
||||
# print poly
|
||||
theta = np.arctan2((poly[3][0] - poly[0][0]), (poly[3][1] - poly[0][1]))
|
||||
poly[0][0] += R * r[0] * np.sin(theta)
|
||||
poly[0][1] += R * r[0] * np.cos(theta)
|
||||
poly[3][0] -= R * r[3] * np.sin(theta)
|
||||
poly[3][1] -= R * r[3] * np.cos(theta)
|
||||
## p1, p2
|
||||
theta = np.arctan2((poly[2][0] - poly[1][0]), (poly[2][1] - poly[1][1]))
|
||||
poly[1][0] += R * r[1] * np.sin(theta)
|
||||
poly[1][1] += R * r[1] * np.cos(theta)
|
||||
poly[2][0] -= R * r[2] * np.sin(theta)
|
||||
poly[2][1] -= R * r[2] * np.cos(theta)
|
||||
## p0, p1
|
||||
theta = np.arctan2((poly[1][1] - poly[0][1]), (poly[1][0] - poly[0][0]))
|
||||
poly[0][0] += R * r[0] * np.cos(theta)
|
||||
poly[0][1] += R * r[0] * np.sin(theta)
|
||||
poly[1][0] -= R * r[1] * np.cos(theta)
|
||||
poly[1][1] -= R * r[1] * np.sin(theta)
|
||||
## p2, p3
|
||||
theta = np.arctan2((poly[2][1] - poly[3][1]), (poly[2][0] - poly[3][0]))
|
||||
poly[3][0] += R * r[3] * np.cos(theta)
|
||||
poly[3][1] += R * r[3] * np.sin(theta)
|
||||
poly[2][0] -= R * r[2] * np.cos(theta)
|
||||
poly[2][1] -= R * r[2] * np.sin(theta)
|
||||
return poly
|
||||
|
||||
def generate_quad(self, im_size, polys, tags):
|
||||
"""
|
||||
Generate quadrangle.
|
||||
"""
|
||||
h, w = im_size
|
||||
poly_mask = np.zeros((h, w), dtype=np.uint8)
|
||||
score_map = np.zeros((h, w), dtype=np.uint8)
|
||||
# (x1, y1, ..., x4, y4, short_edge_norm)
|
||||
geo_map = np.zeros((h, w, 9), dtype=np.float32)
|
||||
# mask used during training, to ignore some hard areas
|
||||
training_mask = np.ones((h, w), dtype=np.uint8)
|
||||
for poly_idx, poly_tag in enumerate(zip(polys, tags)):
|
||||
poly = poly_tag[0]
|
||||
tag = poly_tag[1]
|
||||
|
||||
r = [None, None, None, None]
|
||||
for i in range(4):
|
||||
dist1 = np.linalg.norm(poly[i] - poly[(i + 1) % 4])
|
||||
dist2 = np.linalg.norm(poly[i] - poly[(i - 1) % 4])
|
||||
r[i] = min(dist1, dist2)
|
||||
# score map
|
||||
shrinked_poly = self.shrink_poly(poly.copy(), r).astype(np.int32)[
|
||||
np.newaxis, :, :
|
||||
]
|
||||
cv2.fillPoly(score_map, shrinked_poly, 1)
|
||||
cv2.fillPoly(poly_mask, shrinked_poly, poly_idx + 1)
|
||||
# if the poly is too small, then ignore it during training
|
||||
poly_h = min(
|
||||
np.linalg.norm(poly[0] - poly[3]), np.linalg.norm(poly[1] - poly[2])
|
||||
)
|
||||
poly_w = min(
|
||||
np.linalg.norm(poly[0] - poly[1]), np.linalg.norm(poly[2] - poly[3])
|
||||
)
|
||||
if min(poly_h, poly_w) < self.min_text_size:
|
||||
cv2.fillPoly(training_mask, poly.astype(np.int32)[np.newaxis, :, :], 0)
|
||||
|
||||
if tag:
|
||||
cv2.fillPoly(training_mask, poly.astype(np.int32)[np.newaxis, :, :], 0)
|
||||
|
||||
xy_in_poly = np.argwhere(poly_mask == (poly_idx + 1))
|
||||
# geo map.
|
||||
y_in_poly = xy_in_poly[:, 0]
|
||||
x_in_poly = xy_in_poly[:, 1]
|
||||
poly[:, 0] = np.minimum(np.maximum(poly[:, 0], 0), w)
|
||||
poly[:, 1] = np.minimum(np.maximum(poly[:, 1], 0), h)
|
||||
for pno in range(4):
|
||||
geo_channel_beg = pno * 2
|
||||
geo_map[y_in_poly, x_in_poly, geo_channel_beg] = (
|
||||
x_in_poly - poly[pno, 0]
|
||||
)
|
||||
geo_map[y_in_poly, x_in_poly, geo_channel_beg + 1] = (
|
||||
y_in_poly - poly[pno, 1]
|
||||
)
|
||||
geo_map[y_in_poly, x_in_poly, 8] = 1.0 / max(min(poly_h, poly_w), 1.0)
|
||||
return score_map, geo_map, training_mask
|
||||
|
||||
def crop_area(self, im, polys, tags, crop_background=False, max_tries=50):
|
||||
"""
|
||||
make random crop from the input image
|
||||
:param im:
|
||||
:param polys:
|
||||
:param tags:
|
||||
:param crop_background:
|
||||
:param max_tries:
|
||||
:return:
|
||||
"""
|
||||
h, w, _ = im.shape
|
||||
pad_h = h // 10
|
||||
pad_w = w // 10
|
||||
h_array = np.zeros((h + pad_h * 2), dtype=np.int32)
|
||||
w_array = np.zeros((w + pad_w * 2), dtype=np.int32)
|
||||
for poly in polys:
|
||||
poly = np.round(poly, decimals=0).astype(np.int32)
|
||||
minx = np.min(poly[:, 0])
|
||||
maxx = np.max(poly[:, 0])
|
||||
w_array[minx + pad_w : maxx + pad_w] = 1
|
||||
miny = np.min(poly[:, 1])
|
||||
maxy = np.max(poly[:, 1])
|
||||
h_array[miny + pad_h : maxy + pad_h] = 1
|
||||
# ensure the cropped area not across a text
|
||||
h_axis = np.where(h_array == 0)[0]
|
||||
w_axis = np.where(w_array == 0)[0]
|
||||
if len(h_axis) == 0 or len(w_axis) == 0:
|
||||
return im, polys, tags
|
||||
|
||||
for i in range(max_tries):
|
||||
xx = np.random.choice(w_axis, size=2)
|
||||
xmin = np.min(xx) - pad_w
|
||||
xmax = np.max(xx) - pad_w
|
||||
xmin = np.clip(xmin, 0, w - 1)
|
||||
xmax = np.clip(xmax, 0, w - 1)
|
||||
yy = np.random.choice(h_axis, size=2)
|
||||
ymin = np.min(yy) - pad_h
|
||||
ymax = np.max(yy) - pad_h
|
||||
ymin = np.clip(ymin, 0, h - 1)
|
||||
ymax = np.clip(ymax, 0, h - 1)
|
||||
if (
|
||||
xmax - xmin < self.min_crop_side_ratio * w
|
||||
or ymax - ymin < self.min_crop_side_ratio * h
|
||||
):
|
||||
# area too small
|
||||
continue
|
||||
if polys.shape[0] != 0:
|
||||
poly_axis_in_area = (
|
||||
(polys[:, :, 0] >= xmin)
|
||||
& (polys[:, :, 0] <= xmax)
|
||||
& (polys[:, :, 1] >= ymin)
|
||||
& (polys[:, :, 1] <= ymax)
|
||||
)
|
||||
selected_polys = np.where(np.sum(poly_axis_in_area, axis=1) == 4)[0]
|
||||
else:
|
||||
selected_polys = []
|
||||
|
||||
if len(selected_polys) == 0:
|
||||
# no text in this area
|
||||
if crop_background:
|
||||
im = im[ymin : ymax + 1, xmin : xmax + 1, :]
|
||||
polys = []
|
||||
tags = []
|
||||
return im, polys, tags
|
||||
else:
|
||||
continue
|
||||
|
||||
im = im[ymin : ymax + 1, xmin : xmax + 1, :]
|
||||
polys = polys[selected_polys]
|
||||
tags = tags[selected_polys]
|
||||
polys[:, :, 0] -= xmin
|
||||
polys[:, :, 1] -= ymin
|
||||
return im, polys, tags
|
||||
return im, polys, tags
|
||||
|
||||
def crop_background_infor(self, im, text_polys, text_tags):
|
||||
im, text_polys, text_tags = self.crop_area(
|
||||
im, text_polys, text_tags, crop_background=True
|
||||
)
|
||||
|
||||
if len(text_polys) > 0:
|
||||
return None
|
||||
# pad and resize image
|
||||
input_size = self.input_size
|
||||
im, ratio = self.preprocess(im)
|
||||
score_map = np.zeros((input_size, input_size), dtype=np.float32)
|
||||
geo_map = np.zeros((input_size, input_size, 9), dtype=np.float32)
|
||||
training_mask = np.ones((input_size, input_size), dtype=np.float32)
|
||||
return im, score_map, geo_map, training_mask
|
||||
|
||||
def crop_foreground_infor(self, im, text_polys, text_tags):
|
||||
im, text_polys, text_tags = self.crop_area(
|
||||
im, text_polys, text_tags, crop_background=False
|
||||
)
|
||||
|
||||
if text_polys.shape[0] == 0:
|
||||
return None
|
||||
# continue for all ignore case
|
||||
if np.sum((text_tags * 1.0)) >= text_tags.size:
|
||||
return None
|
||||
# pad and resize image
|
||||
input_size = self.input_size
|
||||
im, ratio = self.preprocess(im)
|
||||
text_polys[:, :, 0] *= ratio
|
||||
text_polys[:, :, 1] *= ratio
|
||||
_, _, new_h, new_w = im.shape
|
||||
# print(im.shape)
|
||||
# self.draw_img_polys(im, text_polys)
|
||||
score_map, geo_map, training_mask = self.generate_quad(
|
||||
(new_h, new_w), text_polys, text_tags
|
||||
)
|
||||
return im, score_map, geo_map, training_mask
|
||||
|
||||
def __call__(self, data):
|
||||
im = data["image"]
|
||||
text_polys = data["polys"]
|
||||
text_tags = data["ignore_tags"]
|
||||
if im is None:
|
||||
return None
|
||||
if text_polys.shape[0] == 0:
|
||||
return None
|
||||
|
||||
# add rotate cases
|
||||
if np.random.rand() < 0.5:
|
||||
im, text_polys = self.rotate_im_poly(im, text_polys)
|
||||
h, w, _ = im.shape
|
||||
text_polys, text_tags = self.check_and_validate_polys(
|
||||
text_polys, text_tags, h, w
|
||||
)
|
||||
if text_polys.shape[0] == 0:
|
||||
return None
|
||||
|
||||
# random scale this image
|
||||
rd_scale = np.random.choice(self.random_scale)
|
||||
im = cv2.resize(im, dsize=None, fx=rd_scale, fy=rd_scale)
|
||||
text_polys *= rd_scale
|
||||
if np.random.rand() < self.background_ratio:
|
||||
outs = self.crop_background_infor(im, text_polys, text_tags)
|
||||
else:
|
||||
outs = self.crop_foreground_infor(im, text_polys, text_tags)
|
||||
|
||||
if outs is None:
|
||||
return None
|
||||
im, score_map, geo_map, training_mask = outs
|
||||
score_map = score_map[np.newaxis, ::4, ::4].astype(np.float32)
|
||||
geo_map = np.swapaxes(geo_map, 1, 2)
|
||||
geo_map = np.swapaxes(geo_map, 1, 0)
|
||||
geo_map = geo_map[:, ::4, ::4].astype(np.float32)
|
||||
training_mask = training_mask[np.newaxis, ::4, ::4]
|
||||
training_mask = training_mask.astype(np.float32)
|
||||
|
||||
data["image"] = im[0]
|
||||
data["score_map"] = score_map
|
||||
data["geo_map"] = geo_map
|
||||
data["training_mask"] = training_mask
|
||||
return data
|
||||
575
ppocr/data/imaug/fce_aug.py
Normal file
575
ppocr/data/imaug/fce_aug.py
Normal file
@@ -0,0 +1,575 @@
|
||||
# copyright (c) 2022 PaddlePaddle Authors. All Rights Reserve.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
"""
|
||||
This code is refer from:
|
||||
https://github.com/open-mmlab/mmocr/blob/main/mmocr/datasets/pipelines/transforms.py
|
||||
"""
|
||||
import numpy as np
|
||||
from PIL import Image, ImageDraw
|
||||
import cv2
|
||||
from shapely.geometry import Polygon
|
||||
import math
|
||||
from ppocr.utils.poly_nms import poly_intersection
|
||||
|
||||
|
||||
class RandomScaling:
|
||||
def __init__(self, size=800, scale=(3.0 / 4, 5.0 / 2), **kwargs):
|
||||
"""Random scale the image while keeping aspect.
|
||||
|
||||
Args:
|
||||
size (int) : Base size before scaling.
|
||||
scale (tuple(float)) : The range of scaling.
|
||||
"""
|
||||
assert isinstance(size, int)
|
||||
assert isinstance(scale, float) or isinstance(scale, tuple)
|
||||
self.size = size
|
||||
self.scale = scale if isinstance(scale, tuple) else (1 - scale, 1 + scale)
|
||||
|
||||
def __call__(self, data):
|
||||
image = data["image"]
|
||||
text_polys = data["polys"]
|
||||
h, w, _ = image.shape
|
||||
|
||||
aspect_ratio = np.random.uniform(min(self.scale), max(self.scale))
|
||||
scales = self.size * 1.0 / max(h, w) * aspect_ratio
|
||||
scales = np.array([scales, scales])
|
||||
out_size = (int(h * scales[1]), int(w * scales[0]))
|
||||
image = cv2.resize(image, out_size[::-1])
|
||||
|
||||
data["image"] = image
|
||||
text_polys[:, :, 0::2] = text_polys[:, :, 0::2] * scales[1]
|
||||
text_polys[:, :, 1::2] = text_polys[:, :, 1::2] * scales[0]
|
||||
data["polys"] = text_polys
|
||||
|
||||
return data
|
||||
|
||||
|
||||
class RandomCropFlip:
|
||||
def __init__(
|
||||
self, pad_ratio=0.1, crop_ratio=0.5, iter_num=1, min_area_ratio=0.2, **kwargs
|
||||
):
|
||||
"""Random crop and flip a patch of the image.
|
||||
|
||||
Args:
|
||||
crop_ratio (float): The ratio of cropping.
|
||||
iter_num (int): Number of operations.
|
||||
min_area_ratio (float): Minimal area ratio between cropped patch
|
||||
and original image.
|
||||
"""
|
||||
assert isinstance(crop_ratio, float)
|
||||
assert isinstance(iter_num, int)
|
||||
assert isinstance(min_area_ratio, float)
|
||||
|
||||
self.pad_ratio = pad_ratio
|
||||
self.epsilon = 1e-2
|
||||
self.crop_ratio = crop_ratio
|
||||
self.iter_num = iter_num
|
||||
self.min_area_ratio = min_area_ratio
|
||||
|
||||
def __call__(self, results):
|
||||
for i in range(self.iter_num):
|
||||
results = self.random_crop_flip(results)
|
||||
|
||||
return results
|
||||
|
||||
def random_crop_flip(self, results):
|
||||
image = results["image"]
|
||||
polygons = results["polys"]
|
||||
ignore_tags = results["ignore_tags"]
|
||||
if len(polygons) == 0:
|
||||
return results
|
||||
|
||||
if np.random.random() >= self.crop_ratio:
|
||||
return results
|
||||
|
||||
h, w, _ = image.shape
|
||||
area = h * w
|
||||
pad_h = int(h * self.pad_ratio)
|
||||
pad_w = int(w * self.pad_ratio)
|
||||
h_axis, w_axis = self.generate_crop_target(image, polygons, pad_h, pad_w)
|
||||
if len(h_axis) == 0 or len(w_axis) == 0:
|
||||
return results
|
||||
|
||||
attempt = 0
|
||||
while attempt < 50:
|
||||
attempt += 1
|
||||
polys_keep = []
|
||||
polys_new = []
|
||||
ignore_tags_keep = []
|
||||
ignore_tags_new = []
|
||||
xx = np.random.choice(w_axis, size=2)
|
||||
xmin = np.min(xx) - pad_w
|
||||
xmax = np.max(xx) - pad_w
|
||||
xmin = np.clip(xmin, 0, w - 1)
|
||||
xmax = np.clip(xmax, 0, w - 1)
|
||||
yy = np.random.choice(h_axis, size=2)
|
||||
ymin = np.min(yy) - pad_h
|
||||
ymax = np.max(yy) - pad_h
|
||||
ymin = np.clip(ymin, 0, h - 1)
|
||||
ymax = np.clip(ymax, 0, h - 1)
|
||||
if (xmax - xmin) * (ymax - ymin) < area * self.min_area_ratio:
|
||||
# area too small
|
||||
continue
|
||||
|
||||
pts = np.stack(
|
||||
[[xmin, xmax, xmax, xmin], [ymin, ymin, ymax, ymax]]
|
||||
).T.astype(np.int32)
|
||||
pp = Polygon(pts)
|
||||
fail_flag = False
|
||||
for polygon, ignore_tag in zip(polygons, ignore_tags):
|
||||
ppi = Polygon(polygon.reshape(-1, 2))
|
||||
ppiou, _ = poly_intersection(ppi, pp, buffer=0)
|
||||
if (
|
||||
np.abs(ppiou - float(ppi.area)) > self.epsilon
|
||||
and np.abs(ppiou) > self.epsilon
|
||||
):
|
||||
fail_flag = True
|
||||
break
|
||||
elif np.abs(ppiou - float(ppi.area)) < self.epsilon:
|
||||
polys_new.append(polygon)
|
||||
ignore_tags_new.append(ignore_tag)
|
||||
else:
|
||||
polys_keep.append(polygon)
|
||||
ignore_tags_keep.append(ignore_tag)
|
||||
|
||||
if fail_flag:
|
||||
continue
|
||||
else:
|
||||
break
|
||||
|
||||
cropped = image[ymin:ymax, xmin:xmax, :]
|
||||
select_type = np.random.randint(3)
|
||||
if select_type == 0:
|
||||
img = np.ascontiguousarray(cropped[:, ::-1])
|
||||
elif select_type == 1:
|
||||
img = np.ascontiguousarray(cropped[::-1, :])
|
||||
else:
|
||||
img = np.ascontiguousarray(cropped[::-1, ::-1])
|
||||
image[ymin:ymax, xmin:xmax, :] = img
|
||||
results["img"] = image
|
||||
|
||||
if len(polys_new) != 0:
|
||||
height, width, _ = cropped.shape
|
||||
if select_type == 0:
|
||||
for idx, polygon in enumerate(polys_new):
|
||||
poly = polygon.reshape(-1, 2)
|
||||
poly[:, 0] = width - poly[:, 0] + 2 * xmin
|
||||
polys_new[idx] = poly
|
||||
elif select_type == 1:
|
||||
for idx, polygon in enumerate(polys_new):
|
||||
poly = polygon.reshape(-1, 2)
|
||||
poly[:, 1] = height - poly[:, 1] + 2 * ymin
|
||||
polys_new[idx] = poly
|
||||
else:
|
||||
for idx, polygon in enumerate(polys_new):
|
||||
poly = polygon.reshape(-1, 2)
|
||||
poly[:, 0] = width - poly[:, 0] + 2 * xmin
|
||||
poly[:, 1] = height - poly[:, 1] + 2 * ymin
|
||||
polys_new[idx] = poly
|
||||
polygons = polys_keep + polys_new
|
||||
ignore_tags = ignore_tags_keep + ignore_tags_new
|
||||
results["polys"] = np.array(polygons)
|
||||
results["ignore_tags"] = ignore_tags
|
||||
|
||||
return results
|
||||
|
||||
def generate_crop_target(self, image, all_polys, pad_h, pad_w):
|
||||
"""Generate crop target and make sure not to crop the polygon
|
||||
instances.
|
||||
|
||||
Args:
|
||||
image (ndarray): The image waited to be crop.
|
||||
all_polys (list[list[ndarray]]): All polygons including ground
|
||||
truth polygons and ground truth ignored polygons.
|
||||
pad_h (int): Padding length of height.
|
||||
pad_w (int): Padding length of width.
|
||||
Returns:
|
||||
h_axis (ndarray): Vertical cropping range.
|
||||
w_axis (ndarray): Horizontal cropping range.
|
||||
"""
|
||||
h, w, _ = image.shape
|
||||
h_array = np.zeros((h + pad_h * 2), dtype=np.int32)
|
||||
w_array = np.zeros((w + pad_w * 2), dtype=np.int32)
|
||||
|
||||
text_polys = []
|
||||
for polygon in all_polys:
|
||||
rect = cv2.minAreaRect(polygon.astype(np.int32).reshape(-1, 2))
|
||||
box = cv2.boxPoints(rect)
|
||||
box = np.int64(box)
|
||||
text_polys.append([box[0], box[1], box[2], box[3]])
|
||||
|
||||
polys = np.array(text_polys, dtype=np.int32)
|
||||
for poly in polys:
|
||||
poly = np.round(poly, decimals=0).astype(np.int32)
|
||||
minx = np.min(poly[:, 0])
|
||||
maxx = np.max(poly[:, 0])
|
||||
w_array[minx + pad_w : maxx + pad_w] = 1
|
||||
miny = np.min(poly[:, 1])
|
||||
maxy = np.max(poly[:, 1])
|
||||
h_array[miny + pad_h : maxy + pad_h] = 1
|
||||
|
||||
h_axis = np.where(h_array == 0)[0]
|
||||
w_axis = np.where(w_array == 0)[0]
|
||||
return h_axis, w_axis
|
||||
|
||||
|
||||
class RandomCropPolyInstances:
|
||||
"""Randomly crop images and make sure to contain at least one intact
|
||||
instance."""
|
||||
|
||||
def __init__(self, crop_ratio=5.0 / 8.0, min_side_ratio=0.4, **kwargs):
|
||||
super().__init__()
|
||||
self.crop_ratio = crop_ratio
|
||||
self.min_side_ratio = min_side_ratio
|
||||
|
||||
def sample_valid_start_end(self, valid_array, min_len, max_start, min_end):
|
||||
assert isinstance(min_len, int)
|
||||
assert len(valid_array) > min_len
|
||||
|
||||
start_array = valid_array.copy()
|
||||
max_start = min(len(start_array) - min_len, max_start)
|
||||
start_array[max_start:] = 0
|
||||
start_array[0] = 1
|
||||
diff_array = np.hstack([0, start_array]) - np.hstack([start_array, 0])
|
||||
region_starts = np.where(diff_array < 0)[0]
|
||||
region_ends = np.where(diff_array > 0)[0]
|
||||
region_ind = np.random.randint(0, len(region_starts))
|
||||
start = np.random.randint(region_starts[region_ind], region_ends[region_ind])
|
||||
|
||||
end_array = valid_array.copy()
|
||||
min_end = max(start + min_len, min_end)
|
||||
end_array[:min_end] = 0
|
||||
end_array[-1] = 1
|
||||
diff_array = np.hstack([0, end_array]) - np.hstack([end_array, 0])
|
||||
region_starts = np.where(diff_array < 0)[0]
|
||||
region_ends = np.where(diff_array > 0)[0]
|
||||
region_ind = np.random.randint(0, len(region_starts))
|
||||
end = np.random.randint(region_starts[region_ind], region_ends[region_ind])
|
||||
return start, end
|
||||
|
||||
def sample_crop_box(self, img_size, results):
|
||||
"""Generate crop box and make sure not to crop the polygon instances.
|
||||
|
||||
Args:
|
||||
img_size (tuple(int)): The image size (h, w).
|
||||
results (dict): The results dict.
|
||||
"""
|
||||
|
||||
assert isinstance(img_size, tuple)
|
||||
h, w = img_size[:2]
|
||||
|
||||
key_masks = results["polys"]
|
||||
|
||||
x_valid_array = np.ones(w, dtype=np.int32)
|
||||
y_valid_array = np.ones(h, dtype=np.int32)
|
||||
|
||||
selected_mask = key_masks[np.random.randint(0, len(key_masks))]
|
||||
selected_mask = selected_mask.reshape((-1, 2)).astype(np.int32)
|
||||
max_x_start = max(np.min(selected_mask[:, 0]) - 2, 0)
|
||||
min_x_end = min(np.max(selected_mask[:, 0]) + 3, w - 1)
|
||||
max_y_start = max(np.min(selected_mask[:, 1]) - 2, 0)
|
||||
min_y_end = min(np.max(selected_mask[:, 1]) + 3, h - 1)
|
||||
|
||||
for mask in key_masks:
|
||||
mask = mask.reshape((-1, 2)).astype(np.int32)
|
||||
clip_x = np.clip(mask[:, 0], 0, w - 1)
|
||||
clip_y = np.clip(mask[:, 1], 0, h - 1)
|
||||
min_x, max_x = np.min(clip_x), np.max(clip_x)
|
||||
min_y, max_y = np.min(clip_y), np.max(clip_y)
|
||||
|
||||
x_valid_array[min_x - 2 : max_x + 3] = 0
|
||||
y_valid_array[min_y - 2 : max_y + 3] = 0
|
||||
|
||||
min_w = int(w * self.min_side_ratio)
|
||||
min_h = int(h * self.min_side_ratio)
|
||||
|
||||
x1, x2 = self.sample_valid_start_end(
|
||||
x_valid_array, min_w, max_x_start, min_x_end
|
||||
)
|
||||
y1, y2 = self.sample_valid_start_end(
|
||||
y_valid_array, min_h, max_y_start, min_y_end
|
||||
)
|
||||
|
||||
return np.array([x1, y1, x2, y2])
|
||||
|
||||
def crop_img(self, img, bbox):
|
||||
assert img.ndim == 3
|
||||
h, w, _ = img.shape
|
||||
assert 0 <= bbox[1] < bbox[3] <= h
|
||||
assert 0 <= bbox[0] < bbox[2] <= w
|
||||
return img[bbox[1] : bbox[3], bbox[0] : bbox[2]]
|
||||
|
||||
def __call__(self, results):
|
||||
image = results["image"]
|
||||
polygons = results["polys"]
|
||||
ignore_tags = results["ignore_tags"]
|
||||
if len(polygons) < 1:
|
||||
return results
|
||||
|
||||
if np.random.random_sample() < self.crop_ratio:
|
||||
crop_box = self.sample_crop_box(image.shape, results)
|
||||
img = self.crop_img(image, crop_box)
|
||||
results["image"] = img
|
||||
# crop and filter masks
|
||||
x1, y1, x2, y2 = crop_box
|
||||
w = max(x2 - x1, 1)
|
||||
h = max(y2 - y1, 1)
|
||||
polygons[:, :, 0::2] = polygons[:, :, 0::2] - x1
|
||||
polygons[:, :, 1::2] = polygons[:, :, 1::2] - y1
|
||||
|
||||
valid_masks_list = []
|
||||
valid_tags_list = []
|
||||
for ind, polygon in enumerate(polygons):
|
||||
if (
|
||||
(polygon[:, ::2] > -4).all()
|
||||
and (polygon[:, ::2] < w + 4).all()
|
||||
and (polygon[:, 1::2] > -4).all()
|
||||
and (polygon[:, 1::2] < h + 4).all()
|
||||
):
|
||||
polygon[:, ::2] = np.clip(polygon[:, ::2], 0, w)
|
||||
polygon[:, 1::2] = np.clip(polygon[:, 1::2], 0, h)
|
||||
valid_masks_list.append(polygon)
|
||||
valid_tags_list.append(ignore_tags[ind])
|
||||
|
||||
results["polys"] = np.array(valid_masks_list)
|
||||
results["ignore_tags"] = valid_tags_list
|
||||
|
||||
return results
|
||||
|
||||
def __repr__(self):
|
||||
repr_str = self.__class__.__name__
|
||||
return repr_str
|
||||
|
||||
|
||||
class RandomRotatePolyInstances:
|
||||
def __init__(
|
||||
self,
|
||||
rotate_ratio=0.5,
|
||||
max_angle=10,
|
||||
pad_with_fixed_color=False,
|
||||
pad_value=(0, 0, 0),
|
||||
**kwargs,
|
||||
):
|
||||
"""Randomly rotate images and polygon masks.
|
||||
|
||||
Args:
|
||||
rotate_ratio (float): The ratio of samples to operate rotation.
|
||||
max_angle (int): The maximum rotation angle.
|
||||
pad_with_fixed_color (bool): The flag for whether to pad rotated
|
||||
image with fixed value. If set to False, the rotated image will
|
||||
be padded onto cropped image.
|
||||
pad_value (tuple(int)): The color value for padding rotated image.
|
||||
"""
|
||||
self.rotate_ratio = rotate_ratio
|
||||
self.max_angle = max_angle
|
||||
self.pad_with_fixed_color = pad_with_fixed_color
|
||||
self.pad_value = pad_value
|
||||
|
||||
def rotate(self, center, points, theta, center_shift=(0, 0)):
|
||||
# rotate points.
|
||||
(center_x, center_y) = center
|
||||
center_y = -center_y
|
||||
x, y = points[:, ::2], points[:, 1::2]
|
||||
y = -y
|
||||
|
||||
theta = theta / 180 * math.pi
|
||||
cos = math.cos(theta)
|
||||
sin = math.sin(theta)
|
||||
|
||||
x = x - center_x
|
||||
y = y - center_y
|
||||
|
||||
_x = center_x + x * cos - y * sin + center_shift[0]
|
||||
_y = -(center_y + x * sin + y * cos) + center_shift[1]
|
||||
|
||||
points[:, ::2], points[:, 1::2] = _x, _y
|
||||
return points
|
||||
|
||||
def cal_canvas_size(self, ori_size, degree):
|
||||
assert isinstance(ori_size, tuple)
|
||||
angle = degree * math.pi / 180.0
|
||||
h, w = ori_size[:2]
|
||||
|
||||
cos = math.cos(angle)
|
||||
sin = math.sin(angle)
|
||||
canvas_h = int(w * math.fabs(sin) + h * math.fabs(cos))
|
||||
canvas_w = int(w * math.fabs(cos) + h * math.fabs(sin))
|
||||
|
||||
canvas_size = (canvas_h, canvas_w)
|
||||
return canvas_size
|
||||
|
||||
def sample_angle(self, max_angle):
|
||||
angle = np.random.random_sample() * 2 * max_angle - max_angle
|
||||
return angle
|
||||
|
||||
def rotate_img(self, img, angle, canvas_size):
|
||||
h, w = img.shape[:2]
|
||||
rotation_matrix = cv2.getRotationMatrix2D((w / 2, h / 2), angle, 1)
|
||||
rotation_matrix[0, 2] += int((canvas_size[1] - w) / 2)
|
||||
rotation_matrix[1, 2] += int((canvas_size[0] - h) / 2)
|
||||
|
||||
if self.pad_with_fixed_color:
|
||||
target_img = cv2.warpAffine(
|
||||
img,
|
||||
rotation_matrix,
|
||||
(canvas_size[1], canvas_size[0]),
|
||||
flags=cv2.INTER_NEAREST,
|
||||
borderValue=self.pad_value,
|
||||
)
|
||||
else:
|
||||
mask = np.zeros_like(img)
|
||||
(h_ind, w_ind) = (
|
||||
np.random.randint(0, h * 7 // 8),
|
||||
np.random.randint(0, w * 7 // 8),
|
||||
)
|
||||
img_cut = img[h_ind : (h_ind + h // 9), w_ind : (w_ind + w // 9)]
|
||||
img_cut = cv2.resize(img_cut, (canvas_size[1], canvas_size[0]))
|
||||
|
||||
mask = cv2.warpAffine(
|
||||
mask,
|
||||
rotation_matrix,
|
||||
(canvas_size[1], canvas_size[0]),
|
||||
borderValue=[1, 1, 1],
|
||||
)
|
||||
target_img = cv2.warpAffine(
|
||||
img,
|
||||
rotation_matrix,
|
||||
(canvas_size[1], canvas_size[0]),
|
||||
borderValue=[0, 0, 0],
|
||||
)
|
||||
target_img = target_img + img_cut * mask
|
||||
|
||||
return target_img
|
||||
|
||||
def __call__(self, results):
|
||||
if np.random.random_sample() < self.rotate_ratio:
|
||||
image = results["image"]
|
||||
polygons = results["polys"]
|
||||
h, w = image.shape[:2]
|
||||
|
||||
angle = self.sample_angle(self.max_angle)
|
||||
canvas_size = self.cal_canvas_size((h, w), angle)
|
||||
center_shift = (
|
||||
int((canvas_size[1] - w) / 2),
|
||||
int((canvas_size[0] - h) / 2),
|
||||
)
|
||||
image = self.rotate_img(image, angle, canvas_size)
|
||||
results["image"] = image
|
||||
# rotate polygons
|
||||
rotated_masks = []
|
||||
for mask in polygons:
|
||||
rotated_mask = self.rotate((w / 2, h / 2), mask, angle, center_shift)
|
||||
rotated_masks.append(rotated_mask)
|
||||
results["polys"] = np.array(rotated_masks)
|
||||
|
||||
return results
|
||||
|
||||
def __repr__(self):
|
||||
repr_str = self.__class__.__name__
|
||||
return repr_str
|
||||
|
||||
|
||||
class SquareResizePad:
|
||||
def __init__(
|
||||
self,
|
||||
target_size,
|
||||
pad_ratio=0.6,
|
||||
pad_with_fixed_color=False,
|
||||
pad_value=(0, 0, 0),
|
||||
**kwargs,
|
||||
):
|
||||
"""Resize or pad images to be square shape.
|
||||
|
||||
Args:
|
||||
target_size (int): The target size of square shaped image.
|
||||
pad_with_fixed_color (bool): The flag for whether to pad rotated
|
||||
image with fixed value. If set to False, the rescales image will
|
||||
be padded onto cropped image.
|
||||
pad_value (tuple(int)): The color value for padding rotated image.
|
||||
"""
|
||||
assert isinstance(target_size, int)
|
||||
assert isinstance(pad_ratio, float)
|
||||
assert isinstance(pad_with_fixed_color, bool)
|
||||
assert isinstance(pad_value, tuple)
|
||||
|
||||
self.target_size = target_size
|
||||
self.pad_ratio = pad_ratio
|
||||
self.pad_with_fixed_color = pad_with_fixed_color
|
||||
self.pad_value = pad_value
|
||||
|
||||
def resize_img(self, img, keep_ratio=True):
|
||||
h, w, _ = img.shape
|
||||
if keep_ratio:
|
||||
t_h = self.target_size if h >= w else int(h * self.target_size / w)
|
||||
t_w = self.target_size if h <= w else int(w * self.target_size / h)
|
||||
else:
|
||||
t_h = t_w = self.target_size
|
||||
img = cv2.resize(img, (t_w, t_h))
|
||||
return img, (t_h, t_w)
|
||||
|
||||
def square_pad(self, img):
|
||||
h, w = img.shape[:2]
|
||||
if h == w:
|
||||
return img, (0, 0)
|
||||
pad_size = max(h, w)
|
||||
if self.pad_with_fixed_color:
|
||||
expand_img = np.ones((pad_size, pad_size, 3), dtype=np.uint8)
|
||||
expand_img[:] = self.pad_value
|
||||
else:
|
||||
(h_ind, w_ind) = (
|
||||
np.random.randint(0, h * 7 // 8),
|
||||
np.random.randint(0, w * 7 // 8),
|
||||
)
|
||||
img_cut = img[h_ind : (h_ind + h // 9), w_ind : (w_ind + w // 9)]
|
||||
expand_img = cv2.resize(img_cut, (pad_size, pad_size))
|
||||
if h > w:
|
||||
y0, x0 = 0, (h - w) // 2
|
||||
else:
|
||||
y0, x0 = (w - h) // 2, 0
|
||||
expand_img[y0 : y0 + h, x0 : x0 + w] = img
|
||||
offset = (x0, y0)
|
||||
|
||||
return expand_img, offset
|
||||
|
||||
def square_pad_mask(self, points, offset):
|
||||
x0, y0 = offset
|
||||
pad_points = points.copy()
|
||||
pad_points[::2] = pad_points[::2] + x0
|
||||
pad_points[1::2] = pad_points[1::2] + y0
|
||||
return pad_points
|
||||
|
||||
def __call__(self, results):
|
||||
image = results["image"]
|
||||
polygons = results["polys"]
|
||||
h, w = image.shape[:2]
|
||||
|
||||
if np.random.random_sample() < self.pad_ratio:
|
||||
image, out_size = self.resize_img(image, keep_ratio=True)
|
||||
image, offset = self.square_pad(image)
|
||||
else:
|
||||
image, out_size = self.resize_img(image, keep_ratio=False)
|
||||
offset = (0, 0)
|
||||
results["image"] = image
|
||||
try:
|
||||
polygons[:, :, 0::2] = polygons[:, :, 0::2] * out_size[1] / w + offset[0]
|
||||
polygons[:, :, 1::2] = polygons[:, :, 1::2] * out_size[0] / h + offset[1]
|
||||
except:
|
||||
pass
|
||||
results["polys"] = polygons
|
||||
|
||||
return results
|
||||
|
||||
def __repr__(self):
|
||||
repr_str = self.__class__.__name__
|
||||
return repr_str
|
||||
697
ppocr/data/imaug/fce_targets.py
Normal file
697
ppocr/data/imaug/fce_targets.py
Normal file
@@ -0,0 +1,697 @@
|
||||
# copyright (c) 2022 PaddlePaddle Authors. All Rights Reserve.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
"""
|
||||
This code is refer from:
|
||||
https://github.com/open-mmlab/mmocr/blob/main/mmocr/datasets/pipelines/textdet_targets/fcenet_targets.py
|
||||
"""
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
from numpy.fft import fft
|
||||
from numpy.linalg import norm
|
||||
import sys
|
||||
|
||||
|
||||
def vector_slope(vec):
|
||||
assert len(vec) == 2
|
||||
return abs(vec[1] / (vec[0] + 1e-8))
|
||||
|
||||
|
||||
class FCENetTargets:
|
||||
"""Generate the ground truth targets of FCENet: Fourier Contour Embedding
|
||||
for Arbitrary-Shaped Text Detection.
|
||||
|
||||
[https://arxiv.org/abs/2104.10442]
|
||||
|
||||
Args:
|
||||
fourier_degree (int): The maximum Fourier transform degree k.
|
||||
resample_step (float): The step size for resampling the text center
|
||||
line (TCL). It's better not to exceed half of the minimum width.
|
||||
center_region_shrink_ratio (float): The shrink ratio of text center
|
||||
region.
|
||||
level_size_divisors (tuple(int)): The downsample ratio on each level.
|
||||
level_proportion_range (tuple(tuple(int))): The range of text sizes
|
||||
assigned to each level.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
fourier_degree=5,
|
||||
resample_step=4.0,
|
||||
center_region_shrink_ratio=0.3,
|
||||
level_size_divisors=(8, 16, 32),
|
||||
level_proportion_range=((0, 0.25), (0.2, 0.65), (0.55, 1.0)),
|
||||
orientation_thr=2.0,
|
||||
**kwargs,
|
||||
):
|
||||
super().__init__()
|
||||
assert isinstance(level_size_divisors, tuple)
|
||||
assert isinstance(level_proportion_range, tuple)
|
||||
assert len(level_size_divisors) == len(level_proportion_range)
|
||||
self.fourier_degree = fourier_degree
|
||||
self.resample_step = resample_step
|
||||
self.center_region_shrink_ratio = center_region_shrink_ratio
|
||||
self.level_size_divisors = level_size_divisors
|
||||
self.level_proportion_range = level_proportion_range
|
||||
|
||||
self.orientation_thr = orientation_thr
|
||||
|
||||
def vector_angle(self, vec1, vec2):
|
||||
if vec1.ndim > 1:
|
||||
unit_vec1 = vec1 / (norm(vec1, axis=-1) + 1e-8).reshape((-1, 1))
|
||||
else:
|
||||
unit_vec1 = vec1 / (norm(vec1, axis=-1) + 1e-8)
|
||||
if vec2.ndim > 1:
|
||||
unit_vec2 = vec2 / (norm(vec2, axis=-1) + 1e-8).reshape((-1, 1))
|
||||
else:
|
||||
unit_vec2 = vec2 / (norm(vec2, axis=-1) + 1e-8)
|
||||
return np.arccos(np.clip(np.sum(unit_vec1 * unit_vec2, axis=-1), -1.0, 1.0))
|
||||
|
||||
def resample_line(self, line, n):
|
||||
"""Resample n points on a line.
|
||||
|
||||
Args:
|
||||
line (ndarray): The points composing a line.
|
||||
n (int): The resampled points number.
|
||||
|
||||
Returns:
|
||||
resampled_line (ndarray): The points composing the resampled line.
|
||||
"""
|
||||
|
||||
assert line.ndim == 2
|
||||
assert line.shape[0] >= 2
|
||||
assert line.shape[1] == 2
|
||||
assert isinstance(n, int)
|
||||
assert n > 0
|
||||
|
||||
length_list = [norm(line[i + 1] - line[i]) for i in range(len(line) - 1)]
|
||||
total_length = sum(length_list)
|
||||
length_cumsum = np.cumsum([0.0] + length_list)
|
||||
delta_length = total_length / (float(n) + 1e-8)
|
||||
|
||||
current_edge_ind = 0
|
||||
resampled_line = [line[0]]
|
||||
|
||||
for i in range(1, n):
|
||||
current_line_len = i * delta_length
|
||||
|
||||
while (
|
||||
current_edge_ind + 1 < len(length_cumsum)
|
||||
and current_line_len >= length_cumsum[current_edge_ind + 1]
|
||||
):
|
||||
current_edge_ind += 1
|
||||
|
||||
current_edge_end_shift = current_line_len - length_cumsum[current_edge_ind]
|
||||
|
||||
if current_edge_ind >= len(length_list):
|
||||
break
|
||||
end_shift_ratio = current_edge_end_shift / length_list[current_edge_ind]
|
||||
current_point = (
|
||||
line[current_edge_ind]
|
||||
+ (line[current_edge_ind + 1] - line[current_edge_ind])
|
||||
* end_shift_ratio
|
||||
)
|
||||
resampled_line.append(current_point)
|
||||
resampled_line.append(line[-1])
|
||||
resampled_line = np.array(resampled_line)
|
||||
|
||||
return resampled_line
|
||||
|
||||
def reorder_poly_edge(self, points):
|
||||
"""Get the respective points composing head edge, tail edge, top
|
||||
sideline and bottom sideline.
|
||||
|
||||
Args:
|
||||
points (ndarray): The points composing a text polygon.
|
||||
|
||||
Returns:
|
||||
head_edge (ndarray): The two points composing the head edge of text
|
||||
polygon.
|
||||
tail_edge (ndarray): The two points composing the tail edge of text
|
||||
polygon.
|
||||
top_sideline (ndarray): The points composing top curved sideline of
|
||||
text polygon.
|
||||
bot_sideline (ndarray): The points composing bottom curved sideline
|
||||
of text polygon.
|
||||
"""
|
||||
|
||||
assert points.ndim == 2
|
||||
assert points.shape[0] >= 4
|
||||
assert points.shape[1] == 2
|
||||
|
||||
head_inds, tail_inds = self.find_head_tail(points, self.orientation_thr)
|
||||
head_edge, tail_edge = points[head_inds], points[tail_inds]
|
||||
|
||||
pad_points = np.vstack([points, points])
|
||||
if tail_inds[1] < 1:
|
||||
tail_inds[1] = len(points)
|
||||
sideline1 = pad_points[head_inds[1] : tail_inds[1]]
|
||||
sideline2 = pad_points[tail_inds[1] : (head_inds[1] + len(points))]
|
||||
sideline_mean_shift = np.mean(sideline1, axis=0) - np.mean(sideline2, axis=0)
|
||||
|
||||
if sideline_mean_shift[1] > 0:
|
||||
top_sideline, bot_sideline = sideline2, sideline1
|
||||
else:
|
||||
top_sideline, bot_sideline = sideline1, sideline2
|
||||
|
||||
return head_edge, tail_edge, top_sideline, bot_sideline
|
||||
|
||||
def find_head_tail(self, points, orientation_thr):
|
||||
"""Find the head edge and tail edge of a text polygon.
|
||||
|
||||
Args:
|
||||
points (ndarray): The points composing a text polygon.
|
||||
orientation_thr (float): The threshold for distinguishing between
|
||||
head edge and tail edge among the horizontal and vertical edges
|
||||
of a quadrangle.
|
||||
|
||||
Returns:
|
||||
head_inds (list): The indexes of two points composing head edge.
|
||||
tail_inds (list): The indexes of two points composing tail edge.
|
||||
"""
|
||||
|
||||
assert points.ndim == 2
|
||||
assert points.shape[0] >= 4
|
||||
assert points.shape[1] == 2
|
||||
assert isinstance(orientation_thr, float)
|
||||
|
||||
if len(points) > 4:
|
||||
pad_points = np.vstack([points, points[0]])
|
||||
edge_vec = pad_points[1:] - pad_points[:-1]
|
||||
|
||||
theta_sum = []
|
||||
adjacent_vec_theta = []
|
||||
for i, edge_vec1 in enumerate(edge_vec):
|
||||
adjacent_ind = [x % len(edge_vec) for x in [i - 1, i + 1]]
|
||||
adjacent_edge_vec = edge_vec[adjacent_ind]
|
||||
temp_theta_sum = np.sum(self.vector_angle(edge_vec1, adjacent_edge_vec))
|
||||
temp_adjacent_theta = self.vector_angle(
|
||||
adjacent_edge_vec[0], adjacent_edge_vec[1]
|
||||
)
|
||||
theta_sum.append(temp_theta_sum)
|
||||
adjacent_vec_theta.append(temp_adjacent_theta)
|
||||
theta_sum_score = np.array(theta_sum) / np.pi
|
||||
adjacent_theta_score = np.array(adjacent_vec_theta) / np.pi
|
||||
poly_center = np.mean(points, axis=0)
|
||||
edge_dist = np.maximum(
|
||||
norm(pad_points[1:] - poly_center, axis=-1),
|
||||
norm(pad_points[:-1] - poly_center, axis=-1),
|
||||
)
|
||||
dist_score = edge_dist / np.max(edge_dist)
|
||||
position_score = np.zeros(len(edge_vec))
|
||||
score = 0.5 * theta_sum_score + 0.15 * adjacent_theta_score
|
||||
score += 0.35 * dist_score
|
||||
if len(points) % 2 == 0:
|
||||
position_score[(len(score) // 2 - 1)] += 1
|
||||
position_score[-1] += 1
|
||||
score += 0.1 * position_score
|
||||
pad_score = np.concatenate([score, score])
|
||||
score_matrix = np.zeros((len(score), len(score) - 3))
|
||||
x = np.arange(len(score) - 3) / float(len(score) - 4)
|
||||
gaussian = (
|
||||
1.0
|
||||
/ (np.sqrt(2.0 * np.pi) * 0.5)
|
||||
* np.exp(-np.power((x - 0.5) / 0.5, 2.0) / 2)
|
||||
)
|
||||
gaussian = gaussian / np.max(gaussian)
|
||||
for i in range(len(score)):
|
||||
score_matrix[i, :] = (
|
||||
score[i]
|
||||
+ pad_score[(i + 2) : (i + len(score) - 1)] * gaussian * 0.3
|
||||
)
|
||||
|
||||
head_start, tail_increment = np.unravel_index(
|
||||
score_matrix.argmax(), score_matrix.shape
|
||||
)
|
||||
tail_start = (head_start + tail_increment + 2) % len(points)
|
||||
head_end = (head_start + 1) % len(points)
|
||||
tail_end = (tail_start + 1) % len(points)
|
||||
|
||||
if head_end > tail_end:
|
||||
head_start, tail_start = tail_start, head_start
|
||||
head_end, tail_end = tail_end, head_end
|
||||
head_inds = [head_start, head_end]
|
||||
tail_inds = [tail_start, tail_end]
|
||||
else:
|
||||
if vector_slope(points[1] - points[0]) + vector_slope(
|
||||
points[3] - points[2]
|
||||
) < vector_slope(points[2] - points[1]) + vector_slope(
|
||||
points[0] - points[3]
|
||||
):
|
||||
horizontal_edge_inds = [[0, 1], [2, 3]]
|
||||
vertical_edge_inds = [[3, 0], [1, 2]]
|
||||
else:
|
||||
horizontal_edge_inds = [[3, 0], [1, 2]]
|
||||
vertical_edge_inds = [[0, 1], [2, 3]]
|
||||
|
||||
vertical_len_sum = norm(
|
||||
points[vertical_edge_inds[0][0]] - points[vertical_edge_inds[0][1]]
|
||||
) + norm(
|
||||
points[vertical_edge_inds[1][0]] - points[vertical_edge_inds[1][1]]
|
||||
)
|
||||
horizontal_len_sum = norm(
|
||||
points[horizontal_edge_inds[0][0]] - points[horizontal_edge_inds[0][1]]
|
||||
) + norm(
|
||||
points[horizontal_edge_inds[1][0]] - points[horizontal_edge_inds[1][1]]
|
||||
)
|
||||
|
||||
if vertical_len_sum > horizontal_len_sum * orientation_thr:
|
||||
head_inds = horizontal_edge_inds[0]
|
||||
tail_inds = horizontal_edge_inds[1]
|
||||
else:
|
||||
head_inds = vertical_edge_inds[0]
|
||||
tail_inds = vertical_edge_inds[1]
|
||||
|
||||
return head_inds, tail_inds
|
||||
|
||||
def resample_sidelines(self, sideline1, sideline2, resample_step):
|
||||
"""Resample two sidelines to be of the same points number according to
|
||||
step size.
|
||||
|
||||
Args:
|
||||
sideline1 (ndarray): The points composing a sideline of a text
|
||||
polygon.
|
||||
sideline2 (ndarray): The points composing another sideline of a
|
||||
text polygon.
|
||||
resample_step (float): The resampled step size.
|
||||
|
||||
Returns:
|
||||
resampled_line1 (ndarray): The resampled line 1.
|
||||
resampled_line2 (ndarray): The resampled line 2.
|
||||
"""
|
||||
|
||||
assert sideline1.ndim == sideline2.ndim == 2
|
||||
assert sideline1.shape[1] == sideline2.shape[1] == 2
|
||||
assert sideline1.shape[0] >= 2
|
||||
assert sideline2.shape[0] >= 2
|
||||
assert isinstance(resample_step, float)
|
||||
|
||||
length1 = sum(
|
||||
[norm(sideline1[i + 1] - sideline1[i]) for i in range(len(sideline1) - 1)]
|
||||
)
|
||||
length2 = sum(
|
||||
[norm(sideline2[i + 1] - sideline2[i]) for i in range(len(sideline2) - 1)]
|
||||
)
|
||||
|
||||
total_length = (length1 + length2) / 2
|
||||
resample_point_num = max(int(float(total_length) / resample_step), 1)
|
||||
|
||||
resampled_line1 = self.resample_line(sideline1, resample_point_num)
|
||||
resampled_line2 = self.resample_line(sideline2, resample_point_num)
|
||||
|
||||
return resampled_line1, resampled_line2
|
||||
|
||||
def generate_center_region_mask(self, img_size, text_polys):
|
||||
"""Generate text center region mask.
|
||||
|
||||
Args:
|
||||
img_size (tuple): The image size of (height, width).
|
||||
text_polys (list[list[ndarray]]): The list of text polygons.
|
||||
|
||||
Returns:
|
||||
center_region_mask (ndarray): The text center region mask.
|
||||
"""
|
||||
|
||||
assert isinstance(img_size, tuple)
|
||||
# assert check_argument.is_2dlist(text_polys)
|
||||
|
||||
h, w = img_size
|
||||
|
||||
center_region_mask = np.zeros((h, w), np.uint8)
|
||||
|
||||
center_region_boxes = []
|
||||
for poly in text_polys:
|
||||
# assert len(poly) == 1
|
||||
polygon_points = poly.reshape(-1, 2)
|
||||
_, _, top_line, bot_line = self.reorder_poly_edge(polygon_points)
|
||||
resampled_top_line, resampled_bot_line = self.resample_sidelines(
|
||||
top_line, bot_line, self.resample_step
|
||||
)
|
||||
resampled_bot_line = resampled_bot_line[::-1]
|
||||
if len(resampled_top_line) != len(resampled_bot_line):
|
||||
continue
|
||||
center_line = (resampled_top_line + resampled_bot_line) / 2
|
||||
|
||||
line_head_shrink_len = (
|
||||
norm(resampled_top_line[0] - resampled_bot_line[0]) / 4.0
|
||||
)
|
||||
line_tail_shrink_len = (
|
||||
norm(resampled_top_line[-1] - resampled_bot_line[-1]) / 4.0
|
||||
)
|
||||
head_shrink_num = int(line_head_shrink_len // self.resample_step)
|
||||
tail_shrink_num = int(line_tail_shrink_len // self.resample_step)
|
||||
if len(center_line) > head_shrink_num + tail_shrink_num + 2:
|
||||
center_line = center_line[
|
||||
head_shrink_num : len(center_line) - tail_shrink_num
|
||||
]
|
||||
resampled_top_line = resampled_top_line[
|
||||
head_shrink_num : len(resampled_top_line) - tail_shrink_num
|
||||
]
|
||||
resampled_bot_line = resampled_bot_line[
|
||||
head_shrink_num : len(resampled_bot_line) - tail_shrink_num
|
||||
]
|
||||
|
||||
for i in range(0, len(center_line) - 1):
|
||||
tl = (
|
||||
center_line[i]
|
||||
+ (resampled_top_line[i] - center_line[i])
|
||||
* self.center_region_shrink_ratio
|
||||
)
|
||||
tr = (
|
||||
center_line[i + 1]
|
||||
+ (resampled_top_line[i + 1] - center_line[i + 1])
|
||||
* self.center_region_shrink_ratio
|
||||
)
|
||||
br = (
|
||||
center_line[i + 1]
|
||||
+ (resampled_bot_line[i + 1] - center_line[i + 1])
|
||||
* self.center_region_shrink_ratio
|
||||
)
|
||||
bl = (
|
||||
center_line[i]
|
||||
+ (resampled_bot_line[i] - center_line[i])
|
||||
* self.center_region_shrink_ratio
|
||||
)
|
||||
current_center_box = np.vstack([tl, tr, br, bl]).astype(np.int32)
|
||||
center_region_boxes.append(current_center_box)
|
||||
|
||||
cv2.fillPoly(center_region_mask, center_region_boxes, 1)
|
||||
return center_region_mask
|
||||
|
||||
def resample_polygon(self, polygon, n=400):
|
||||
"""Resample one polygon with n points on its boundary.
|
||||
|
||||
Args:
|
||||
polygon (list[float]): The input polygon.
|
||||
n (int): The number of resampled points.
|
||||
Returns:
|
||||
resampled_polygon (list[float]): The resampled polygon.
|
||||
"""
|
||||
length = []
|
||||
|
||||
for i in range(len(polygon)):
|
||||
p1 = polygon[i]
|
||||
if i == len(polygon) - 1:
|
||||
p2 = polygon[0]
|
||||
else:
|
||||
p2 = polygon[i + 1]
|
||||
length.append(((p1[0] - p2[0]) ** 2 + (p1[1] - p2[1]) ** 2) ** 0.5)
|
||||
|
||||
total_length = sum(length)
|
||||
n_on_each_line = (np.array(length) / (total_length + 1e-8)) * n
|
||||
n_on_each_line = n_on_each_line.astype(np.int32)
|
||||
new_polygon = []
|
||||
|
||||
for i in range(len(polygon)):
|
||||
num = n_on_each_line[i]
|
||||
p1 = polygon[i]
|
||||
if i == len(polygon) - 1:
|
||||
p2 = polygon[0]
|
||||
else:
|
||||
p2 = polygon[i + 1]
|
||||
|
||||
if num == 0:
|
||||
continue
|
||||
|
||||
dxdy = (p2 - p1) / num
|
||||
for j in range(num):
|
||||
point = p1 + dxdy * j
|
||||
new_polygon.append(point)
|
||||
|
||||
return np.array(new_polygon)
|
||||
|
||||
def normalize_polygon(self, polygon):
|
||||
"""Normalize one polygon so that its start point is at right most.
|
||||
|
||||
Args:
|
||||
polygon (list[float]): The origin polygon.
|
||||
Returns:
|
||||
new_polygon (lost[float]): The polygon with start point at right.
|
||||
"""
|
||||
temp_polygon = polygon - polygon.mean(axis=0)
|
||||
x = np.abs(temp_polygon[:, 0])
|
||||
y = temp_polygon[:, 1]
|
||||
index_x = np.argsort(x)
|
||||
index_y = np.argmin(y[index_x[:8]])
|
||||
index = index_x[index_y]
|
||||
new_polygon = np.concatenate([polygon[index:], polygon[:index]])
|
||||
return new_polygon
|
||||
|
||||
def poly2fourier(self, polygon, fourier_degree):
|
||||
"""Perform Fourier transformation to generate Fourier coefficients ck
|
||||
from polygon.
|
||||
|
||||
Args:
|
||||
polygon (ndarray): An input polygon.
|
||||
fourier_degree (int): The maximum Fourier degree K.
|
||||
Returns:
|
||||
c (ndarray(complex)): Fourier coefficients.
|
||||
"""
|
||||
points = polygon[:, 0] + polygon[:, 1] * 1j
|
||||
c_fft = fft(points) / len(points)
|
||||
c = np.hstack((c_fft[-fourier_degree:], c_fft[: fourier_degree + 1]))
|
||||
return c
|
||||
|
||||
def clockwise(self, c, fourier_degree):
|
||||
"""Make sure the polygon reconstructed from Fourier coefficients c in
|
||||
the clockwise direction.
|
||||
|
||||
Args:
|
||||
polygon (list[float]): The origin polygon.
|
||||
Returns:
|
||||
new_polygon (lost[float]): The polygon in clockwise point order.
|
||||
"""
|
||||
if np.abs(c[fourier_degree + 1]) > np.abs(c[fourier_degree - 1]):
|
||||
return c
|
||||
elif np.abs(c[fourier_degree + 1]) < np.abs(c[fourier_degree - 1]):
|
||||
return c[::-1]
|
||||
else:
|
||||
if np.abs(c[fourier_degree + 2]) > np.abs(c[fourier_degree - 2]):
|
||||
return c
|
||||
else:
|
||||
return c[::-1]
|
||||
|
||||
def cal_fourier_signature(self, polygon, fourier_degree):
|
||||
"""Calculate Fourier signature from input polygon.
|
||||
|
||||
Args:
|
||||
polygon (ndarray): The input polygon.
|
||||
fourier_degree (int): The maximum Fourier degree K.
|
||||
Returns:
|
||||
fourier_signature (ndarray): An array shaped (2k+1, 2) containing
|
||||
real part and image part of 2k+1 Fourier coefficients.
|
||||
"""
|
||||
resampled_polygon = self.resample_polygon(polygon)
|
||||
resampled_polygon = self.normalize_polygon(resampled_polygon)
|
||||
|
||||
fourier_coeff = self.poly2fourier(resampled_polygon, fourier_degree)
|
||||
fourier_coeff = self.clockwise(fourier_coeff, fourier_degree)
|
||||
|
||||
real_part = np.real(fourier_coeff).reshape((-1, 1))
|
||||
image_part = np.imag(fourier_coeff).reshape((-1, 1))
|
||||
fourier_signature = np.hstack([real_part, image_part])
|
||||
|
||||
return fourier_signature
|
||||
|
||||
def generate_fourier_maps(self, img_size, text_polys):
|
||||
"""Generate Fourier coefficient maps.
|
||||
|
||||
Args:
|
||||
img_size (tuple): The image size of (height, width).
|
||||
text_polys (list[list[ndarray]]): The list of text polygons.
|
||||
|
||||
Returns:
|
||||
fourier_real_map (ndarray): The Fourier coefficient real part maps.
|
||||
fourier_image_map (ndarray): The Fourier coefficient image part
|
||||
maps.
|
||||
"""
|
||||
|
||||
assert isinstance(img_size, tuple)
|
||||
|
||||
h, w = img_size
|
||||
k = self.fourier_degree
|
||||
real_map = np.zeros((k * 2 + 1, h, w), dtype=np.float32)
|
||||
imag_map = np.zeros((k * 2 + 1, h, w), dtype=np.float32)
|
||||
|
||||
for poly in text_polys:
|
||||
mask = np.zeros((h, w), dtype=np.uint8)
|
||||
polygon = np.array(poly).reshape((1, -1, 2))
|
||||
cv2.fillPoly(mask, polygon.astype(np.int32), 1)
|
||||
fourier_coeff = self.cal_fourier_signature(polygon[0], k)
|
||||
for i in range(-k, k + 1):
|
||||
if i != 0:
|
||||
real_map[i + k, :, :] = (
|
||||
mask * fourier_coeff[i + k, 0]
|
||||
+ (1 - mask) * real_map[i + k, :, :]
|
||||
)
|
||||
imag_map[i + k, :, :] = (
|
||||
mask * fourier_coeff[i + k, 1]
|
||||
+ (1 - mask) * imag_map[i + k, :, :]
|
||||
)
|
||||
else:
|
||||
yx = np.argwhere(mask > 0.5)
|
||||
k_ind = np.ones((len(yx)), dtype=np.int64) * k
|
||||
y, x = yx[:, 0], yx[:, 1]
|
||||
real_map[k_ind, y, x] = fourier_coeff[k, 0] - x
|
||||
imag_map[k_ind, y, x] = fourier_coeff[k, 1] - y
|
||||
|
||||
return real_map, imag_map
|
||||
|
||||
def generate_text_region_mask(self, img_size, text_polys):
|
||||
"""Generate text center region mask and geometry attribute maps.
|
||||
|
||||
Args:
|
||||
img_size (tuple): The image size (height, width).
|
||||
text_polys (list[list[ndarray]]): The list of text polygons.
|
||||
|
||||
Returns:
|
||||
text_region_mask (ndarray): The text region mask.
|
||||
"""
|
||||
|
||||
assert isinstance(img_size, tuple)
|
||||
|
||||
h, w = img_size
|
||||
text_region_mask = np.zeros((h, w), dtype=np.uint8)
|
||||
|
||||
for poly in text_polys:
|
||||
polygon = np.array(poly, dtype=np.int32).reshape((1, -1, 2))
|
||||
cv2.fillPoly(text_region_mask, polygon, 1)
|
||||
|
||||
return text_region_mask
|
||||
|
||||
def generate_effective_mask(self, mask_size: tuple, polygons_ignore):
|
||||
"""Generate effective mask by setting the ineffective regions to 0 and
|
||||
effective regions to 1.
|
||||
|
||||
Args:
|
||||
mask_size (tuple): The mask size.
|
||||
polygons_ignore (list[[ndarray]]: The list of ignored text
|
||||
polygons.
|
||||
|
||||
Returns:
|
||||
mask (ndarray): The effective mask of (height, width).
|
||||
"""
|
||||
|
||||
mask = np.ones(mask_size, dtype=np.uint8)
|
||||
|
||||
for poly in polygons_ignore:
|
||||
instance = poly.reshape(-1, 2).astype(np.int32).reshape(1, -1, 2)
|
||||
cv2.fillPoly(mask, instance, 0)
|
||||
|
||||
return mask
|
||||
|
||||
def generate_level_targets(self, img_size, text_polys, ignore_polys):
|
||||
"""Generate ground truth target on each level.
|
||||
|
||||
Args:
|
||||
img_size (list[int]): Shape of input image.
|
||||
text_polys (list[list[ndarray]]): A list of ground truth polygons.
|
||||
ignore_polys (list[list[ndarray]]): A list of ignored polygons.
|
||||
Returns:
|
||||
level_maps (list(ndarray)): A list of ground target on each level.
|
||||
"""
|
||||
h, w = img_size
|
||||
lv_size_divs = self.level_size_divisors
|
||||
lv_proportion_range = self.level_proportion_range
|
||||
lv_text_polys = [[] for i in range(len(lv_size_divs))]
|
||||
lv_ignore_polys = [[] for i in range(len(lv_size_divs))]
|
||||
level_maps = []
|
||||
for poly in text_polys:
|
||||
polygon = np.array(poly, dtype=np.int32).reshape((1, -1, 2))
|
||||
_, _, box_w, box_h = cv2.boundingRect(polygon)
|
||||
proportion = max(box_h, box_w) / (h + 1e-8)
|
||||
|
||||
for ind, proportion_range in enumerate(lv_proportion_range):
|
||||
if proportion_range[0] < proportion < proportion_range[1]:
|
||||
lv_text_polys[ind].append(poly / lv_size_divs[ind])
|
||||
|
||||
for ignore_poly in ignore_polys:
|
||||
polygon = np.array(ignore_poly, dtype=np.int32).reshape((1, -1, 2))
|
||||
_, _, box_w, box_h = cv2.boundingRect(polygon)
|
||||
proportion = max(box_h, box_w) / (h + 1e-8)
|
||||
|
||||
for ind, proportion_range in enumerate(lv_proportion_range):
|
||||
if proportion_range[0] < proportion < proportion_range[1]:
|
||||
lv_ignore_polys[ind].append(ignore_poly / lv_size_divs[ind])
|
||||
|
||||
for ind, size_divisor in enumerate(lv_size_divs):
|
||||
current_level_maps = []
|
||||
level_img_size = (h // size_divisor, w // size_divisor)
|
||||
|
||||
text_region = self.generate_text_region_mask(
|
||||
level_img_size, lv_text_polys[ind]
|
||||
)[None]
|
||||
current_level_maps.append(text_region)
|
||||
|
||||
center_region = self.generate_center_region_mask(
|
||||
level_img_size, lv_text_polys[ind]
|
||||
)[None]
|
||||
current_level_maps.append(center_region)
|
||||
|
||||
effective_mask = self.generate_effective_mask(
|
||||
level_img_size, lv_ignore_polys[ind]
|
||||
)[None]
|
||||
current_level_maps.append(effective_mask)
|
||||
|
||||
fourier_real_map, fourier_image_maps = self.generate_fourier_maps(
|
||||
level_img_size, lv_text_polys[ind]
|
||||
)
|
||||
current_level_maps.append(fourier_real_map)
|
||||
current_level_maps.append(fourier_image_maps)
|
||||
|
||||
level_maps.append(np.concatenate(current_level_maps))
|
||||
|
||||
return level_maps
|
||||
|
||||
def generate_targets(self, results):
|
||||
"""Generate the ground truth targets for FCENet.
|
||||
|
||||
Args:
|
||||
results (dict): The input result dictionary.
|
||||
|
||||
Returns:
|
||||
results (dict): The output result dictionary.
|
||||
"""
|
||||
|
||||
assert isinstance(results, dict)
|
||||
image = results["image"]
|
||||
polygons = results["polys"]
|
||||
ignore_tags = results["ignore_tags"]
|
||||
h, w, _ = image.shape
|
||||
|
||||
polygon_masks = []
|
||||
polygon_masks_ignore = []
|
||||
for tag, polygon in zip(ignore_tags, polygons):
|
||||
if tag is True:
|
||||
polygon_masks_ignore.append(polygon)
|
||||
else:
|
||||
polygon_masks.append(polygon)
|
||||
|
||||
level_maps = self.generate_level_targets(
|
||||
(h, w), polygon_masks, polygon_masks_ignore
|
||||
)
|
||||
|
||||
mapping = {
|
||||
"p3_maps": level_maps[0],
|
||||
"p4_maps": level_maps[1],
|
||||
"p5_maps": level_maps[2],
|
||||
}
|
||||
for key, value in mapping.items():
|
||||
results[key] = value
|
||||
|
||||
return results
|
||||
|
||||
def __call__(self, results):
|
||||
results = self.generate_targets(results)
|
||||
return results
|
||||
BIN
ppocr/data/imaug/frost_img/frost1.jpg
Normal file
BIN
ppocr/data/imaug/frost_img/frost1.jpg
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 273 KiB |
BIN
ppocr/data/imaug/frost_img/frost2.png
Normal file
BIN
ppocr/data/imaug/frost_img/frost2.png
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 292 KiB |
BIN
ppocr/data/imaug/frost_img/frost3.png
Normal file
BIN
ppocr/data/imaug/frost_img/frost3.png
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 292 KiB |
BIN
ppocr/data/imaug/frost_img/frost4.jpg
Normal file
BIN
ppocr/data/imaug/frost_img/frost4.jpg
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 36 KiB |
BIN
ppocr/data/imaug/frost_img/frost5.jpg
Normal file
BIN
ppocr/data/imaug/frost_img/frost5.jpg
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 152 KiB |
BIN
ppocr/data/imaug/frost_img/frost6.jpg
Normal file
BIN
ppocr/data/imaug/frost_img/frost6.jpg
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 88 KiB |
213
ppocr/data/imaug/iaa_augment.py
Normal file
213
ppocr/data/imaug/iaa_augment.py
Normal file
@@ -0,0 +1,213 @@
|
||||
# copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
"""
|
||||
This code is refer from:
|
||||
https://github.com/WenmuZhou/DBNet.pytorch/blob/master/data_loader/modules/iaa_augment.py
|
||||
"""
|
||||
import os
|
||||
|
||||
# Prevent automatic updates in Albumentations for stability in augmentation behavior
|
||||
os.environ["NO_ALBUMENTATIONS_UPDATE"] = "1"
|
||||
|
||||
import numpy as np
|
||||
import albumentations as A
|
||||
from albumentations.core.transforms_interface import DualTransform
|
||||
from albumentations.augmentations.geometric import functional as fgeometric
|
||||
from packaging import version
|
||||
|
||||
ALBU_VERSION = version.parse(A.__version__)
|
||||
IS_ALBU_NEW_VERSION = ALBU_VERSION >= version.parse("1.4.15")
|
||||
|
||||
|
||||
# Custom resize transformation mimicking Imgaug's behavior with scaling
|
||||
class ImgaugLikeResize(DualTransform):
|
||||
def __init__(self, scale_range=(0.5, 3.0), interpolation=1, p=1.0):
|
||||
super(ImgaugLikeResize, self).__init__(p)
|
||||
self.scale_range = scale_range
|
||||
self.interpolation = interpolation
|
||||
|
||||
# Resize the image based on a randomly chosen scale within the scale range
|
||||
def apply(self, img, scale=1.0, **params):
|
||||
height, width = img.shape[:2]
|
||||
new_height = int(height * scale)
|
||||
new_width = int(width * scale)
|
||||
|
||||
if IS_ALBU_NEW_VERSION:
|
||||
return fgeometric.resize(
|
||||
img, (new_height, new_width), interpolation=self.interpolation
|
||||
)
|
||||
return fgeometric.resize(
|
||||
img, new_height, new_width, interpolation=self.interpolation
|
||||
)
|
||||
|
||||
# Apply the same scaling transformation to keypoints (e.g., polygon points)
|
||||
def apply_to_keypoints(self, keypoints, scale=1.0, **params):
|
||||
return np.array(
|
||||
[(x * scale, y * scale) + tuple(rest) for x, y, *rest in keypoints]
|
||||
)
|
||||
|
||||
# Get random scale parameter within the specified range
|
||||
def get_params(self):
|
||||
scale = np.random.uniform(self.scale_range[0], self.scale_range[1])
|
||||
return {"scale": scale}
|
||||
|
||||
|
||||
# Builder class to translate custom augmenter arguments into Albumentations-compatible format
|
||||
class AugmenterBuilder(object):
|
||||
def __init__(self):
|
||||
# Map common Imgaug transformations to equivalent Albumentations transforms
|
||||
self.imgaug_to_albu = {
|
||||
"Fliplr": "HorizontalFlip",
|
||||
"Flipud": "VerticalFlip",
|
||||
"Affine": "Affine",
|
||||
# Additional mappings can be added here if needed
|
||||
}
|
||||
|
||||
# Recursive method to construct augmentation pipeline based on provided arguments
|
||||
def build(self, args, root=True):
|
||||
if args is None or len(args) == 0:
|
||||
return None
|
||||
elif isinstance(args, list):
|
||||
# Build the full augmentation sequence if it's a root-level call
|
||||
if root:
|
||||
sequence = [self.build(value, root=False) for value in args]
|
||||
return A.Compose(
|
||||
sequence,
|
||||
keypoint_params=A.KeypointParams(
|
||||
format="xy", remove_invisible=False
|
||||
),
|
||||
)
|
||||
else:
|
||||
# Build individual augmenters for nested arguments
|
||||
augmenter_type = args[0]
|
||||
augmenter_args = args[1] if len(args) > 1 else {}
|
||||
augmenter_args_mapped = self.map_arguments(
|
||||
augmenter_type, augmenter_args
|
||||
)
|
||||
augmenter_type_mapped = self.imgaug_to_albu.get(
|
||||
augmenter_type, augmenter_type
|
||||
)
|
||||
if augmenter_type_mapped == "Resize":
|
||||
return ImgaugLikeResize(**augmenter_args_mapped)
|
||||
else:
|
||||
cls = getattr(A, augmenter_type_mapped)
|
||||
return cls(
|
||||
**{
|
||||
k: self.to_tuple_if_list(v)
|
||||
for k, v in augmenter_args_mapped.items()
|
||||
}
|
||||
)
|
||||
elif isinstance(args, dict):
|
||||
# Process individual transformation specified as dictionary
|
||||
augmenter_type = args["type"]
|
||||
augmenter_args = args.get("args", {})
|
||||
augmenter_args_mapped = self.map_arguments(augmenter_type, augmenter_args)
|
||||
augmenter_type_mapped = self.imgaug_to_albu.get(
|
||||
augmenter_type, augmenter_type
|
||||
)
|
||||
if augmenter_type_mapped == "Resize":
|
||||
return ImgaugLikeResize(**augmenter_args_mapped)
|
||||
else:
|
||||
cls = getattr(A, augmenter_type_mapped)
|
||||
return cls(
|
||||
**{
|
||||
k: self.to_tuple_if_list(v)
|
||||
for k, v in augmenter_args_mapped.items()
|
||||
}
|
||||
)
|
||||
else:
|
||||
raise RuntimeError("Unknown augmenter arg: " + str(args))
|
||||
|
||||
# Map arguments to expected format for each augmenter type
|
||||
def map_arguments(self, augmenter_type, augmenter_args):
|
||||
augmenter_args = augmenter_args.copy() # Avoid modifying the original arguments
|
||||
if augmenter_type == "Resize":
|
||||
# Ensure size is a valid 2-element list or tuple
|
||||
size = augmenter_args.get("size")
|
||||
if size:
|
||||
if not isinstance(size, (list, tuple)) or len(size) != 2:
|
||||
raise ValueError(
|
||||
f"'size' must be a list or tuple of two numbers, but got {size}"
|
||||
)
|
||||
min_scale, max_scale = size
|
||||
return {
|
||||
"scale_range": (min_scale, max_scale),
|
||||
"interpolation": 1, # Linear interpolation
|
||||
"p": 1.0,
|
||||
}
|
||||
else:
|
||||
return {"scale_range": (1.0, 1.0), "interpolation": 1, "p": 1.0}
|
||||
elif augmenter_type == "Affine":
|
||||
# Map rotation to a tuple and ensure p=1.0 to apply transformation
|
||||
rotate = augmenter_args.get("rotate", 0)
|
||||
if isinstance(rotate, list):
|
||||
rotate = tuple(rotate)
|
||||
elif isinstance(rotate, (int, float)):
|
||||
rotate = (float(rotate), float(rotate))
|
||||
augmenter_args["rotate"] = rotate
|
||||
augmenter_args["p"] = 1.0
|
||||
return augmenter_args
|
||||
else:
|
||||
# For other augmenters, ensure 'p' probability is specified
|
||||
p = augmenter_args.get("p", 1.0)
|
||||
augmenter_args["p"] = p
|
||||
return augmenter_args
|
||||
|
||||
# Convert lists to tuples for Albumentations compatibility
|
||||
def to_tuple_if_list(self, obj):
|
||||
if isinstance(obj, list):
|
||||
return tuple(obj)
|
||||
return obj
|
||||
|
||||
|
||||
# Wrapper class for image and polygon transformations using Imgaug-style augmentation
|
||||
class IaaAugment:
|
||||
def __init__(self, augmenter_args=None, **kwargs):
|
||||
if augmenter_args is None:
|
||||
# Default augmenters if none are specified
|
||||
augmenter_args = [
|
||||
{"type": "Fliplr", "args": {"p": 0.5}},
|
||||
{"type": "Affine", "args": {"rotate": [-10, 10]}},
|
||||
{"type": "Resize", "args": {"size": [0.5, 3]}},
|
||||
]
|
||||
self.augmenter = AugmenterBuilder().build(augmenter_args)
|
||||
|
||||
# Apply the augmentations to image and polygon data
|
||||
def __call__(self, data):
|
||||
image = data["image"]
|
||||
|
||||
if self.augmenter:
|
||||
# Flatten polygons to individual keypoints for transformation
|
||||
keypoints = []
|
||||
keypoints_lengths = []
|
||||
for poly in data["polys"]:
|
||||
keypoints.extend([tuple(point) for point in poly])
|
||||
keypoints_lengths.append(len(poly))
|
||||
|
||||
# Apply the augmentation pipeline to image and keypoints
|
||||
transformed = self.augmenter(image=image, keypoints=keypoints)
|
||||
data["image"] = transformed["image"]
|
||||
|
||||
# Extract transformed keypoints and reconstruct polygon structures
|
||||
transformed_keypoints = transformed["keypoints"]
|
||||
|
||||
# Reassemble polygons from transformed keypoints
|
||||
new_polys = []
|
||||
idx = 0
|
||||
for length in keypoints_lengths:
|
||||
new_poly = transformed_keypoints[idx : idx + length]
|
||||
new_polys.append(np.array([kp[:2] for kp in new_poly]))
|
||||
idx += length
|
||||
data["polys"] = np.array(new_polys)
|
||||
return data
|
||||
2193
ppocr/data/imaug/label_ops.py
Normal file
2193
ppocr/data/imaug/label_ops.py
Normal file
File diff suppressed because it is too large
Load Diff
183
ppocr/data/imaug/latex_ocr_aug.py
Normal file
183
ppocr/data/imaug/latex_ocr_aug.py
Normal file
@@ -0,0 +1,183 @@
|
||||
# copyright (c) 2024 PaddlePaddle Authors. All Rights Reserve.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
"""
|
||||
This code is refer from:
|
||||
https://github.com/lukas-blecher/LaTeX-OCR/blob/main/pix2tex/dataset/transforms.py
|
||||
"""
|
||||
|
||||
from __future__ import absolute_import
|
||||
from __future__ import division
|
||||
from __future__ import print_function
|
||||
from __future__ import unicode_literals
|
||||
|
||||
import os
|
||||
|
||||
os.environ["NO_ALBUMENTATIONS_UPDATE"] = "1"
|
||||
|
||||
import math
|
||||
import cv2
|
||||
import numpy as np
|
||||
import albumentations as A
|
||||
from PIL import Image
|
||||
|
||||
|
||||
class LatexTrainTransform:
|
||||
def __init__(self, bitmap_prob=0.04, **kwargs):
|
||||
# your init code
|
||||
self.bitmap_prob = bitmap_prob
|
||||
self.train_transform = A.Compose(
|
||||
[
|
||||
A.Compose(
|
||||
[
|
||||
A.ShiftScaleRotate(
|
||||
shift_limit=0,
|
||||
scale_limit=(-0.15, 0),
|
||||
rotate_limit=1,
|
||||
border_mode=0,
|
||||
interpolation=3,
|
||||
value=[255, 255, 255],
|
||||
p=1,
|
||||
),
|
||||
A.GridDistortion(
|
||||
distort_limit=0.1,
|
||||
border_mode=0,
|
||||
interpolation=3,
|
||||
value=[255, 255, 255],
|
||||
p=0.5,
|
||||
),
|
||||
],
|
||||
p=0.15,
|
||||
),
|
||||
A.RGBShift(r_shift_limit=15, g_shift_limit=15, b_shift_limit=15, p=0.3),
|
||||
A.GaussNoise(10, p=0.2),
|
||||
A.RandomBrightnessContrast(0.05, (-0.2, 0), True, p=0.2),
|
||||
A.ImageCompression(95, p=0.3),
|
||||
A.ToGray(always_apply=True),
|
||||
]
|
||||
)
|
||||
|
||||
def __call__(self, data):
|
||||
img = data["image"]
|
||||
if np.random.random() < self.bitmap_prob:
|
||||
img[img != 255] = 0
|
||||
img = self.train_transform(image=img)["image"]
|
||||
data["image"] = img
|
||||
return data
|
||||
|
||||
|
||||
class LatexTestTransform:
|
||||
def __init__(self, **kwargs):
|
||||
# your init code
|
||||
self.test_transform = A.Compose(
|
||||
[
|
||||
A.ToGray(always_apply=True),
|
||||
]
|
||||
)
|
||||
|
||||
def __call__(self, data):
|
||||
img = data["image"]
|
||||
img = self.test_transform(image=img)["image"]
|
||||
data["image"] = img
|
||||
return data
|
||||
|
||||
|
||||
class MinMaxResize:
|
||||
def __init__(self, min_dimensions=[32, 32], max_dimensions=[672, 192], **kwargs):
|
||||
# your init code
|
||||
self.min_dimensions = min_dimensions
|
||||
self.max_dimensions = max_dimensions
|
||||
# pass
|
||||
|
||||
def pad_(self, img, divable=32):
|
||||
threshold = 128
|
||||
data = np.array(img.convert("LA"))
|
||||
if data[..., -1].var() == 0:
|
||||
data = (data[..., 0]).astype(np.uint8)
|
||||
else:
|
||||
data = (255 - data[..., -1]).astype(np.uint8)
|
||||
data = (data - data.min()) / (data.max() - data.min()) * 255
|
||||
if data.mean() > threshold:
|
||||
# To invert the text to white
|
||||
gray = 255 * (data < threshold).astype(np.uint8)
|
||||
else:
|
||||
gray = 255 * (data > threshold).astype(np.uint8)
|
||||
data = 255 - data
|
||||
|
||||
coords = cv2.findNonZero(gray) # Find all non-zero points (text)
|
||||
a, b, w, h = cv2.boundingRect(coords) # Find minimum spanning bounding box
|
||||
rect = data[b : b + h, a : a + w]
|
||||
im = Image.fromarray(rect).convert("L")
|
||||
dims = []
|
||||
for x in [w, h]:
|
||||
div, mod = divmod(x, divable)
|
||||
dims.append(divable * (div + (1 if mod > 0 else 0)))
|
||||
padded = Image.new("L", dims, 255)
|
||||
padded.paste(im, (0, 0, im.size[0], im.size[1]))
|
||||
return padded
|
||||
|
||||
def minmax_size_(self, img, max_dimensions, min_dimensions):
|
||||
if max_dimensions is not None:
|
||||
ratios = [a / b for a, b in zip(img.size, max_dimensions)]
|
||||
if any([r > 1 for r in ratios]):
|
||||
size = np.array(img.size) // max(ratios)
|
||||
img = img.resize(tuple(size.astype(int)), Image.BILINEAR)
|
||||
if min_dimensions is not None:
|
||||
# hypothesis: there is a dim in img smaller than min_dimensions, and return a proper dim >= min_dimensions
|
||||
padded_size = [
|
||||
max(img_dim, min_dim)
|
||||
for img_dim, min_dim in zip(img.size, min_dimensions)
|
||||
]
|
||||
if padded_size != list(img.size): # assert hypothesis
|
||||
padded_im = Image.new("L", padded_size, 255)
|
||||
padded_im.paste(img, img.getbbox())
|
||||
img = padded_im
|
||||
return img
|
||||
|
||||
def __call__(self, data):
|
||||
img = data["image"]
|
||||
h, w = img.shape[:2]
|
||||
if (
|
||||
self.min_dimensions[0] <= w <= self.max_dimensions[0]
|
||||
and self.min_dimensions[1] <= h <= self.max_dimensions[1]
|
||||
):
|
||||
return data
|
||||
else:
|
||||
im = Image.fromarray(np.uint8(img))
|
||||
im = self.minmax_size_(
|
||||
self.pad_(im), self.max_dimensions, self.min_dimensions
|
||||
)
|
||||
im = np.array(im)
|
||||
im = np.dstack((im, im, im))
|
||||
data["image"] = im
|
||||
return data
|
||||
|
||||
|
||||
class LatexImageFormat:
|
||||
def __init__(self, **kwargs):
|
||||
# your init code
|
||||
pass
|
||||
|
||||
def __call__(self, data):
|
||||
img = data["image"]
|
||||
im_h, im_w = img.shape[:2]
|
||||
divide_h = math.ceil(im_h / 16) * 16
|
||||
divide_w = math.ceil(im_w / 16) * 16
|
||||
img = img[:, :, 0]
|
||||
img = np.pad(
|
||||
img, ((0, divide_h - im_h), (0, divide_w - im_w)), constant_values=(1, 1)
|
||||
)
|
||||
img_expanded = img[:, :, np.newaxis].transpose(2, 0, 1)
|
||||
data["image"] = img_expanded
|
||||
return data
|
||||
179
ppocr/data/imaug/make_border_map.py
Normal file
179
ppocr/data/imaug/make_border_map.py
Normal file
@@ -0,0 +1,179 @@
|
||||
# copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
"""
|
||||
This code is refer from:
|
||||
https://github.com/WenmuZhou/DBNet.pytorch/blob/master/data_loader/modules/make_border_map.py
|
||||
"""
|
||||
|
||||
from __future__ import absolute_import
|
||||
from __future__ import division
|
||||
from __future__ import print_function
|
||||
from __future__ import unicode_literals
|
||||
|
||||
import numpy as np
|
||||
import cv2
|
||||
|
||||
np.seterr(divide="ignore", invalid="ignore")
|
||||
import pyclipper
|
||||
from shapely.geometry import Polygon
|
||||
import sys
|
||||
import warnings
|
||||
|
||||
warnings.simplefilter("ignore")
|
||||
|
||||
__all__ = ["MakeBorderMap"]
|
||||
|
||||
|
||||
class MakeBorderMap(object):
|
||||
def __init__(self, shrink_ratio=0.4, thresh_min=0.3, thresh_max=0.7, **kwargs):
|
||||
self.shrink_ratio = shrink_ratio
|
||||
self.thresh_min = thresh_min
|
||||
self.thresh_max = thresh_max
|
||||
if "total_epoch" in kwargs and "epoch" in kwargs and kwargs["epoch"] != "None":
|
||||
self.shrink_ratio = self.shrink_ratio + 0.2 * kwargs["epoch"] / float(
|
||||
kwargs["total_epoch"]
|
||||
)
|
||||
|
||||
def __call__(self, data):
|
||||
img = data["image"]
|
||||
text_polys = data["polys"]
|
||||
ignore_tags = data["ignore_tags"]
|
||||
|
||||
canvas = np.zeros(img.shape[:2], dtype=np.float32)
|
||||
mask = np.zeros(img.shape[:2], dtype=np.float32)
|
||||
|
||||
for i in range(len(text_polys)):
|
||||
if ignore_tags[i]:
|
||||
continue
|
||||
self.draw_border_map(text_polys[i], canvas, mask=mask)
|
||||
canvas = canvas * (self.thresh_max - self.thresh_min) + self.thresh_min
|
||||
|
||||
data["threshold_map"] = canvas
|
||||
data["threshold_mask"] = mask
|
||||
return data
|
||||
|
||||
def draw_border_map(self, polygon, canvas, mask):
|
||||
polygon = np.array(polygon)
|
||||
assert polygon.ndim == 2
|
||||
assert polygon.shape[1] == 2
|
||||
|
||||
polygon_shape = Polygon(polygon)
|
||||
if polygon_shape.area <= 0:
|
||||
return
|
||||
distance = (
|
||||
polygon_shape.area
|
||||
* (1 - np.power(self.shrink_ratio, 2))
|
||||
/ polygon_shape.length
|
||||
)
|
||||
subject = [tuple(l) for l in polygon]
|
||||
padding = pyclipper.PyclipperOffset()
|
||||
padding.AddPath(subject, pyclipper.JT_ROUND, pyclipper.ET_CLOSEDPOLYGON)
|
||||
|
||||
padded_polygon = np.array(padding.Execute(distance)[0])
|
||||
cv2.fillPoly(mask, [padded_polygon.astype(np.int32)], 1.0)
|
||||
|
||||
xmin = padded_polygon[:, 0].min()
|
||||
xmax = padded_polygon[:, 0].max()
|
||||
ymin = padded_polygon[:, 1].min()
|
||||
ymax = padded_polygon[:, 1].max()
|
||||
width = xmax - xmin + 1
|
||||
height = ymax - ymin + 1
|
||||
|
||||
polygon[:, 0] = polygon[:, 0] - xmin
|
||||
polygon[:, 1] = polygon[:, 1] - ymin
|
||||
|
||||
xs = np.broadcast_to(
|
||||
np.linspace(0, width - 1, num=width).reshape(1, width), (height, width)
|
||||
)
|
||||
ys = np.broadcast_to(
|
||||
np.linspace(0, height - 1, num=height).reshape(height, 1), (height, width)
|
||||
)
|
||||
|
||||
distance_map = np.zeros((polygon.shape[0], height, width), dtype=np.float32)
|
||||
for i in range(polygon.shape[0]):
|
||||
j = (i + 1) % polygon.shape[0]
|
||||
absolute_distance = self._distance(xs, ys, polygon[i], polygon[j])
|
||||
distance_map[i] = np.clip(absolute_distance / distance, 0, 1)
|
||||
distance_map = distance_map.min(axis=0)
|
||||
|
||||
xmin_valid = min(max(0, xmin), canvas.shape[1] - 1)
|
||||
xmax_valid = min(max(0, xmax), canvas.shape[1] - 1)
|
||||
ymin_valid = min(max(0, ymin), canvas.shape[0] - 1)
|
||||
ymax_valid = min(max(0, ymax), canvas.shape[0] - 1)
|
||||
canvas[ymin_valid : ymax_valid + 1, xmin_valid : xmax_valid + 1] = np.fmax(
|
||||
1
|
||||
- distance_map[
|
||||
ymin_valid - ymin : ymax_valid - ymax + height,
|
||||
xmin_valid - xmin : xmax_valid - xmax + width,
|
||||
],
|
||||
canvas[ymin_valid : ymax_valid + 1, xmin_valid : xmax_valid + 1],
|
||||
)
|
||||
|
||||
def _distance(self, xs, ys, point_1, point_2):
|
||||
"""
|
||||
compute the distance from point to a line
|
||||
ys: coordinates in the first axis
|
||||
xs: coordinates in the second axis
|
||||
point_1, point_2: (x, y), the end of the line
|
||||
"""
|
||||
height, width = xs.shape[:2]
|
||||
square_distance_1 = np.square(xs - point_1[0]) + np.square(ys - point_1[1])
|
||||
square_distance_2 = np.square(xs - point_2[0]) + np.square(ys - point_2[1])
|
||||
square_distance = np.square(point_1[0] - point_2[0]) + np.square(
|
||||
point_1[1] - point_2[1]
|
||||
)
|
||||
|
||||
cosin = (square_distance - square_distance_1 - square_distance_2) / (
|
||||
2 * np.sqrt(square_distance_1 * square_distance_2)
|
||||
)
|
||||
square_sin = 1 - np.square(cosin)
|
||||
square_sin = np.nan_to_num(square_sin)
|
||||
result = np.sqrt(
|
||||
square_distance_1 * square_distance_2 * square_sin / square_distance
|
||||
)
|
||||
|
||||
result[cosin < 0] = np.sqrt(np.fmin(square_distance_1, square_distance_2))[
|
||||
cosin < 0
|
||||
]
|
||||
# self.extend_line(point_1, point_2, result)
|
||||
return result
|
||||
|
||||
def extend_line(self, point_1, point_2, result, shrink_ratio):
|
||||
ex_point_1 = (
|
||||
int(round(point_1[0] + (point_1[0] - point_2[0]) * (1 + shrink_ratio))),
|
||||
int(round(point_1[1] + (point_1[1] - point_2[1]) * (1 + shrink_ratio))),
|
||||
)
|
||||
cv2.line(
|
||||
result,
|
||||
tuple(ex_point_1),
|
||||
tuple(point_1),
|
||||
4096.0,
|
||||
1,
|
||||
lineType=cv2.LINE_AA,
|
||||
shift=0,
|
||||
)
|
||||
ex_point_2 = (
|
||||
int(round(point_2[0] + (point_2[0] - point_1[0]) * (1 + shrink_ratio))),
|
||||
int(round(point_2[1] + (point_2[1] - point_1[1]) * (1 + shrink_ratio))),
|
||||
)
|
||||
cv2.line(
|
||||
result,
|
||||
tuple(ex_point_2),
|
||||
tuple(point_2),
|
||||
4096.0,
|
||||
1,
|
||||
lineType=cv2.LINE_AA,
|
||||
shift=0,
|
||||
)
|
||||
return ex_point_1, ex_point_2
|
||||
104
ppocr/data/imaug/make_pse_gt.py
Normal file
104
ppocr/data/imaug/make_pse_gt.py
Normal file
@@ -0,0 +1,104 @@
|
||||
# copyright (c) 2021 PaddlePaddle Authors. All Rights Reserve.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from __future__ import absolute_import
|
||||
from __future__ import division
|
||||
from __future__ import print_function
|
||||
from __future__ import unicode_literals
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
import pyclipper
|
||||
from shapely.geometry import Polygon
|
||||
|
||||
__all__ = ["MakePseGt"]
|
||||
|
||||
|
||||
class MakePseGt(object):
|
||||
def __init__(self, kernel_num=7, size=640, min_shrink_ratio=0.4, **kwargs):
|
||||
self.kernel_num = kernel_num
|
||||
self.min_shrink_ratio = min_shrink_ratio
|
||||
self.size = size
|
||||
|
||||
def __call__(self, data):
|
||||
image = data["image"]
|
||||
text_polys = data["polys"]
|
||||
ignore_tags = data["ignore_tags"]
|
||||
|
||||
h, w, _ = image.shape
|
||||
short_edge = min(h, w)
|
||||
if short_edge < self.size:
|
||||
# keep short_size >= self.size
|
||||
scale = self.size / short_edge
|
||||
image = cv2.resize(image, dsize=None, fx=scale, fy=scale)
|
||||
text_polys *= scale
|
||||
|
||||
gt_kernels = []
|
||||
for i in range(1, self.kernel_num + 1):
|
||||
# s1->sn, from big to small
|
||||
rate = 1.0 - (1.0 - self.min_shrink_ratio) / (self.kernel_num - 1) * i
|
||||
text_kernel, ignore_tags = self.generate_kernel(
|
||||
image.shape[0:2], rate, text_polys, ignore_tags
|
||||
)
|
||||
gt_kernels.append(text_kernel)
|
||||
|
||||
training_mask = np.ones(image.shape[0:2], dtype="uint8")
|
||||
for i in range(text_polys.shape[0]):
|
||||
if ignore_tags[i]:
|
||||
cv2.fillPoly(
|
||||
training_mask, text_polys[i].astype(np.int32)[np.newaxis, :, :], 0
|
||||
)
|
||||
|
||||
gt_kernels = np.array(gt_kernels)
|
||||
gt_kernels[gt_kernels > 0] = 1
|
||||
|
||||
data["image"] = image
|
||||
data["polys"] = text_polys
|
||||
data["gt_kernels"] = gt_kernels[0:]
|
||||
data["gt_text"] = gt_kernels[0]
|
||||
data["mask"] = training_mask.astype("float32")
|
||||
return data
|
||||
|
||||
def generate_kernel(self, img_size, shrink_ratio, text_polys, ignore_tags=None):
|
||||
"""
|
||||
Refer to part of the code:
|
||||
https://github.com/open-mmlab/mmocr/blob/main/mmocr/datasets/pipelines/textdet_targets/base_textdet_targets.py
|
||||
"""
|
||||
|
||||
h, w = img_size
|
||||
text_kernel = np.zeros((h, w), dtype=np.float32)
|
||||
for i, poly in enumerate(text_polys):
|
||||
polygon = Polygon(poly)
|
||||
distance = (
|
||||
polygon.area
|
||||
* (1 - shrink_ratio * shrink_ratio)
|
||||
/ (polygon.length + 1e-6)
|
||||
)
|
||||
subject = [tuple(l) for l in poly]
|
||||
pco = pyclipper.PyclipperOffset()
|
||||
pco.AddPath(subject, pyclipper.JT_ROUND, pyclipper.ET_CLOSEDPOLYGON)
|
||||
shrunk = np.array(pco.Execute(-distance))
|
||||
|
||||
if len(shrunk) == 0 or shrunk.size == 0:
|
||||
if ignore_tags is not None:
|
||||
ignore_tags[i] = True
|
||||
continue
|
||||
try:
|
||||
shrunk = np.array(shrunk[0]).reshape(-1, 2)
|
||||
except:
|
||||
if ignore_tags is not None:
|
||||
ignore_tags[i] = True
|
||||
continue
|
||||
cv2.fillPoly(text_kernel, [shrunk.astype(np.int32)], i + 1)
|
||||
return text_kernel, ignore_tags
|
||||
125
ppocr/data/imaug/make_shrink_map.py
Normal file
125
ppocr/data/imaug/make_shrink_map.py
Normal file
@@ -0,0 +1,125 @@
|
||||
# copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
"""
|
||||
This code is refer from:
|
||||
https://github.com/WenmuZhou/DBNet.pytorch/blob/master/data_loader/modules/make_shrink_map.py
|
||||
"""
|
||||
|
||||
from __future__ import absolute_import
|
||||
from __future__ import division
|
||||
from __future__ import print_function
|
||||
from __future__ import unicode_literals
|
||||
|
||||
import numpy as np
|
||||
import cv2
|
||||
from shapely.geometry import Polygon
|
||||
import pyclipper
|
||||
|
||||
__all__ = ["MakeShrinkMap"]
|
||||
|
||||
|
||||
class MakeShrinkMap(object):
|
||||
r"""
|
||||
Making binary mask from detection data with ICDAR format.
|
||||
Typically following the process of class `MakeICDARData`.
|
||||
"""
|
||||
|
||||
def __init__(self, min_text_size=8, shrink_ratio=0.4, **kwargs):
|
||||
self.min_text_size = min_text_size
|
||||
self.shrink_ratio = shrink_ratio
|
||||
if "total_epoch" in kwargs and "epoch" in kwargs and kwargs["epoch"] != "None":
|
||||
self.shrink_ratio = self.shrink_ratio + 0.2 * kwargs["epoch"] / float(
|
||||
kwargs["total_epoch"]
|
||||
)
|
||||
|
||||
def __call__(self, data):
|
||||
image = data["image"]
|
||||
text_polys = data["polys"]
|
||||
ignore_tags = data["ignore_tags"]
|
||||
|
||||
h, w = image.shape[:2]
|
||||
text_polys, ignore_tags = self.validate_polygons(text_polys, ignore_tags, h, w)
|
||||
gt = np.zeros((h, w), dtype=np.float32)
|
||||
mask = np.ones((h, w), dtype=np.float32)
|
||||
for i in range(len(text_polys)):
|
||||
polygon = text_polys[i]
|
||||
height = max(polygon[:, 1]) - min(polygon[:, 1])
|
||||
width = max(polygon[:, 0]) - min(polygon[:, 0])
|
||||
if ignore_tags[i] or min(height, width) < self.min_text_size:
|
||||
cv2.fillPoly(mask, polygon.astype(np.int32)[np.newaxis, :, :], 0)
|
||||
ignore_tags[i] = True
|
||||
else:
|
||||
polygon_shape = Polygon(polygon)
|
||||
subject = [tuple(l) for l in polygon]
|
||||
padding = pyclipper.PyclipperOffset()
|
||||
padding.AddPath(subject, pyclipper.JT_ROUND, pyclipper.ET_CLOSEDPOLYGON)
|
||||
shrunk = []
|
||||
|
||||
# Increase the shrink ratio every time we get multiple polygon returned back
|
||||
possible_ratios = np.arange(self.shrink_ratio, 1, self.shrink_ratio)
|
||||
np.append(possible_ratios, 1)
|
||||
# print(possible_ratios)
|
||||
for ratio in possible_ratios:
|
||||
# print(f"Change shrink ratio to {ratio}")
|
||||
distance = (
|
||||
polygon_shape.area
|
||||
* (1 - np.power(ratio, 2))
|
||||
/ polygon_shape.length
|
||||
)
|
||||
shrunk = padding.Execute(-distance)
|
||||
if len(shrunk) == 1:
|
||||
break
|
||||
|
||||
if shrunk == []:
|
||||
cv2.fillPoly(mask, polygon.astype(np.int32)[np.newaxis, :, :], 0)
|
||||
ignore_tags[i] = True
|
||||
continue
|
||||
|
||||
for each_shrink in shrunk:
|
||||
shrink = np.array(each_shrink).reshape(-1, 2)
|
||||
cv2.fillPoly(gt, [shrink.astype(np.int32)], 1)
|
||||
|
||||
data["shrink_map"] = gt
|
||||
data["shrink_mask"] = mask
|
||||
return data
|
||||
|
||||
def validate_polygons(self, polygons, ignore_tags, h, w):
|
||||
"""
|
||||
polygons (numpy.array, required): of shape (num_instances, num_points, 2)
|
||||
"""
|
||||
if len(polygons) == 0:
|
||||
return polygons, ignore_tags
|
||||
assert len(polygons) == len(ignore_tags)
|
||||
for polygon in polygons:
|
||||
polygon[:, 0] = np.clip(polygon[:, 0], 0, w - 1)
|
||||
polygon[:, 1] = np.clip(polygon[:, 1], 0, h - 1)
|
||||
|
||||
for i in range(len(polygons)):
|
||||
area = self.polygon_area(polygons[i])
|
||||
if abs(area) < 1:
|
||||
ignore_tags[i] = True
|
||||
if area > 0:
|
||||
polygons[i] = polygons[i][::-1, :]
|
||||
return polygons, ignore_tags
|
||||
|
||||
def polygon_area(self, polygon):
|
||||
"""
|
||||
compute polygon area
|
||||
"""
|
||||
area = 0
|
||||
q = polygon[-1]
|
||||
for p in polygon:
|
||||
area += p[0] * q[1] - p[1] * q[0]
|
||||
q = p
|
||||
return area / 2.0
|
||||
523
ppocr/data/imaug/operators.py
Normal file
523
ppocr/data/imaug/operators.py
Normal file
@@ -0,0 +1,523 @@
|
||||
"""
|
||||
# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
"""
|
||||
|
||||
from __future__ import absolute_import
|
||||
from __future__ import division
|
||||
from __future__ import print_function
|
||||
from __future__ import unicode_literals
|
||||
|
||||
import sys
|
||||
import cv2
|
||||
import numpy as np
|
||||
import math
|
||||
from PIL import Image
|
||||
|
||||
|
||||
class DecodeImage(object):
|
||||
"""decode image"""
|
||||
|
||||
def __init__(
|
||||
self, img_mode="RGB", channel_first=False, ignore_orientation=False, **kwargs
|
||||
):
|
||||
self.img_mode = img_mode
|
||||
self.channel_first = channel_first
|
||||
self.ignore_orientation = ignore_orientation
|
||||
|
||||
def __call__(self, data):
|
||||
img = data["image"]
|
||||
assert type(img) is bytes and len(img) > 0, "invalid input 'img' in DecodeImage"
|
||||
img = np.frombuffer(img, dtype="uint8")
|
||||
if self.ignore_orientation:
|
||||
img = cv2.imdecode(img, cv2.IMREAD_IGNORE_ORIENTATION | cv2.IMREAD_COLOR)
|
||||
else:
|
||||
img = cv2.imdecode(img, 1)
|
||||
if img is None:
|
||||
return None
|
||||
if self.img_mode == "GRAY":
|
||||
img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
|
||||
elif self.img_mode == "RGB":
|
||||
assert img.shape[2] == 3, "invalid shape of image[%s]" % (img.shape)
|
||||
img = img[:, :, ::-1]
|
||||
|
||||
if self.channel_first:
|
||||
img = img.transpose((2, 0, 1))
|
||||
|
||||
data["image"] = img
|
||||
return data
|
||||
|
||||
|
||||
class NormalizeImage(object):
|
||||
"""normalize image such as subtract mean, divide std"""
|
||||
|
||||
def __init__(self, scale=None, mean=None, std=None, order="chw", **kwargs):
|
||||
if isinstance(scale, str):
|
||||
scale = eval(scale)
|
||||
self.scale = np.float32(scale if scale is not None else 1.0 / 255.0)
|
||||
mean = mean if mean is not None else [0.485, 0.456, 0.406]
|
||||
std = std if std is not None else [0.229, 0.224, 0.225]
|
||||
|
||||
shape = (3, 1, 1) if order == "chw" else (1, 1, 3)
|
||||
self.mean = np.array(mean).reshape(shape).astype("float32")
|
||||
self.std = np.array(std).reshape(shape).astype("float32")
|
||||
|
||||
def __call__(self, data):
|
||||
img = data["image"]
|
||||
from PIL import Image
|
||||
|
||||
if isinstance(img, Image.Image):
|
||||
img = np.array(img)
|
||||
assert isinstance(img, np.ndarray), "invalid input 'img' in NormalizeImage"
|
||||
data["image"] = (img.astype("float32") * self.scale - self.mean) / self.std
|
||||
return data
|
||||
|
||||
|
||||
class ToCHWImage(object):
|
||||
"""convert hwc image to chw image"""
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
pass
|
||||
|
||||
def __call__(self, data):
|
||||
img = data["image"]
|
||||
from PIL import Image
|
||||
|
||||
if isinstance(img, Image.Image):
|
||||
img = np.array(img)
|
||||
data["image"] = img.transpose((2, 0, 1))
|
||||
return data
|
||||
|
||||
|
||||
class Fasttext(object):
|
||||
def __init__(self, path="None", **kwargs):
|
||||
import fasttext
|
||||
|
||||
self.fast_model = fasttext.load_model(path)
|
||||
|
||||
def __call__(self, data):
|
||||
label = data["label"]
|
||||
fast_label = self.fast_model[label]
|
||||
data["fast_label"] = fast_label
|
||||
return data
|
||||
|
||||
|
||||
class KeepKeys(object):
|
||||
def __init__(self, keep_keys, **kwargs):
|
||||
self.keep_keys = keep_keys
|
||||
|
||||
def __call__(self, data):
|
||||
data_list = []
|
||||
for key in self.keep_keys:
|
||||
data_list.append(data[key])
|
||||
return data_list
|
||||
|
||||
|
||||
class Pad(object):
|
||||
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"]
|
||||
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_h,
|
||||
0,
|
||||
resize_w2 - img_w,
|
||||
cv2.BORDER_CONSTANT,
|
||||
value=0,
|
||||
)
|
||||
data["image"] = img
|
||||
return data
|
||||
|
||||
|
||||
class Resize(object):
|
||||
def __init__(self, size=(640, 640), **kwargs):
|
||||
self.size = size
|
||||
|
||||
def resize_image(self, img):
|
||||
resize_h, resize_w = self.size
|
||||
ori_h, ori_w = img.shape[:2] # (h, w, c)
|
||||
ratio_h = float(resize_h) / ori_h
|
||||
ratio_w = float(resize_w) / ori_w
|
||||
img = cv2.resize(img, (int(resize_w), int(resize_h)))
|
||||
return img, [ratio_h, ratio_w]
|
||||
|
||||
def __call__(self, data):
|
||||
img = data["image"]
|
||||
if "polys" in data:
|
||||
text_polys = data["polys"]
|
||||
|
||||
img_resize, [ratio_h, ratio_w] = self.resize_image(img)
|
||||
if "polys" in data:
|
||||
new_boxes = []
|
||||
for box in text_polys:
|
||||
new_box = []
|
||||
for cord in box:
|
||||
new_box.append([cord[0] * ratio_w, cord[1] * ratio_h])
|
||||
new_boxes.append(new_box)
|
||||
data["polys"] = np.array(new_boxes, dtype=np.float32)
|
||||
data["image"] = img_resize
|
||||
return data
|
||||
|
||||
|
||||
class DetResizeForTest(object):
|
||||
def __init__(self, **kwargs):
|
||||
super(DetResizeForTest, self).__init__()
|
||||
self.resize_type = 0
|
||||
self.keep_ratio = False
|
||||
if "image_shape" in kwargs:
|
||||
self.image_shape = kwargs["image_shape"]
|
||||
self.resize_type = 1
|
||||
if "keep_ratio" in kwargs:
|
||||
self.keep_ratio = kwargs["keep_ratio"]
|
||||
elif "limit_side_len" in kwargs:
|
||||
self.limit_side_len = kwargs["limit_side_len"]
|
||||
self.limit_type = kwargs.get("limit_type", "min")
|
||||
elif "resize_long" in kwargs:
|
||||
self.resize_type = 2
|
||||
self.resize_long = kwargs.get("resize_long", 960)
|
||||
else:
|
||||
self.limit_side_len = 736
|
||||
self.limit_type = "min"
|
||||
|
||||
def __call__(self, data):
|
||||
img = data["image"]
|
||||
src_h, src_w, _ = img.shape
|
||||
if sum([src_h, src_w]) < 64:
|
||||
img = self.image_padding(img)
|
||||
|
||||
if self.resize_type == 0:
|
||||
# img, shape = self.resize_image_type0(img)
|
||||
img, [ratio_h, ratio_w] = self.resize_image_type0(img)
|
||||
elif self.resize_type == 2:
|
||||
img, [ratio_h, ratio_w] = self.resize_image_type2(img)
|
||||
else:
|
||||
# img, shape = self.resize_image_type1(img)
|
||||
img, [ratio_h, ratio_w] = self.resize_image_type1(img)
|
||||
data["image"] = img
|
||||
data["shape"] = np.array([src_h, src_w, ratio_h, ratio_w])
|
||||
return data
|
||||
|
||||
def image_padding(self, im, value=0):
|
||||
h, w, c = im.shape
|
||||
im_pad = np.zeros((max(32, h), max(32, w), c), np.uint8) + value
|
||||
im_pad[:h, :w, :] = im
|
||||
return im_pad
|
||||
|
||||
def resize_image_type1(self, img):
|
||||
resize_h, resize_w = self.image_shape
|
||||
ori_h, ori_w = img.shape[:2] # (h, w, c)
|
||||
if self.keep_ratio is True:
|
||||
resize_w = ori_w * resize_h / ori_h
|
||||
N = math.ceil(resize_w / 32)
|
||||
resize_w = N * 32
|
||||
ratio_h = float(resize_h) / ori_h
|
||||
ratio_w = float(resize_w) / ori_w
|
||||
img = cv2.resize(img, (int(resize_w), int(resize_h)))
|
||||
# return img, np.array([ori_h, ori_w])
|
||||
return img, [ratio_h, ratio_w]
|
||||
|
||||
def resize_image_type0(self, img):
|
||||
"""
|
||||
resize image to a size multiple of 32 which is required by the network
|
||||
args:
|
||||
img(array): array with shape [h, w, c]
|
||||
return(tuple):
|
||||
img, (ratio_h, ratio_w)
|
||||
"""
|
||||
limit_side_len = self.limit_side_len
|
||||
h, w, c = img.shape
|
||||
|
||||
# limit the max side
|
||||
if self.limit_type == "max":
|
||||
if max(h, w) > limit_side_len:
|
||||
if h > w:
|
||||
ratio = float(limit_side_len) / h
|
||||
else:
|
||||
ratio = float(limit_side_len) / w
|
||||
else:
|
||||
ratio = 1.0
|
||||
elif self.limit_type == "min":
|
||||
if min(h, w) < limit_side_len:
|
||||
if h < w:
|
||||
ratio = float(limit_side_len) / h
|
||||
else:
|
||||
ratio = float(limit_side_len) / w
|
||||
else:
|
||||
ratio = 1.0
|
||||
elif self.limit_type == "resize_long":
|
||||
ratio = float(limit_side_len) / max(h, w)
|
||||
else:
|
||||
raise Exception("not support limit type, image ")
|
||||
resize_h = int(h * ratio)
|
||||
resize_w = int(w * ratio)
|
||||
|
||||
resize_h = max(int(round(resize_h / 32) * 32), 32)
|
||||
resize_w = max(int(round(resize_w / 32) * 32), 32)
|
||||
|
||||
try:
|
||||
if int(resize_w) <= 0 or int(resize_h) <= 0:
|
||||
return None, (None, None)
|
||||
img = cv2.resize(img, (int(resize_w), int(resize_h)))
|
||||
except:
|
||||
print(img.shape, resize_w, resize_h)
|
||||
sys.exit(0)
|
||||
ratio_h = resize_h / float(h)
|
||||
ratio_w = resize_w / float(w)
|
||||
return img, [ratio_h, ratio_w]
|
||||
|
||||
def resize_image_type2(self, img):
|
||||
h, w, _ = img.shape
|
||||
|
||||
resize_w = w
|
||||
resize_h = h
|
||||
|
||||
if resize_h > resize_w:
|
||||
ratio = float(self.resize_long) / resize_h
|
||||
else:
|
||||
ratio = float(self.resize_long) / resize_w
|
||||
|
||||
resize_h = int(resize_h * ratio)
|
||||
resize_w = int(resize_w * ratio)
|
||||
|
||||
max_stride = 128
|
||||
resize_h = (resize_h + max_stride - 1) // max_stride * max_stride
|
||||
resize_w = (resize_w + max_stride - 1) // max_stride * max_stride
|
||||
img = cv2.resize(img, (int(resize_w), int(resize_h)))
|
||||
ratio_h = resize_h / float(h)
|
||||
ratio_w = resize_w / float(w)
|
||||
|
||||
return img, [ratio_h, ratio_w]
|
||||
|
||||
|
||||
class E2EResizeForTest(object):
|
||||
def __init__(self, **kwargs):
|
||||
super(E2EResizeForTest, self).__init__()
|
||||
self.max_side_len = kwargs["max_side_len"]
|
||||
self.valid_set = kwargs["valid_set"]
|
||||
|
||||
def __call__(self, data):
|
||||
img = data["image"]
|
||||
src_h, src_w, _ = img.shape
|
||||
if self.valid_set == "totaltext":
|
||||
im_resized, [ratio_h, ratio_w] = self.resize_image_for_totaltext(
|
||||
img, max_side_len=self.max_side_len
|
||||
)
|
||||
else:
|
||||
im_resized, (ratio_h, ratio_w) = self.resize_image(
|
||||
img, max_side_len=self.max_side_len
|
||||
)
|
||||
data["image"] = im_resized
|
||||
data["shape"] = np.array([src_h, src_w, ratio_h, ratio_w])
|
||||
return data
|
||||
|
||||
def resize_image_for_totaltext(self, im, max_side_len=512):
|
||||
h, w, _ = im.shape
|
||||
resize_w = w
|
||||
resize_h = h
|
||||
ratio = 1.25
|
||||
if h * ratio > max_side_len:
|
||||
ratio = float(max_side_len) / resize_h
|
||||
resize_h = int(resize_h * ratio)
|
||||
resize_w = int(resize_w * ratio)
|
||||
|
||||
max_stride = 128
|
||||
resize_h = (resize_h + max_stride - 1) // max_stride * max_stride
|
||||
resize_w = (resize_w + max_stride - 1) // max_stride * max_stride
|
||||
im = cv2.resize(im, (int(resize_w), int(resize_h)))
|
||||
ratio_h = resize_h / float(h)
|
||||
ratio_w = resize_w / float(w)
|
||||
return im, (ratio_h, ratio_w)
|
||||
|
||||
def resize_image(self, im, max_side_len=512):
|
||||
"""
|
||||
resize image to a size multiple of max_stride which is required by the network
|
||||
:param im: the resized image
|
||||
:param max_side_len: limit of max image size to avoid out of memory in gpu
|
||||
:return: the resized image and the resize ratio
|
||||
"""
|
||||
h, w, _ = im.shape
|
||||
|
||||
resize_w = w
|
||||
resize_h = h
|
||||
|
||||
# Fix the longer side
|
||||
if resize_h > resize_w:
|
||||
ratio = float(max_side_len) / resize_h
|
||||
else:
|
||||
ratio = float(max_side_len) / resize_w
|
||||
|
||||
resize_h = int(resize_h * ratio)
|
||||
resize_w = int(resize_w * ratio)
|
||||
|
||||
max_stride = 128
|
||||
resize_h = (resize_h + max_stride - 1) // max_stride * max_stride
|
||||
resize_w = (resize_w + max_stride - 1) // max_stride * max_stride
|
||||
im = cv2.resize(im, (int(resize_w), int(resize_h)))
|
||||
ratio_h = resize_h / float(h)
|
||||
ratio_w = resize_w / float(w)
|
||||
|
||||
return im, (ratio_h, ratio_w)
|
||||
|
||||
|
||||
class KieResize(object):
|
||||
def __init__(self, **kwargs):
|
||||
super(KieResize, self).__init__()
|
||||
self.max_side, self.min_side = kwargs["img_scale"][0], kwargs["img_scale"][1]
|
||||
|
||||
def __call__(self, data):
|
||||
img = data["image"]
|
||||
points = data["points"]
|
||||
src_h, src_w, _ = img.shape
|
||||
(
|
||||
im_resized,
|
||||
scale_factor,
|
||||
[ratio_h, ratio_w],
|
||||
[new_h, new_w],
|
||||
) = self.resize_image(img)
|
||||
resize_points = self.resize_boxes(img, points, scale_factor)
|
||||
data["ori_image"] = img
|
||||
data["ori_boxes"] = points
|
||||
data["points"] = resize_points
|
||||
data["image"] = im_resized
|
||||
data["shape"] = np.array([new_h, new_w])
|
||||
return data
|
||||
|
||||
def resize_image(self, img):
|
||||
norm_img = np.zeros([1024, 1024, 3], dtype="float32")
|
||||
scale = [512, 1024]
|
||||
h, w = img.shape[:2]
|
||||
max_long_edge = max(scale)
|
||||
max_short_edge = min(scale)
|
||||
scale_factor = min(max_long_edge / max(h, w), max_short_edge / min(h, w))
|
||||
resize_w, resize_h = int(w * float(scale_factor) + 0.5), int(
|
||||
h * float(scale_factor) + 0.5
|
||||
)
|
||||
max_stride = 32
|
||||
resize_h = (resize_h + max_stride - 1) // max_stride * max_stride
|
||||
resize_w = (resize_w + max_stride - 1) // max_stride * max_stride
|
||||
im = cv2.resize(img, (resize_w, resize_h))
|
||||
new_h, new_w = im.shape[:2]
|
||||
w_scale = new_w / w
|
||||
h_scale = new_h / h
|
||||
scale_factor = np.array([w_scale, h_scale, w_scale, h_scale], dtype=np.float32)
|
||||
norm_img[:new_h, :new_w, :] = im
|
||||
return norm_img, scale_factor, [h_scale, w_scale], [new_h, new_w]
|
||||
|
||||
def resize_boxes(self, im, points, scale_factor):
|
||||
points = points * scale_factor
|
||||
img_shape = im.shape[:2]
|
||||
points[:, 0::2] = np.clip(points[:, 0::2], 0, img_shape[1])
|
||||
points[:, 1::2] = np.clip(points[:, 1::2], 0, img_shape[0])
|
||||
return points
|
||||
|
||||
|
||||
class SRResize(object):
|
||||
def __init__(
|
||||
self,
|
||||
imgH=32,
|
||||
imgW=128,
|
||||
down_sample_scale=4,
|
||||
keep_ratio=False,
|
||||
min_ratio=1,
|
||||
mask=False,
|
||||
infer_mode=False,
|
||||
**kwargs,
|
||||
):
|
||||
self.imgH = imgH
|
||||
self.imgW = imgW
|
||||
self.keep_ratio = keep_ratio
|
||||
self.min_ratio = min_ratio
|
||||
self.down_sample_scale = down_sample_scale
|
||||
self.mask = mask
|
||||
self.infer_mode = infer_mode
|
||||
|
||||
def __call__(self, data):
|
||||
imgH = self.imgH
|
||||
imgW = self.imgW
|
||||
images_lr = data["image_lr"]
|
||||
transform2 = ResizeNormalize(
|
||||
(imgW // self.down_sample_scale, imgH // self.down_sample_scale)
|
||||
)
|
||||
images_lr = transform2(images_lr)
|
||||
data["img_lr"] = images_lr
|
||||
if self.infer_mode:
|
||||
return data
|
||||
|
||||
images_HR = data["image_hr"]
|
||||
label_strs = data["label"]
|
||||
transform = ResizeNormalize((imgW, imgH))
|
||||
images_HR = transform(images_HR)
|
||||
data["img_hr"] = images_HR
|
||||
return data
|
||||
|
||||
|
||||
class ResizeNormalize(object):
|
||||
def __init__(self, size, interpolation=Image.BICUBIC):
|
||||
self.size = size
|
||||
self.interpolation = interpolation
|
||||
|
||||
def __call__(self, img):
|
||||
img = img.resize(self.size, self.interpolation)
|
||||
img_numpy = np.array(img).astype("float32")
|
||||
img_numpy = img_numpy.transpose((2, 0, 1)) / 255
|
||||
return img_numpy
|
||||
|
||||
|
||||
class GrayImageChannelFormat(object):
|
||||
"""
|
||||
format gray scale image's channel: (3,h,w) -> (1,h,w)
|
||||
Args:
|
||||
inverse: inverse gray image
|
||||
"""
|
||||
|
||||
def __init__(self, inverse=False, **kwargs):
|
||||
self.inverse = inverse
|
||||
|
||||
def __call__(self, data):
|
||||
img = data["image"]
|
||||
img_single_channel = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
|
||||
img_expanded = np.expand_dims(img_single_channel, 0)
|
||||
|
||||
if self.inverse:
|
||||
data["image"] = np.abs(img_expanded - 1)
|
||||
else:
|
||||
data["image"] = img_expanded
|
||||
|
||||
data["src_image"] = img
|
||||
return data
|
||||
1116
ppocr/data/imaug/pg_process.py
Normal file
1116
ppocr/data/imaug/pg_process.py
Normal file
File diff suppressed because it is too large
Load Diff
134
ppocr/data/imaug/randaugment.py
Normal file
134
ppocr/data/imaug/randaugment.py
Normal file
@@ -0,0 +1,134 @@
|
||||
# copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from __future__ import absolute_import
|
||||
from __future__ import division
|
||||
from __future__ import print_function
|
||||
from __future__ import unicode_literals
|
||||
|
||||
from PIL import Image, ImageEnhance, ImageOps
|
||||
import numpy as np
|
||||
import random
|
||||
|
||||
|
||||
class RawRandAugment(object):
|
||||
def __init__(self, num_layers=2, magnitude=5, fillcolor=(128, 128, 128), **kwargs):
|
||||
self.num_layers = num_layers
|
||||
self.magnitude = magnitude
|
||||
self.max_level = 10
|
||||
|
||||
abso_level = self.magnitude / self.max_level
|
||||
self.level_map = {
|
||||
"shearX": 0.3 * abso_level,
|
||||
"shearY": 0.3 * abso_level,
|
||||
"translateX": 150.0 / 331 * abso_level,
|
||||
"translateY": 150.0 / 331 * abso_level,
|
||||
"rotate": 30 * abso_level,
|
||||
"color": 0.9 * abso_level,
|
||||
"posterize": int(4.0 * abso_level),
|
||||
"solarize": 256.0 * abso_level,
|
||||
"contrast": 0.9 * abso_level,
|
||||
"sharpness": 0.9 * abso_level,
|
||||
"brightness": 0.9 * abso_level,
|
||||
"autocontrast": 0,
|
||||
"equalize": 0,
|
||||
"invert": 0,
|
||||
}
|
||||
|
||||
# from https://stackoverflow.com/questions/5252170/
|
||||
# specify-image-filling-color-when-rotating-in-python-with-pil-and-setting-expand
|
||||
def rotate_with_fill(img, magnitude):
|
||||
rot = img.convert("RGBA").rotate(magnitude)
|
||||
return Image.composite(
|
||||
rot, Image.new("RGBA", rot.size, (128,) * 4), rot
|
||||
).convert(img.mode)
|
||||
|
||||
rnd_ch_op = random.choice
|
||||
|
||||
self.func = {
|
||||
"shearX": lambda img, magnitude: img.transform(
|
||||
img.size,
|
||||
Image.AFFINE,
|
||||
(1, magnitude * rnd_ch_op([-1, 1]), 0, 0, 1, 0),
|
||||
Image.BICUBIC,
|
||||
fillcolor=fillcolor,
|
||||
),
|
||||
"shearY": lambda img, magnitude: img.transform(
|
||||
img.size,
|
||||
Image.AFFINE,
|
||||
(1, 0, 0, magnitude * rnd_ch_op([-1, 1]), 1, 0),
|
||||
Image.BICUBIC,
|
||||
fillcolor=fillcolor,
|
||||
),
|
||||
"translateX": lambda img, magnitude: img.transform(
|
||||
img.size,
|
||||
Image.AFFINE,
|
||||
(1, 0, magnitude * img.size[0] * rnd_ch_op([-1, 1]), 0, 1, 0),
|
||||
fillcolor=fillcolor,
|
||||
),
|
||||
"translateY": lambda img, magnitude: img.transform(
|
||||
img.size,
|
||||
Image.AFFINE,
|
||||
(1, 0, 0, 0, 1, magnitude * img.size[1] * rnd_ch_op([-1, 1])),
|
||||
fillcolor=fillcolor,
|
||||
),
|
||||
"rotate": lambda img, magnitude: rotate_with_fill(img, magnitude),
|
||||
"color": lambda img, magnitude: ImageEnhance.Color(img).enhance(
|
||||
1 + magnitude * rnd_ch_op([-1, 1])
|
||||
),
|
||||
"posterize": lambda img, magnitude: ImageOps.posterize(img, magnitude),
|
||||
"solarize": lambda img, magnitude: ImageOps.solarize(img, magnitude),
|
||||
"contrast": lambda img, magnitude: ImageEnhance.Contrast(img).enhance(
|
||||
1 + magnitude * rnd_ch_op([-1, 1])
|
||||
),
|
||||
"sharpness": lambda img, magnitude: ImageEnhance.Sharpness(img).enhance(
|
||||
1 + magnitude * rnd_ch_op([-1, 1])
|
||||
),
|
||||
"brightness": lambda img, magnitude: ImageEnhance.Brightness(img).enhance(
|
||||
1 + magnitude * rnd_ch_op([-1, 1])
|
||||
),
|
||||
"autocontrast": lambda img, magnitude: ImageOps.autocontrast(img),
|
||||
"equalize": lambda img, magnitude: ImageOps.equalize(img),
|
||||
"invert": lambda img, magnitude: ImageOps.invert(img),
|
||||
}
|
||||
|
||||
def __call__(self, img):
|
||||
avaiable_op_names = list(self.level_map.keys())
|
||||
for layer_num in range(self.num_layers):
|
||||
op_name = np.random.choice(avaiable_op_names)
|
||||
img = self.func[op_name](img, self.level_map[op_name])
|
||||
return img
|
||||
|
||||
|
||||
class RandAugment(RawRandAugment):
|
||||
"""RandAugment wrapper to auto fit different img types"""
|
||||
|
||||
def __init__(self, prob=0.5, *args, **kwargs):
|
||||
self.prob = prob
|
||||
super().__init__(*args, **kwargs)
|
||||
|
||||
def __call__(self, data):
|
||||
if np.random.rand() > self.prob:
|
||||
return data
|
||||
img = data["image"]
|
||||
if not isinstance(img, Image.Image):
|
||||
img = np.ascontiguousarray(img)
|
||||
img = Image.fromarray(img)
|
||||
|
||||
img = super().__call__(img)
|
||||
|
||||
if isinstance(img, Image.Image):
|
||||
img = np.asarray(img)
|
||||
data["image"] = img
|
||||
return data
|
||||
238
ppocr/data/imaug/random_crop_data.py
Normal file
238
ppocr/data/imaug/random_crop_data.py
Normal file
@@ -0,0 +1,238 @@
|
||||
# copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
"""
|
||||
This code is refer from:
|
||||
https://github.com/WenmuZhou/DBNet.pytorch/blob/master/data_loader/modules/random_crop_data.py
|
||||
"""
|
||||
|
||||
from __future__ import absolute_import
|
||||
from __future__ import division
|
||||
from __future__ import print_function
|
||||
from __future__ import unicode_literals
|
||||
|
||||
import numpy as np
|
||||
import cv2
|
||||
import random
|
||||
|
||||
|
||||
def is_poly_in_rect(poly, x, y, w, h):
|
||||
poly = np.array(poly)
|
||||
if poly[:, 0].min() < x or poly[:, 0].max() > x + w:
|
||||
return False
|
||||
if poly[:, 1].min() < y or poly[:, 1].max() > y + h:
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def is_poly_outside_rect(poly, x, y, w, h):
|
||||
poly = np.array(poly)
|
||||
if poly[:, 0].max() < x or poly[:, 0].min() > x + w:
|
||||
return True
|
||||
if poly[:, 1].max() < y or poly[:, 1].min() > y + h:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def split_regions(axis):
|
||||
regions = []
|
||||
min_axis = 0
|
||||
for i in range(1, axis.shape[0]):
|
||||
if axis[i] != axis[i - 1] + 1:
|
||||
region = axis[min_axis:i]
|
||||
min_axis = i
|
||||
regions.append(region)
|
||||
return regions
|
||||
|
||||
|
||||
def random_select(axis, max_size):
|
||||
xx = np.random.choice(axis, size=2)
|
||||
xmin = np.min(xx)
|
||||
xmax = np.max(xx)
|
||||
xmin = np.clip(xmin, 0, max_size - 1)
|
||||
xmax = np.clip(xmax, 0, max_size - 1)
|
||||
return xmin, xmax
|
||||
|
||||
|
||||
def region_wise_random_select(regions, max_size):
|
||||
selected_index = list(np.random.choice(len(regions), 2))
|
||||
selected_values = []
|
||||
for index in selected_index:
|
||||
axis = regions[index]
|
||||
xx = int(np.random.choice(axis, size=1))
|
||||
selected_values.append(xx)
|
||||
xmin = min(selected_values)
|
||||
xmax = max(selected_values)
|
||||
return xmin, xmax
|
||||
|
||||
|
||||
def crop_area(im, text_polys, min_crop_side_ratio, max_tries):
|
||||
h, w, _ = im.shape
|
||||
h_array = np.zeros(h, dtype=np.int32)
|
||||
w_array = np.zeros(w, dtype=np.int32)
|
||||
for points in text_polys:
|
||||
points = np.round(points, decimals=0).astype(np.int32)
|
||||
minx = np.min(points[:, 0])
|
||||
maxx = np.max(points[:, 0])
|
||||
w_array[minx:maxx] = 1
|
||||
miny = np.min(points[:, 1])
|
||||
maxy = np.max(points[:, 1])
|
||||
h_array[miny:maxy] = 1
|
||||
# ensure the cropped area not across a text
|
||||
h_axis = np.where(h_array == 0)[0]
|
||||
w_axis = np.where(w_array == 0)[0]
|
||||
|
||||
if len(h_axis) == 0 or len(w_axis) == 0:
|
||||
return 0, 0, w, h
|
||||
|
||||
h_regions = split_regions(h_axis)
|
||||
w_regions = split_regions(w_axis)
|
||||
|
||||
for i in range(max_tries):
|
||||
if len(w_regions) > 1:
|
||||
xmin, xmax = region_wise_random_select(w_regions, w)
|
||||
else:
|
||||
xmin, xmax = random_select(w_axis, w)
|
||||
if len(h_regions) > 1:
|
||||
ymin, ymax = region_wise_random_select(h_regions, h)
|
||||
else:
|
||||
ymin, ymax = random_select(h_axis, h)
|
||||
|
||||
if (
|
||||
xmax - xmin < min_crop_side_ratio * w
|
||||
or ymax - ymin < min_crop_side_ratio * h
|
||||
):
|
||||
# area too small
|
||||
continue
|
||||
num_poly_in_rect = 0
|
||||
for poly in text_polys:
|
||||
if not is_poly_outside_rect(poly, xmin, ymin, xmax - xmin, ymax - ymin):
|
||||
num_poly_in_rect += 1
|
||||
break
|
||||
|
||||
if num_poly_in_rect > 0:
|
||||
return xmin, ymin, xmax - xmin, ymax - ymin
|
||||
|
||||
return 0, 0, w, h
|
||||
|
||||
|
||||
class EastRandomCropData(object):
|
||||
def __init__(
|
||||
self,
|
||||
size=(640, 640),
|
||||
max_tries=10,
|
||||
min_crop_side_ratio=0.1,
|
||||
keep_ratio=True,
|
||||
**kwargs,
|
||||
):
|
||||
self.size = size
|
||||
self.max_tries = max_tries
|
||||
self.min_crop_side_ratio = min_crop_side_ratio
|
||||
self.keep_ratio = keep_ratio
|
||||
|
||||
def __call__(self, data):
|
||||
img = data["image"]
|
||||
text_polys = data["polys"]
|
||||
ignore_tags = data["ignore_tags"]
|
||||
texts = data["texts"]
|
||||
all_care_polys = [text_polys[i] for i, tag in enumerate(ignore_tags) if not tag]
|
||||
# 计算crop区域
|
||||
crop_x, crop_y, crop_w, crop_h = crop_area(
|
||||
img, all_care_polys, self.min_crop_side_ratio, self.max_tries
|
||||
)
|
||||
# crop 图片 保持比例填充
|
||||
scale_w = self.size[0] / crop_w
|
||||
scale_h = self.size[1] / crop_h
|
||||
scale = min(scale_w, scale_h)
|
||||
h = int(crop_h * scale)
|
||||
w = int(crop_w * scale)
|
||||
if self.keep_ratio:
|
||||
padimg = np.zeros((self.size[1], self.size[0], img.shape[2]), img.dtype)
|
||||
padimg[:h, :w] = cv2.resize(
|
||||
img[crop_y : crop_y + crop_h, crop_x : crop_x + crop_w], (w, h)
|
||||
)
|
||||
img = padimg
|
||||
else:
|
||||
img = cv2.resize(
|
||||
img[crop_y : crop_y + crop_h, crop_x : crop_x + crop_w],
|
||||
tuple(self.size),
|
||||
)
|
||||
# crop 文本框
|
||||
text_polys_crop = []
|
||||
ignore_tags_crop = []
|
||||
texts_crop = []
|
||||
for poly, text, tag in zip(text_polys, texts, ignore_tags):
|
||||
poly = ((poly - (crop_x, crop_y)) * scale).tolist()
|
||||
if not is_poly_outside_rect(poly, 0, 0, w, h):
|
||||
text_polys_crop.append(poly)
|
||||
ignore_tags_crop.append(tag)
|
||||
texts_crop.append(text)
|
||||
data["image"] = img
|
||||
data["polys"] = np.array(text_polys_crop)
|
||||
data["ignore_tags"] = ignore_tags_crop
|
||||
data["texts"] = texts_crop
|
||||
return data
|
||||
|
||||
|
||||
class RandomCropImgMask(object):
|
||||
def __init__(self, size, main_key, crop_keys, p=3 / 8, **kwargs):
|
||||
self.size = size
|
||||
self.main_key = main_key
|
||||
self.crop_keys = crop_keys
|
||||
self.p = p
|
||||
|
||||
def __call__(self, data):
|
||||
image = data["image"]
|
||||
|
||||
h, w = image.shape[0:2]
|
||||
th, tw = self.size
|
||||
if w == tw and h == th:
|
||||
return data
|
||||
|
||||
mask = data[self.main_key]
|
||||
if np.max(mask) > 0 and random.random() > self.p:
|
||||
# make sure to crop the text region
|
||||
tl = np.min(np.where(mask > 0), axis=1) - (th, tw)
|
||||
tl[tl < 0] = 0
|
||||
br = np.max(np.where(mask > 0), axis=1) - (th, tw)
|
||||
br[br < 0] = 0
|
||||
|
||||
br[0] = min(br[0], h - th)
|
||||
br[1] = min(br[1], w - tw)
|
||||
|
||||
i = random.randint(tl[0], br[0]) if tl[0] < br[0] else 0
|
||||
j = random.randint(tl[1], br[1]) if tl[1] < br[1] else 0
|
||||
else:
|
||||
i = random.randint(0, h - th) if h - th > 0 else 0
|
||||
j = random.randint(0, w - tw) if w - tw > 0 else 0
|
||||
|
||||
# return i, j, th, tw
|
||||
for k in data:
|
||||
if k in self.crop_keys:
|
||||
if len(data[k].shape) == 3:
|
||||
if np.argmin(data[k].shape) == 0:
|
||||
img = data[k][:, i : i + th, j : j + tw]
|
||||
if img.shape[1] != img.shape[2]:
|
||||
a = 1
|
||||
elif np.argmin(data[k].shape) == 2:
|
||||
img = data[k][i : i + th, j : j + tw, :]
|
||||
if img.shape[1] != img.shape[0]:
|
||||
a = 1
|
||||
else:
|
||||
img = data[k]
|
||||
else:
|
||||
img = data[k][i : i + th, j : j + tw]
|
||||
if img.shape[0] != img.shape[1]:
|
||||
a = 1
|
||||
data[k] = img
|
||||
return data
|
||||
932
ppocr/data/imaug/rec_img_aug.py
Normal file
932
ppocr/data/imaug/rec_img_aug.py
Normal file
@@ -0,0 +1,932 @@
|
||||
# copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import math
|
||||
import cv2
|
||||
import numpy as np
|
||||
import random
|
||||
import copy
|
||||
from PIL import Image
|
||||
import PIL
|
||||
from .text_image_aug import tia_perspective, tia_stretch, tia_distort
|
||||
from .abinet_aug import (
|
||||
CVGeometry,
|
||||
CVDeterioration,
|
||||
CVColorJitter,
|
||||
SVTRGeometry,
|
||||
SVTRDeterioration,
|
||||
ParseQDeterioration,
|
||||
)
|
||||
from paddle.vision.transforms import Compose
|
||||
|
||||
|
||||
class RecAug(object):
|
||||
def __init__(
|
||||
self,
|
||||
tia_prob=0.4,
|
||||
crop_prob=0.4,
|
||||
reverse_prob=0.4,
|
||||
noise_prob=0.4,
|
||||
jitter_prob=0.4,
|
||||
blur_prob=0.4,
|
||||
hsv_aug_prob=0.4,
|
||||
**kwargs,
|
||||
):
|
||||
self.tia_prob = tia_prob
|
||||
self.bda = BaseDataAugmentation(
|
||||
crop_prob, reverse_prob, noise_prob, jitter_prob, blur_prob, hsv_aug_prob
|
||||
)
|
||||
|
||||
def __call__(self, data):
|
||||
img = data["image"]
|
||||
h, w, _ = img.shape
|
||||
|
||||
# tia
|
||||
if random.random() <= self.tia_prob:
|
||||
if h >= 20 and w >= 20:
|
||||
img = tia_distort(img, random.randint(3, 6))
|
||||
img = tia_stretch(img, random.randint(3, 6))
|
||||
img = tia_perspective(img)
|
||||
|
||||
# bda
|
||||
data["image"] = img
|
||||
data = self.bda(data)
|
||||
return data
|
||||
|
||||
|
||||
class BaseDataAugmentation(object):
|
||||
def __init__(
|
||||
self,
|
||||
crop_prob=0.4,
|
||||
reverse_prob=0.4,
|
||||
noise_prob=0.4,
|
||||
jitter_prob=0.4,
|
||||
blur_prob=0.4,
|
||||
hsv_aug_prob=0.4,
|
||||
**kwargs,
|
||||
):
|
||||
self.crop_prob = crop_prob
|
||||
self.reverse_prob = reverse_prob
|
||||
self.noise_prob = noise_prob
|
||||
self.jitter_prob = jitter_prob
|
||||
self.blur_prob = blur_prob
|
||||
self.hsv_aug_prob = hsv_aug_prob
|
||||
# for GaussianBlur
|
||||
self.fil = cv2.getGaussianKernel(ksize=5, sigma=1, ktype=cv2.CV_32F)
|
||||
|
||||
def __call__(self, data):
|
||||
img = data["image"]
|
||||
h, w, _ = img.shape
|
||||
|
||||
if random.random() <= self.crop_prob and h >= 20 and w >= 20:
|
||||
img = get_crop(img)
|
||||
|
||||
if random.random() <= self.blur_prob:
|
||||
# GaussianBlur
|
||||
img = cv2.sepFilter2D(img, -1, self.fil, self.fil)
|
||||
|
||||
if random.random() <= self.hsv_aug_prob:
|
||||
img = hsv_aug(img)
|
||||
|
||||
if random.random() <= self.jitter_prob:
|
||||
img = jitter(img)
|
||||
|
||||
if random.random() <= self.noise_prob:
|
||||
img = add_gasuss_noise(img)
|
||||
|
||||
if random.random() <= self.reverse_prob:
|
||||
img = 255 - img
|
||||
|
||||
data["image"] = img
|
||||
return data
|
||||
|
||||
|
||||
class ABINetRecAug(object):
|
||||
def __init__(
|
||||
self, geometry_p=0.5, deterioration_p=0.25, colorjitter_p=0.25, **kwargs
|
||||
):
|
||||
self.transforms = Compose(
|
||||
[
|
||||
CVGeometry(
|
||||
degrees=45,
|
||||
translate=(0.0, 0.0),
|
||||
scale=(0.5, 2.0),
|
||||
shear=(45, 15),
|
||||
distortion=0.5,
|
||||
p=geometry_p,
|
||||
),
|
||||
CVDeterioration(var=20, degrees=6, factor=4, p=deterioration_p),
|
||||
CVColorJitter(
|
||||
brightness=0.5,
|
||||
contrast=0.5,
|
||||
saturation=0.5,
|
||||
hue=0.1,
|
||||
p=colorjitter_p,
|
||||
),
|
||||
]
|
||||
)
|
||||
|
||||
def __call__(self, data):
|
||||
img = data["image"]
|
||||
img = self.transforms(img)
|
||||
data["image"] = img
|
||||
return data
|
||||
|
||||
|
||||
class RecConAug(object):
|
||||
def __init__(
|
||||
self,
|
||||
prob=0.5,
|
||||
image_shape=(32, 320, 3),
|
||||
max_text_length=25,
|
||||
ext_data_num=1,
|
||||
**kwargs,
|
||||
):
|
||||
self.ext_data_num = ext_data_num
|
||||
self.prob = prob
|
||||
self.max_text_length = max_text_length
|
||||
self.image_shape = image_shape
|
||||
self.max_wh_ratio = self.image_shape[1] / self.image_shape[0]
|
||||
|
||||
def merge_ext_data(self, data, ext_data):
|
||||
ori_w = round(
|
||||
data["image"].shape[1] / data["image"].shape[0] * self.image_shape[0]
|
||||
)
|
||||
ext_w = round(
|
||||
ext_data["image"].shape[1]
|
||||
/ ext_data["image"].shape[0]
|
||||
* self.image_shape[0]
|
||||
)
|
||||
data["image"] = cv2.resize(data["image"], (ori_w, self.image_shape[0]))
|
||||
ext_data["image"] = cv2.resize(ext_data["image"], (ext_w, self.image_shape[0]))
|
||||
data["image"] = np.concatenate([data["image"], ext_data["image"]], axis=1)
|
||||
data["label"] += ext_data["label"]
|
||||
return data
|
||||
|
||||
def __call__(self, data):
|
||||
rnd_num = random.random()
|
||||
if rnd_num > self.prob:
|
||||
return data
|
||||
for idx, ext_data in enumerate(data["ext_data"]):
|
||||
if len(data["label"]) + len(ext_data["label"]) > self.max_text_length:
|
||||
break
|
||||
concat_ratio = (
|
||||
data["image"].shape[1] / data["image"].shape[0]
|
||||
+ ext_data["image"].shape[1] / ext_data["image"].shape[0]
|
||||
)
|
||||
if concat_ratio > self.max_wh_ratio:
|
||||
break
|
||||
data = self.merge_ext_data(data, ext_data)
|
||||
data.pop("ext_data")
|
||||
return data
|
||||
|
||||
|
||||
class SVTRRecAug(object):
|
||||
def __init__(
|
||||
self,
|
||||
aug_type=0,
|
||||
geometry_p=0.5,
|
||||
deterioration_p=0.25,
|
||||
colorjitter_p=0.25,
|
||||
**kwargs,
|
||||
):
|
||||
self.transforms = Compose(
|
||||
[
|
||||
SVTRGeometry(
|
||||
aug_type=aug_type,
|
||||
degrees=45,
|
||||
translate=(0.0, 0.0),
|
||||
scale=(0.5, 2.0),
|
||||
shear=(45, 15),
|
||||
distortion=0.5,
|
||||
p=geometry_p,
|
||||
),
|
||||
SVTRDeterioration(var=20, degrees=6, factor=4, p=deterioration_p),
|
||||
CVColorJitter(
|
||||
brightness=0.5,
|
||||
contrast=0.5,
|
||||
saturation=0.5,
|
||||
hue=0.1,
|
||||
p=colorjitter_p,
|
||||
),
|
||||
]
|
||||
)
|
||||
|
||||
def __call__(self, data):
|
||||
img = data["image"]
|
||||
img = self.transforms(img)
|
||||
data["image"] = img
|
||||
return data
|
||||
|
||||
|
||||
class ParseQRecAug(object):
|
||||
def __init__(
|
||||
self,
|
||||
aug_type=0,
|
||||
geometry_p=0.5,
|
||||
deterioration_p=0.25,
|
||||
colorjitter_p=0.25,
|
||||
**kwargs,
|
||||
):
|
||||
self.transforms = Compose(
|
||||
[
|
||||
SVTRGeometry(
|
||||
aug_type=aug_type,
|
||||
degrees=45,
|
||||
translate=(0.0, 0.0),
|
||||
scale=(0.5, 2.0),
|
||||
shear=(45, 15),
|
||||
distortion=0.5,
|
||||
p=geometry_p,
|
||||
),
|
||||
ParseQDeterioration(
|
||||
var=20, degrees=6, lam=20, radius=2.0, factor=4, p=deterioration_p
|
||||
),
|
||||
CVColorJitter(
|
||||
brightness=0.5,
|
||||
contrast=0.5,
|
||||
saturation=0.5,
|
||||
hue=0.1,
|
||||
p=colorjitter_p,
|
||||
),
|
||||
]
|
||||
)
|
||||
|
||||
def __call__(self, data):
|
||||
img = data["image"]
|
||||
img = self.transforms(img)
|
||||
data["image"] = img
|
||||
return data
|
||||
|
||||
|
||||
class ClsResizeImg(object):
|
||||
def __init__(self, image_shape, **kwargs):
|
||||
self.image_shape = image_shape
|
||||
|
||||
def __call__(self, data):
|
||||
img = data["image"]
|
||||
norm_img, _ = resize_norm_img(img, self.image_shape)
|
||||
data["image"] = norm_img
|
||||
return data
|
||||
|
||||
|
||||
class RecResizeImg(object):
|
||||
def __init__(
|
||||
self,
|
||||
image_shape,
|
||||
infer_mode=False,
|
||||
eval_mode=False,
|
||||
character_dict_path="./ppocr/utils/ppocr_keys_v1.txt",
|
||||
padding=True,
|
||||
**kwargs,
|
||||
):
|
||||
self.image_shape = image_shape
|
||||
self.infer_mode = infer_mode
|
||||
self.eval_mode = eval_mode
|
||||
self.character_dict_path = character_dict_path
|
||||
self.padding = padding
|
||||
|
||||
def __call__(self, data):
|
||||
img = data["image"]
|
||||
if self.eval_mode or (self.infer_mode and self.character_dict_path is not None):
|
||||
norm_img, valid_ratio = resize_norm_img_chinese(img, self.image_shape)
|
||||
else:
|
||||
norm_img, valid_ratio = resize_norm_img(img, self.image_shape, self.padding)
|
||||
data["image"] = norm_img
|
||||
data["valid_ratio"] = valid_ratio
|
||||
return data
|
||||
|
||||
|
||||
class VLRecResizeImg(object):
|
||||
def __init__(
|
||||
self,
|
||||
image_shape,
|
||||
infer_mode=False,
|
||||
character_dict_path="./ppocr/utils/ppocr_keys_v1.txt",
|
||||
padding=True,
|
||||
**kwargs,
|
||||
):
|
||||
self.image_shape = image_shape
|
||||
self.infer_mode = infer_mode
|
||||
self.character_dict_path = character_dict_path
|
||||
self.padding = padding
|
||||
|
||||
def __call__(self, data):
|
||||
img = data["image"]
|
||||
|
||||
imgC, imgH, imgW = self.image_shape
|
||||
resized_image = cv2.resize(img, (imgW, imgH), interpolation=cv2.INTER_LINEAR)
|
||||
resized_w = imgW
|
||||
resized_image = resized_image.astype("float32")
|
||||
if self.image_shape[0] == 1:
|
||||
resized_image = resized_image / 255
|
||||
norm_img = resized_image[np.newaxis, :]
|
||||
else:
|
||||
norm_img = resized_image.transpose((2, 0, 1)) / 255
|
||||
valid_ratio = min(1.0, float(resized_w / imgW))
|
||||
|
||||
data["image"] = norm_img
|
||||
data["valid_ratio"] = valid_ratio
|
||||
return data
|
||||
|
||||
|
||||
class RFLRecResizeImg(object):
|
||||
def __init__(self, image_shape, padding=True, interpolation=1, **kwargs):
|
||||
self.image_shape = image_shape
|
||||
self.padding = padding
|
||||
|
||||
self.interpolation = interpolation
|
||||
if self.interpolation == 0:
|
||||
self.interpolation = cv2.INTER_NEAREST
|
||||
elif self.interpolation == 1:
|
||||
self.interpolation = cv2.INTER_LINEAR
|
||||
elif self.interpolation == 2:
|
||||
self.interpolation = cv2.INTER_CUBIC
|
||||
elif self.interpolation == 3:
|
||||
self.interpolation = cv2.INTER_AREA
|
||||
else:
|
||||
raise Exception("Unsupported interpolation type !!!")
|
||||
|
||||
def __call__(self, data):
|
||||
img = data["image"]
|
||||
img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
|
||||
norm_img, valid_ratio = resize_norm_img(
|
||||
img, self.image_shape, self.padding, self.interpolation
|
||||
)
|
||||
data["image"] = norm_img
|
||||
data["valid_ratio"] = valid_ratio
|
||||
return data
|
||||
|
||||
|
||||
class SRNRecResizeImg(object):
|
||||
def __init__(self, image_shape, num_heads, max_text_length, **kwargs):
|
||||
self.image_shape = image_shape
|
||||
self.num_heads = num_heads
|
||||
self.max_text_length = max_text_length
|
||||
|
||||
def __call__(self, data):
|
||||
img = data["image"]
|
||||
norm_img = resize_norm_img_srn(img, self.image_shape)
|
||||
data["image"] = norm_img
|
||||
[
|
||||
encoder_word_pos,
|
||||
gsrm_word_pos,
|
||||
gsrm_slf_attn_bias1,
|
||||
gsrm_slf_attn_bias2,
|
||||
] = srn_other_inputs(self.image_shape, self.num_heads, self.max_text_length)
|
||||
|
||||
data["encoder_word_pos"] = encoder_word_pos
|
||||
data["gsrm_word_pos"] = gsrm_word_pos
|
||||
data["gsrm_slf_attn_bias1"] = gsrm_slf_attn_bias1
|
||||
data["gsrm_slf_attn_bias2"] = gsrm_slf_attn_bias2
|
||||
return data
|
||||
|
||||
|
||||
class SARRecResizeImg(object):
|
||||
def __init__(self, image_shape, width_downsample_ratio=0.25, **kwargs):
|
||||
self.image_shape = image_shape
|
||||
self.width_downsample_ratio = width_downsample_ratio
|
||||
|
||||
def __call__(self, data):
|
||||
img = data["image"]
|
||||
norm_img, resize_shape, pad_shape, valid_ratio = resize_norm_img_sar(
|
||||
img, self.image_shape, self.width_downsample_ratio
|
||||
)
|
||||
data["image"] = norm_img
|
||||
data["resized_shape"] = resize_shape
|
||||
data["pad_shape"] = pad_shape
|
||||
data["valid_ratio"] = valid_ratio
|
||||
return data
|
||||
|
||||
|
||||
class PRENResizeImg(object):
|
||||
def __init__(self, image_shape, **kwargs):
|
||||
"""
|
||||
According to original paper's realization, it's a hard resize method here.
|
||||
So maybe you should optimize it to fit for your task better.
|
||||
"""
|
||||
self.dst_h, self.dst_w = image_shape
|
||||
|
||||
def __call__(self, data):
|
||||
img = data["image"]
|
||||
resized_img = cv2.resize(
|
||||
img, (self.dst_w, self.dst_h), interpolation=cv2.INTER_LINEAR
|
||||
)
|
||||
resized_img = resized_img.transpose((2, 0, 1)) / 255
|
||||
resized_img -= 0.5
|
||||
resized_img /= 0.5
|
||||
data["image"] = resized_img.astype(np.float32)
|
||||
return data
|
||||
|
||||
|
||||
class SPINRecResizeImg(object):
|
||||
def __init__(
|
||||
self,
|
||||
image_shape,
|
||||
interpolation=2,
|
||||
mean=(127.5, 127.5, 127.5),
|
||||
std=(127.5, 127.5, 127.5),
|
||||
**kwargs,
|
||||
):
|
||||
self.image_shape = image_shape
|
||||
|
||||
self.mean = np.array(mean, dtype=np.float32)
|
||||
self.std = np.array(std, dtype=np.float32)
|
||||
self.interpolation = interpolation
|
||||
|
||||
def __call__(self, data):
|
||||
img = data["image"]
|
||||
img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
|
||||
# different interpolation type corresponding the OpenCV
|
||||
if self.interpolation == 0:
|
||||
interpolation = cv2.INTER_NEAREST
|
||||
elif self.interpolation == 1:
|
||||
interpolation = cv2.INTER_LINEAR
|
||||
elif self.interpolation == 2:
|
||||
interpolation = cv2.INTER_CUBIC
|
||||
elif self.interpolation == 3:
|
||||
interpolation = cv2.INTER_AREA
|
||||
else:
|
||||
raise Exception("Unsupported interpolation type !!!")
|
||||
# Deal with the image error during image loading
|
||||
if img is None:
|
||||
return None
|
||||
|
||||
img = cv2.resize(img, tuple(self.image_shape), interpolation)
|
||||
img = np.array(img, np.float32)
|
||||
img = np.expand_dims(img, -1)
|
||||
img = img.transpose((2, 0, 1))
|
||||
# normalize the image
|
||||
img = img.copy().astype(np.float32)
|
||||
mean = np.float64(self.mean.reshape(1, -1))
|
||||
stdinv = 1 / np.float64(self.std.reshape(1, -1))
|
||||
img -= mean
|
||||
img *= stdinv
|
||||
data["image"] = img
|
||||
return data
|
||||
|
||||
|
||||
class GrayRecResizeImg(object):
|
||||
def __init__(
|
||||
self,
|
||||
image_shape,
|
||||
resize_type,
|
||||
inter_type="Image.Resampling.LANCZOS",
|
||||
scale=True,
|
||||
padding=False,
|
||||
**kwargs,
|
||||
):
|
||||
self.image_shape = image_shape
|
||||
self.resize_type = resize_type
|
||||
self.padding = padding
|
||||
self.inter_type = eval(inter_type)
|
||||
self.scale = scale
|
||||
|
||||
def __call__(self, data):
|
||||
img = data["image"]
|
||||
img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
|
||||
image_shape = self.image_shape
|
||||
if self.padding:
|
||||
imgC, imgH, imgW = image_shape
|
||||
# todo: change to 0 and modified image shape
|
||||
h = img.shape[0]
|
||||
w = img.shape[1]
|
||||
ratio = w / float(h)
|
||||
if math.ceil(imgH * ratio) > imgW:
|
||||
resized_w = imgW
|
||||
else:
|
||||
resized_w = int(math.ceil(imgH * ratio))
|
||||
resized_image = cv2.resize(img, (resized_w, imgH))
|
||||
norm_img = np.expand_dims(resized_image, -1)
|
||||
norm_img = norm_img.transpose((2, 0, 1))
|
||||
resized_image = norm_img.astype(np.float32) / 128.0 - 1.0
|
||||
padding_im = np.zeros((imgC, imgH, imgW), dtype=np.float32)
|
||||
padding_im[:, :, 0:resized_w] = resized_image
|
||||
data["image"] = padding_im
|
||||
return data
|
||||
if self.resize_type == "PIL":
|
||||
image_pil = Image.fromarray(np.uint8(img))
|
||||
img = image_pil.resize(self.image_shape, self.inter_type)
|
||||
img = np.array(img)
|
||||
if self.resize_type == "OpenCV":
|
||||
img = cv2.resize(img, self.image_shape)
|
||||
norm_img = np.expand_dims(img, -1)
|
||||
norm_img = norm_img.transpose((2, 0, 1))
|
||||
if self.scale:
|
||||
data["image"] = norm_img.astype(np.float32) / 128.0 - 1.0
|
||||
else:
|
||||
data["image"] = norm_img.astype(np.float32) / 255.0
|
||||
return data
|
||||
|
||||
|
||||
class ABINetRecResizeImg(object):
|
||||
def __init__(self, image_shape, **kwargs):
|
||||
self.image_shape = image_shape
|
||||
|
||||
def __call__(self, data):
|
||||
img = data["image"]
|
||||
norm_img, valid_ratio = resize_norm_img_abinet(img, self.image_shape)
|
||||
data["image"] = norm_img
|
||||
data["valid_ratio"] = valid_ratio
|
||||
return data
|
||||
|
||||
|
||||
class SVTRRecResizeImg(object):
|
||||
def __init__(self, image_shape, padding=True, **kwargs):
|
||||
self.image_shape = image_shape
|
||||
self.padding = padding
|
||||
|
||||
def __call__(self, data):
|
||||
img = data["image"]
|
||||
|
||||
norm_img, valid_ratio = resize_norm_img(img, self.image_shape, self.padding)
|
||||
data["image"] = norm_img
|
||||
data["valid_ratio"] = valid_ratio
|
||||
return data
|
||||
|
||||
|
||||
class RobustScannerRecResizeImg(object):
|
||||
def __init__(
|
||||
self, image_shape, max_text_length, width_downsample_ratio=0.25, **kwargs
|
||||
):
|
||||
self.image_shape = image_shape
|
||||
self.width_downsample_ratio = width_downsample_ratio
|
||||
self.max_text_length = max_text_length
|
||||
|
||||
def __call__(self, data):
|
||||
img = data["image"]
|
||||
norm_img, resize_shape, pad_shape, valid_ratio = resize_norm_img_sar(
|
||||
img, self.image_shape, self.width_downsample_ratio
|
||||
)
|
||||
word_positons = np.array(range(0, self.max_text_length)).astype("int64")
|
||||
data["image"] = norm_img
|
||||
data["resized_shape"] = resize_shape
|
||||
data["pad_shape"] = pad_shape
|
||||
data["valid_ratio"] = valid_ratio
|
||||
data["word_positons"] = word_positons
|
||||
return data
|
||||
|
||||
|
||||
def resize_norm_img_sar(img, image_shape, width_downsample_ratio=0.25):
|
||||
imgC, imgH, imgW_min, imgW_max = image_shape
|
||||
h = img.shape[0]
|
||||
w = img.shape[1]
|
||||
valid_ratio = 1.0
|
||||
# make sure new_width is an integral multiple of width_divisor.
|
||||
width_divisor = int(1 / width_downsample_ratio)
|
||||
# resize
|
||||
ratio = w / float(h)
|
||||
resize_w = math.ceil(imgH * ratio)
|
||||
if resize_w % width_divisor != 0:
|
||||
resize_w = round(resize_w / width_divisor) * width_divisor
|
||||
if imgW_min is not None:
|
||||
resize_w = max(imgW_min, resize_w)
|
||||
if imgW_max is not None:
|
||||
valid_ratio = min(1.0, 1.0 * resize_w / imgW_max)
|
||||
resize_w = min(imgW_max, resize_w)
|
||||
resized_image = cv2.resize(img, (resize_w, imgH))
|
||||
resized_image = resized_image.astype("float32")
|
||||
# norm
|
||||
if image_shape[0] == 1:
|
||||
resized_image = resized_image / 255
|
||||
resized_image = resized_image[np.newaxis, :]
|
||||
else:
|
||||
resized_image = resized_image.transpose((2, 0, 1)) / 255
|
||||
resized_image -= 0.5
|
||||
resized_image /= 0.5
|
||||
resize_shape = resized_image.shape
|
||||
padding_im = -1.0 * np.ones((imgC, imgH, imgW_max), dtype=np.float32)
|
||||
padding_im[:, :, 0:resize_w] = resized_image
|
||||
pad_shape = padding_im.shape
|
||||
|
||||
return padding_im, resize_shape, pad_shape, valid_ratio
|
||||
|
||||
|
||||
def resize_norm_img(img, image_shape, padding=True, interpolation=cv2.INTER_LINEAR):
|
||||
imgC, imgH, imgW = image_shape
|
||||
h = img.shape[0]
|
||||
w = img.shape[1]
|
||||
if not padding:
|
||||
resized_image = cv2.resize(img, (imgW, imgH), interpolation=interpolation)
|
||||
resized_w = imgW
|
||||
else:
|
||||
ratio = w / float(h)
|
||||
if math.ceil(imgH * ratio) > imgW:
|
||||
resized_w = imgW
|
||||
else:
|
||||
resized_w = int(math.ceil(imgH * ratio))
|
||||
resized_image = cv2.resize(img, (resized_w, imgH))
|
||||
resized_image = resized_image.astype("float32")
|
||||
if image_shape[0] == 1:
|
||||
resized_image = resized_image / 255
|
||||
resized_image = resized_image[np.newaxis, :]
|
||||
else:
|
||||
resized_image = resized_image.transpose((2, 0, 1)) / 255
|
||||
resized_image -= 0.5
|
||||
resized_image /= 0.5
|
||||
padding_im = np.zeros((imgC, imgH, imgW), dtype=np.float32)
|
||||
padding_im[:, :, 0:resized_w] = resized_image
|
||||
valid_ratio = min(1.0, float(resized_w / imgW))
|
||||
return padding_im, valid_ratio
|
||||
|
||||
|
||||
def resize_norm_img_chinese(img, image_shape):
|
||||
imgC, imgH, imgW = image_shape
|
||||
# todo: change to 0 and modified image shape
|
||||
max_wh_ratio = imgW * 1.0 / imgH
|
||||
h, w = img.shape[0], img.shape[1]
|
||||
ratio = w * 1.0 / h
|
||||
max_wh_ratio = max(max_wh_ratio, ratio)
|
||||
imgW = int(imgH * max_wh_ratio)
|
||||
if math.ceil(imgH * ratio) > imgW:
|
||||
resized_w = imgW
|
||||
else:
|
||||
resized_w = int(math.ceil(imgH * ratio))
|
||||
resized_image = cv2.resize(img, (resized_w, imgH))
|
||||
resized_image = resized_image.astype("float32")
|
||||
if image_shape[0] == 1:
|
||||
resized_image = resized_image / 255
|
||||
resized_image = resized_image[np.newaxis, :]
|
||||
else:
|
||||
resized_image = resized_image.transpose((2, 0, 1)) / 255
|
||||
resized_image -= 0.5
|
||||
resized_image /= 0.5
|
||||
padding_im = np.zeros((imgC, imgH, imgW), dtype=np.float32)
|
||||
padding_im[:, :, 0:resized_w] = resized_image
|
||||
valid_ratio = min(1.0, float(resized_w / imgW))
|
||||
return padding_im, valid_ratio
|
||||
|
||||
|
||||
def resize_norm_img_srn(img, image_shape):
|
||||
imgC, imgH, imgW = image_shape
|
||||
|
||||
img_black = np.zeros((imgH, imgW))
|
||||
im_hei = img.shape[0]
|
||||
im_wid = img.shape[1]
|
||||
|
||||
if im_wid <= im_hei * 1:
|
||||
img_new = cv2.resize(img, (imgH * 1, imgH))
|
||||
elif im_wid <= im_hei * 2:
|
||||
img_new = cv2.resize(img, (imgH * 2, imgH))
|
||||
elif im_wid <= im_hei * 3:
|
||||
img_new = cv2.resize(img, (imgH * 3, imgH))
|
||||
else:
|
||||
img_new = cv2.resize(img, (imgW, imgH))
|
||||
|
||||
img_np = np.asarray(img_new)
|
||||
img_np = cv2.cvtColor(img_np, cv2.COLOR_BGR2GRAY)
|
||||
img_black[:, 0 : img_np.shape[1]] = img_np
|
||||
img_black = img_black[:, :, np.newaxis]
|
||||
|
||||
row, col, c = img_black.shape
|
||||
c = 1
|
||||
|
||||
return np.reshape(img_black, (c, row, col)).astype(np.float32)
|
||||
|
||||
|
||||
def resize_norm_img_abinet(img, image_shape):
|
||||
imgC, imgH, imgW = image_shape
|
||||
|
||||
resized_image = cv2.resize(img, (imgW, imgH), interpolation=cv2.INTER_LINEAR)
|
||||
resized_w = imgW
|
||||
resized_image = resized_image.astype("float32")
|
||||
resized_image = resized_image / 255.0
|
||||
|
||||
mean = np.array([0.485, 0.456, 0.406])
|
||||
std = np.array([0.229, 0.224, 0.225])
|
||||
resized_image = (resized_image - mean[None, None, ...]) / std[None, None, ...]
|
||||
resized_image = resized_image.transpose((2, 0, 1))
|
||||
resized_image = resized_image.astype("float32")
|
||||
|
||||
valid_ratio = min(1.0, float(resized_w / imgW))
|
||||
return resized_image, valid_ratio
|
||||
|
||||
|
||||
def srn_other_inputs(image_shape, num_heads, max_text_length):
|
||||
imgC, imgH, imgW = image_shape
|
||||
feature_dim = int((imgH / 8) * (imgW / 8))
|
||||
|
||||
encoder_word_pos = (
|
||||
np.array(range(0, feature_dim)).reshape((feature_dim, 1)).astype("int64")
|
||||
)
|
||||
gsrm_word_pos = (
|
||||
np.array(range(0, max_text_length))
|
||||
.reshape((max_text_length, 1))
|
||||
.astype("int64")
|
||||
)
|
||||
|
||||
gsrm_attn_bias_data = np.ones((1, max_text_length, max_text_length))
|
||||
gsrm_slf_attn_bias1 = np.triu(gsrm_attn_bias_data, 1).reshape(
|
||||
[1, max_text_length, max_text_length]
|
||||
)
|
||||
gsrm_slf_attn_bias1 = np.tile(gsrm_slf_attn_bias1, [num_heads, 1, 1]) * [-1e9]
|
||||
|
||||
gsrm_slf_attn_bias2 = np.tril(gsrm_attn_bias_data, -1).reshape(
|
||||
[1, max_text_length, max_text_length]
|
||||
)
|
||||
gsrm_slf_attn_bias2 = np.tile(gsrm_slf_attn_bias2, [num_heads, 1, 1]) * [-1e9]
|
||||
|
||||
return [encoder_word_pos, gsrm_word_pos, gsrm_slf_attn_bias1, gsrm_slf_attn_bias2]
|
||||
|
||||
|
||||
def flag():
|
||||
"""
|
||||
flag
|
||||
"""
|
||||
return 1 if random.random() > 0.5000001 else -1
|
||||
|
||||
|
||||
def hsv_aug(img):
|
||||
"""
|
||||
cvtColor
|
||||
"""
|
||||
hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
|
||||
delta = 0.001 * random.random() * flag()
|
||||
hsv[:, :, 2] = hsv[:, :, 2] * (1 + delta)
|
||||
new_img = cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR)
|
||||
return new_img
|
||||
|
||||
|
||||
def blur(img):
|
||||
"""
|
||||
blur
|
||||
"""
|
||||
h, w, _ = img.shape
|
||||
if h > 10 and w > 10:
|
||||
return cv2.GaussianBlur(img, (5, 5), 1)
|
||||
else:
|
||||
return img
|
||||
|
||||
|
||||
def jitter(img):
|
||||
"""
|
||||
jitter
|
||||
"""
|
||||
w, h, _ = img.shape
|
||||
if h > 10 and w > 10:
|
||||
thres = min(w, h)
|
||||
s = int(random.random() * thres * 0.01)
|
||||
src_img = img.copy()
|
||||
for i in range(s):
|
||||
img[i:, i:, :] = src_img[: w - i, : h - i, :]
|
||||
return img
|
||||
else:
|
||||
return img
|
||||
|
||||
|
||||
def add_gasuss_noise(image, mean=0, var=0.1):
|
||||
"""
|
||||
Gasuss noise
|
||||
"""
|
||||
|
||||
noise = np.random.normal(mean, var**0.5, image.shape)
|
||||
out = image + 0.5 * noise
|
||||
out = np.clip(out, 0, 255)
|
||||
out = np.uint8(out)
|
||||
return out
|
||||
|
||||
|
||||
def get_crop(image):
|
||||
"""
|
||||
random crop
|
||||
"""
|
||||
h, w, _ = image.shape
|
||||
top_min = 1
|
||||
top_max = 8
|
||||
top_crop = int(random.randint(top_min, top_max))
|
||||
top_crop = min(top_crop, h - 1)
|
||||
crop_img = image.copy()
|
||||
ratio = random.randint(0, 1)
|
||||
if ratio:
|
||||
crop_img = crop_img[top_crop:h, :, :]
|
||||
else:
|
||||
crop_img = crop_img[0 : h - top_crop, :, :]
|
||||
return crop_img
|
||||
|
||||
|
||||
def rad(x):
|
||||
"""
|
||||
rad
|
||||
"""
|
||||
return x * np.pi / 180
|
||||
|
||||
|
||||
def get_warpR(config):
|
||||
"""
|
||||
get_warpR
|
||||
"""
|
||||
anglex, angley, anglez, fov, w, h, r = (
|
||||
config.anglex,
|
||||
config.angley,
|
||||
config.anglez,
|
||||
config.fov,
|
||||
config.w,
|
||||
config.h,
|
||||
config.r,
|
||||
)
|
||||
if w > 69 and w < 112:
|
||||
anglex = anglex * 1.5
|
||||
|
||||
z = np.sqrt(w**2 + h**2) / 2 / np.tan(rad(fov / 2))
|
||||
# Homogeneous coordinate transformation matrix
|
||||
rx = np.array(
|
||||
[
|
||||
[1, 0, 0, 0],
|
||||
[0, np.cos(rad(anglex)), -np.sin(rad(anglex)), 0],
|
||||
[
|
||||
0,
|
||||
-np.sin(rad(anglex)),
|
||||
np.cos(rad(anglex)),
|
||||
0,
|
||||
],
|
||||
[0, 0, 0, 1],
|
||||
],
|
||||
np.float32,
|
||||
)
|
||||
ry = np.array(
|
||||
[
|
||||
[np.cos(rad(angley)), 0, np.sin(rad(angley)), 0],
|
||||
[0, 1, 0, 0],
|
||||
[
|
||||
-np.sin(rad(angley)),
|
||||
0,
|
||||
np.cos(rad(angley)),
|
||||
0,
|
||||
],
|
||||
[0, 0, 0, 1],
|
||||
],
|
||||
np.float32,
|
||||
)
|
||||
rz = np.array(
|
||||
[
|
||||
[np.cos(rad(anglez)), np.sin(rad(anglez)), 0, 0],
|
||||
[-np.sin(rad(anglez)), np.cos(rad(anglez)), 0, 0],
|
||||
[0, 0, 1, 0],
|
||||
[0, 0, 0, 1],
|
||||
],
|
||||
np.float32,
|
||||
)
|
||||
r = rx.dot(ry).dot(rz)
|
||||
# generate 4 points
|
||||
pcenter = np.array([h / 2, w / 2, 0, 0], np.float32)
|
||||
p1 = np.array([0, 0, 0, 0], np.float32) - pcenter
|
||||
p2 = np.array([w, 0, 0, 0], np.float32) - pcenter
|
||||
p3 = np.array([0, h, 0, 0], np.float32) - pcenter
|
||||
p4 = np.array([w, h, 0, 0], np.float32) - pcenter
|
||||
dst1 = r.dot(p1)
|
||||
dst2 = r.dot(p2)
|
||||
dst3 = r.dot(p3)
|
||||
dst4 = r.dot(p4)
|
||||
list_dst = np.array([dst1, dst2, dst3, dst4])
|
||||
org = np.array([[0, 0], [w, 0], [0, h], [w, h]], np.float32)
|
||||
dst = np.zeros((4, 2), np.float32)
|
||||
# Project onto the image plane
|
||||
dst[:, 0] = list_dst[:, 0] * z / (z - list_dst[:, 2]) + pcenter[0]
|
||||
dst[:, 1] = list_dst[:, 1] * z / (z - list_dst[:, 2]) + pcenter[1]
|
||||
|
||||
warpR = cv2.getPerspectiveTransform(org, dst)
|
||||
|
||||
dst1, dst2, dst3, dst4 = dst
|
||||
r1 = int(min(dst1[1], dst2[1]))
|
||||
r2 = int(max(dst3[1], dst4[1]))
|
||||
c1 = int(min(dst1[0], dst3[0]))
|
||||
c2 = int(max(dst2[0], dst4[0]))
|
||||
|
||||
try:
|
||||
ratio = min(1.0 * h / (r2 - r1), 1.0 * w / (c2 - c1))
|
||||
|
||||
dx = -c1
|
||||
dy = -r1
|
||||
T1 = np.float32([[1.0, 0, dx], [0, 1.0, dy], [0, 0, 1.0 / ratio]])
|
||||
ret = T1.dot(warpR)
|
||||
except:
|
||||
ratio = 1.0
|
||||
T1 = np.float32([[1.0, 0, 0], [0, 1.0, 0], [0, 0, 1.0]])
|
||||
ret = T1
|
||||
return ret, (-r1, -c1), ratio, dst
|
||||
|
||||
|
||||
def get_warpAffine(config):
|
||||
"""
|
||||
get_warpAffine
|
||||
"""
|
||||
anglez = config.anglez
|
||||
rz = np.array(
|
||||
[
|
||||
[np.cos(rad(anglez)), np.sin(rad(anglez)), 0],
|
||||
[-np.sin(rad(anglez)), np.cos(rad(anglez)), 0],
|
||||
],
|
||||
np.float32,
|
||||
)
|
||||
return rz
|
||||
810
ppocr/data/imaug/sast_process.py
Normal file
810
ppocr/data/imaug/sast_process.py
Normal file
@@ -0,0 +1,810 @@
|
||||
# copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
"""
|
||||
This part code is referred from:
|
||||
https://github.com/songdejia/EAST/blob/master/data_utils.py
|
||||
"""
|
||||
import math
|
||||
import cv2
|
||||
import numpy as np
|
||||
import json
|
||||
import sys
|
||||
import os
|
||||
|
||||
__all__ = ["SASTProcessTrain"]
|
||||
|
||||
|
||||
class SASTProcessTrain(object):
|
||||
def __init__(
|
||||
self,
|
||||
image_shape=[512, 512],
|
||||
min_crop_size=24,
|
||||
min_crop_side_ratio=0.3,
|
||||
min_text_size=10,
|
||||
max_text_size=512,
|
||||
**kwargs,
|
||||
):
|
||||
self.input_size = image_shape[1]
|
||||
self.min_crop_size = min_crop_size
|
||||
self.min_crop_side_ratio = min_crop_side_ratio
|
||||
self.min_text_size = min_text_size
|
||||
self.max_text_size = max_text_size
|
||||
|
||||
def quad_area(self, poly):
|
||||
"""
|
||||
compute area of a polygon
|
||||
:param poly:
|
||||
:return:
|
||||
"""
|
||||
edge = [
|
||||
(poly[1][0] - poly[0][0]) * (poly[1][1] + poly[0][1]),
|
||||
(poly[2][0] - poly[1][0]) * (poly[2][1] + poly[1][1]),
|
||||
(poly[3][0] - poly[2][0]) * (poly[3][1] + poly[2][1]),
|
||||
(poly[0][0] - poly[3][0]) * (poly[0][1] + poly[3][1]),
|
||||
]
|
||||
return np.sum(edge) / 2.0
|
||||
|
||||
def gen_quad_from_poly(self, poly):
|
||||
"""
|
||||
Generate min area quad from poly.
|
||||
"""
|
||||
point_num = poly.shape[0]
|
||||
min_area_quad = np.zeros((4, 2), dtype=np.float32)
|
||||
if True:
|
||||
rect = cv2.minAreaRect(
|
||||
poly.astype(np.int32)
|
||||
) # (center (x,y), (width, height), angle of rotation)
|
||||
center_point = rect[0]
|
||||
box = np.array(cv2.boxPoints(rect))
|
||||
|
||||
first_point_idx = 0
|
||||
min_dist = 1e4
|
||||
for i in range(4):
|
||||
dist = (
|
||||
np.linalg.norm(box[(i + 0) % 4] - poly[0])
|
||||
+ np.linalg.norm(box[(i + 1) % 4] - poly[point_num // 2 - 1])
|
||||
+ np.linalg.norm(box[(i + 2) % 4] - poly[point_num // 2])
|
||||
+ np.linalg.norm(box[(i + 3) % 4] - poly[-1])
|
||||
)
|
||||
if dist < min_dist:
|
||||
min_dist = dist
|
||||
first_point_idx = i
|
||||
for i in range(4):
|
||||
min_area_quad[i] = box[(first_point_idx + i) % 4]
|
||||
|
||||
return min_area_quad
|
||||
|
||||
def check_and_validate_polys(self, polys, tags, xxx_todo_changeme):
|
||||
"""
|
||||
check so that the text poly is in the same direction,
|
||||
and also filter some invalid polygons
|
||||
:param polys:
|
||||
:param tags:
|
||||
:return:
|
||||
"""
|
||||
(h, w) = xxx_todo_changeme
|
||||
if polys.shape[0] == 0:
|
||||
return polys, np.array([]), np.array([])
|
||||
polys[:, :, 0] = np.clip(polys[:, :, 0], 0, w - 1)
|
||||
polys[:, :, 1] = np.clip(polys[:, :, 1], 0, h - 1)
|
||||
|
||||
validated_polys = []
|
||||
validated_tags = []
|
||||
hv_tags = []
|
||||
for poly, tag in zip(polys, tags):
|
||||
quad = self.gen_quad_from_poly(poly)
|
||||
p_area = self.quad_area(quad)
|
||||
if abs(p_area) < 1:
|
||||
print("invalid poly")
|
||||
continue
|
||||
if p_area > 0:
|
||||
if tag == False:
|
||||
print("poly in wrong direction")
|
||||
tag = True # reversed cases should be ignore
|
||||
poly = poly[(0, 15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1), :]
|
||||
quad = quad[(0, 3, 2, 1), :]
|
||||
|
||||
len_w = np.linalg.norm(quad[0] - quad[1]) + np.linalg.norm(
|
||||
quad[3] - quad[2]
|
||||
)
|
||||
len_h = np.linalg.norm(quad[0] - quad[3]) + np.linalg.norm(
|
||||
quad[1] - quad[2]
|
||||
)
|
||||
hv_tag = 1
|
||||
|
||||
if len_w * 2.0 < len_h:
|
||||
hv_tag = 0
|
||||
|
||||
validated_polys.append(poly)
|
||||
validated_tags.append(tag)
|
||||
hv_tags.append(hv_tag)
|
||||
return np.array(validated_polys), np.array(validated_tags), np.array(hv_tags)
|
||||
|
||||
def crop_area(self, im, polys, tags, hv_tags, crop_background=False, max_tries=25):
|
||||
"""
|
||||
make random crop from the input image
|
||||
:param im:
|
||||
:param polys:
|
||||
:param tags:
|
||||
:param crop_background:
|
||||
:param max_tries: 50 -> 25
|
||||
:return:
|
||||
"""
|
||||
h, w, _ = im.shape
|
||||
pad_h = h // 10
|
||||
pad_w = w // 10
|
||||
h_array = np.zeros((h + pad_h * 2), dtype=np.int32)
|
||||
w_array = np.zeros((w + pad_w * 2), dtype=np.int32)
|
||||
for poly in polys:
|
||||
poly = np.round(poly, decimals=0).astype(np.int32)
|
||||
minx = np.min(poly[:, 0])
|
||||
maxx = np.max(poly[:, 0])
|
||||
w_array[minx + pad_w : maxx + pad_w] = 1
|
||||
miny = np.min(poly[:, 1])
|
||||
maxy = np.max(poly[:, 1])
|
||||
h_array[miny + pad_h : maxy + pad_h] = 1
|
||||
# ensure the cropped area not across a text
|
||||
h_axis = np.where(h_array == 0)[0]
|
||||
w_axis = np.where(w_array == 0)[0]
|
||||
if len(h_axis) == 0 or len(w_axis) == 0:
|
||||
return im, polys, tags, hv_tags
|
||||
for i in range(max_tries):
|
||||
xx = np.random.choice(w_axis, size=2)
|
||||
xmin = np.min(xx) - pad_w
|
||||
xmax = np.max(xx) - pad_w
|
||||
xmin = np.clip(xmin, 0, w - 1)
|
||||
xmax = np.clip(xmax, 0, w - 1)
|
||||
yy = np.random.choice(h_axis, size=2)
|
||||
ymin = np.min(yy) - pad_h
|
||||
ymax = np.max(yy) - pad_h
|
||||
ymin = np.clip(ymin, 0, h - 1)
|
||||
ymax = np.clip(ymax, 0, h - 1)
|
||||
# if xmax - xmin < ARGS.min_crop_side_ratio * w or \
|
||||
# ymax - ymin < ARGS.min_crop_side_ratio * h:
|
||||
if xmax - xmin < self.min_crop_size or ymax - ymin < self.min_crop_size:
|
||||
# area too small
|
||||
continue
|
||||
if polys.shape[0] != 0:
|
||||
poly_axis_in_area = (
|
||||
(polys[:, :, 0] >= xmin)
|
||||
& (polys[:, :, 0] <= xmax)
|
||||
& (polys[:, :, 1] >= ymin)
|
||||
& (polys[:, :, 1] <= ymax)
|
||||
)
|
||||
selected_polys = np.where(np.sum(poly_axis_in_area, axis=1) == 4)[0]
|
||||
else:
|
||||
selected_polys = []
|
||||
if len(selected_polys) == 0:
|
||||
# no text in this area
|
||||
if crop_background:
|
||||
return (
|
||||
im[ymin : ymax + 1, xmin : xmax + 1, :],
|
||||
polys[selected_polys],
|
||||
tags[selected_polys],
|
||||
hv_tags[selected_polys],
|
||||
)
|
||||
else:
|
||||
continue
|
||||
im = im[ymin : ymax + 1, xmin : xmax + 1, :]
|
||||
polys = polys[selected_polys]
|
||||
tags = tags[selected_polys]
|
||||
hv_tags = hv_tags[selected_polys]
|
||||
polys[:, :, 0] -= xmin
|
||||
polys[:, :, 1] -= ymin
|
||||
return im, polys, tags, hv_tags
|
||||
|
||||
return im, polys, tags, hv_tags
|
||||
|
||||
def generate_direction_map(self, poly_quads, direction_map):
|
||||
""" """
|
||||
width_list = []
|
||||
height_list = []
|
||||
for quad in poly_quads:
|
||||
quad_w = (
|
||||
np.linalg.norm(quad[0] - quad[1]) + np.linalg.norm(quad[2] - quad[3])
|
||||
) / 2.0
|
||||
quad_h = (
|
||||
np.linalg.norm(quad[0] - quad[3]) + np.linalg.norm(quad[2] - quad[1])
|
||||
) / 2.0
|
||||
width_list.append(quad_w)
|
||||
height_list.append(quad_h)
|
||||
norm_width = max(sum(width_list) / (len(width_list) + 1e-6), 1.0)
|
||||
average_height = max(sum(height_list) / (len(height_list) + 1e-6), 1.0)
|
||||
|
||||
for quad in poly_quads:
|
||||
direct_vector_full = ((quad[1] + quad[2]) - (quad[0] + quad[3])) / 2.0
|
||||
direct_vector = (
|
||||
direct_vector_full
|
||||
/ (np.linalg.norm(direct_vector_full) + 1e-6)
|
||||
* norm_width
|
||||
)
|
||||
direction_label = tuple(
|
||||
map(
|
||||
float,
|
||||
[direct_vector[0], direct_vector[1], 1.0 / (average_height + 1e-6)],
|
||||
)
|
||||
)
|
||||
cv2.fillPoly(
|
||||
direction_map,
|
||||
quad.round().astype(np.int32)[np.newaxis, :, :],
|
||||
direction_label,
|
||||
)
|
||||
return direction_map
|
||||
|
||||
def calculate_average_height(self, poly_quads):
|
||||
""" """
|
||||
height_list = []
|
||||
for quad in poly_quads:
|
||||
quad_h = (
|
||||
np.linalg.norm(quad[0] - quad[3]) + np.linalg.norm(quad[2] - quad[1])
|
||||
) / 2.0
|
||||
height_list.append(quad_h)
|
||||
average_height = max(sum(height_list) / len(height_list), 1.0)
|
||||
return average_height
|
||||
|
||||
def generate_tcl_label(
|
||||
self, hw, polys, tags, ds_ratio, tcl_ratio=0.3, shrink_ratio_of_width=0.15
|
||||
):
|
||||
"""
|
||||
Generate polygon.
|
||||
"""
|
||||
h, w = hw
|
||||
h, w = int(h * ds_ratio), int(w * ds_ratio)
|
||||
polys = polys * ds_ratio
|
||||
|
||||
score_map = np.zeros(
|
||||
(
|
||||
h,
|
||||
w,
|
||||
),
|
||||
dtype=np.float32,
|
||||
)
|
||||
tbo_map = np.zeros((h, w, 5), dtype=np.float32)
|
||||
training_mask = np.ones(
|
||||
(
|
||||
h,
|
||||
w,
|
||||
),
|
||||
dtype=np.float32,
|
||||
)
|
||||
direction_map = np.ones((h, w, 3)) * np.array([0, 0, 1]).reshape(
|
||||
[1, 1, 3]
|
||||
).astype(np.float32)
|
||||
|
||||
for poly_idx, poly_tag in enumerate(zip(polys, tags)):
|
||||
poly = poly_tag[0]
|
||||
tag = poly_tag[1]
|
||||
|
||||
# generate min_area_quad
|
||||
min_area_quad, center_point = self.gen_min_area_quad_from_poly(poly)
|
||||
min_area_quad_h = 0.5 * (
|
||||
np.linalg.norm(min_area_quad[0] - min_area_quad[3])
|
||||
+ np.linalg.norm(min_area_quad[1] - min_area_quad[2])
|
||||
)
|
||||
min_area_quad_w = 0.5 * (
|
||||
np.linalg.norm(min_area_quad[0] - min_area_quad[1])
|
||||
+ np.linalg.norm(min_area_quad[2] - min_area_quad[3])
|
||||
)
|
||||
|
||||
if (
|
||||
min(min_area_quad_h, min_area_quad_w) < self.min_text_size * ds_ratio
|
||||
or min(min_area_quad_h, min_area_quad_w) > self.max_text_size * ds_ratio
|
||||
):
|
||||
continue
|
||||
|
||||
if tag:
|
||||
# continue
|
||||
cv2.fillPoly(
|
||||
training_mask, poly.astype(np.int32)[np.newaxis, :, :], 0.15
|
||||
)
|
||||
else:
|
||||
tcl_poly = self.poly2tcl(poly, tcl_ratio)
|
||||
tcl_quads = self.poly2quads(tcl_poly)
|
||||
poly_quads = self.poly2quads(poly)
|
||||
# stcl map
|
||||
stcl_quads, quad_index = self.shrink_poly_along_width(
|
||||
tcl_quads,
|
||||
shrink_ratio_of_width=shrink_ratio_of_width,
|
||||
expand_height_ratio=1.0 / tcl_ratio,
|
||||
)
|
||||
# generate tcl map
|
||||
cv2.fillPoly(score_map, np.round(stcl_quads).astype(np.int32), 1.0)
|
||||
|
||||
# generate tbo map
|
||||
for idx, quad in enumerate(stcl_quads):
|
||||
quad_mask = np.zeros((h, w), dtype=np.float32)
|
||||
quad_mask = cv2.fillPoly(
|
||||
quad_mask,
|
||||
np.round(quad[np.newaxis, :, :]).astype(np.int32),
|
||||
1.0,
|
||||
)
|
||||
tbo_map = self.gen_quad_tbo(
|
||||
poly_quads[quad_index[idx]], quad_mask, tbo_map
|
||||
)
|
||||
return score_map, tbo_map, training_mask
|
||||
|
||||
def generate_tvo_and_tco(self, hw, polys, tags, tcl_ratio=0.3, ds_ratio=0.25):
|
||||
"""
|
||||
Generate tcl map, tvo map and tbo map.
|
||||
"""
|
||||
h, w = hw
|
||||
h, w = int(h * ds_ratio), int(w * ds_ratio)
|
||||
polys = polys * ds_ratio
|
||||
poly_mask = np.zeros((h, w), dtype=np.float32)
|
||||
|
||||
tvo_map = np.ones((9, h, w), dtype=np.float32)
|
||||
tvo_map[0:-1:2] = np.tile(np.arange(0, w), (h, 1))
|
||||
tvo_map[1:-1:2] = np.tile(np.arange(0, w), (h, 1)).T
|
||||
poly_tv_xy_map = np.zeros((8, h, w), dtype=np.float32)
|
||||
|
||||
# tco map
|
||||
tco_map = np.ones((3, h, w), dtype=np.float32)
|
||||
tco_map[0] = np.tile(np.arange(0, w), (h, 1))
|
||||
tco_map[1] = np.tile(np.arange(0, w), (h, 1)).T
|
||||
poly_tc_xy_map = np.zeros((2, h, w), dtype=np.float32)
|
||||
|
||||
poly_short_edge_map = np.ones((h, w), dtype=np.float32)
|
||||
|
||||
for poly, poly_tag in zip(polys, tags):
|
||||
if poly_tag == True:
|
||||
continue
|
||||
|
||||
# adjust point order for vertical poly
|
||||
poly = self.adjust_point(poly)
|
||||
|
||||
# generate min_area_quad
|
||||
min_area_quad, center_point = self.gen_min_area_quad_from_poly(poly)
|
||||
min_area_quad_h = 0.5 * (
|
||||
np.linalg.norm(min_area_quad[0] - min_area_quad[3])
|
||||
+ np.linalg.norm(min_area_quad[1] - min_area_quad[2])
|
||||
)
|
||||
min_area_quad_w = 0.5 * (
|
||||
np.linalg.norm(min_area_quad[0] - min_area_quad[1])
|
||||
+ np.linalg.norm(min_area_quad[2] - min_area_quad[3])
|
||||
)
|
||||
|
||||
# generate tcl map and text, 128 * 128
|
||||
tcl_poly = self.poly2tcl(poly, tcl_ratio)
|
||||
|
||||
# generate poly_tv_xy_map
|
||||
for idx in range(4):
|
||||
cv2.fillPoly(
|
||||
poly_tv_xy_map[2 * idx],
|
||||
np.round(tcl_poly[np.newaxis, :, :]).astype(np.int32),
|
||||
float(min(max(min_area_quad[idx, 0], 0), w)),
|
||||
)
|
||||
cv2.fillPoly(
|
||||
poly_tv_xy_map[2 * idx + 1],
|
||||
np.round(tcl_poly[np.newaxis, :, :]).astype(np.int32),
|
||||
float(min(max(min_area_quad[idx, 1], 0), h)),
|
||||
)
|
||||
|
||||
# generate poly_tc_xy_map
|
||||
for idx in range(2):
|
||||
cv2.fillPoly(
|
||||
poly_tc_xy_map[idx],
|
||||
np.round(tcl_poly[np.newaxis, :, :]).astype(np.int32),
|
||||
float(center_point[idx]),
|
||||
)
|
||||
|
||||
# generate poly_short_edge_map
|
||||
cv2.fillPoly(
|
||||
poly_short_edge_map,
|
||||
np.round(tcl_poly[np.newaxis, :, :]).astype(np.int32),
|
||||
float(max(min(min_area_quad_h, min_area_quad_w), 1.0)),
|
||||
)
|
||||
|
||||
# generate poly_mask and training_mask
|
||||
cv2.fillPoly(
|
||||
poly_mask, np.round(tcl_poly[np.newaxis, :, :]).astype(np.int32), 1
|
||||
)
|
||||
|
||||
tvo_map *= poly_mask
|
||||
tvo_map[:8] -= poly_tv_xy_map
|
||||
tvo_map[-1] /= poly_short_edge_map
|
||||
tvo_map = tvo_map.transpose((1, 2, 0))
|
||||
|
||||
tco_map *= poly_mask
|
||||
tco_map[:2] -= poly_tc_xy_map
|
||||
tco_map[-1] /= poly_short_edge_map
|
||||
tco_map = tco_map.transpose((1, 2, 0))
|
||||
|
||||
return tvo_map, tco_map
|
||||
|
||||
def adjust_point(self, poly):
|
||||
"""
|
||||
adjust point order.
|
||||
"""
|
||||
point_num = poly.shape[0]
|
||||
if point_num == 4:
|
||||
len_1 = np.linalg.norm(poly[0] - poly[1])
|
||||
len_2 = np.linalg.norm(poly[1] - poly[2])
|
||||
len_3 = np.linalg.norm(poly[2] - poly[3])
|
||||
len_4 = np.linalg.norm(poly[3] - poly[0])
|
||||
|
||||
if (len_1 + len_3) * 1.5 < (len_2 + len_4):
|
||||
poly = poly[[1, 2, 3, 0], :]
|
||||
|
||||
elif point_num > 4:
|
||||
vector_1 = poly[0] - poly[1]
|
||||
vector_2 = poly[1] - poly[2]
|
||||
cos_theta = np.dot(vector_1, vector_2) / (
|
||||
np.linalg.norm(vector_1) * np.linalg.norm(vector_2) + 1e-6
|
||||
)
|
||||
theta = np.arccos(np.round(cos_theta, decimals=4))
|
||||
|
||||
if abs(theta) > (70 / 180 * math.pi):
|
||||
index = list(range(1, point_num)) + [0]
|
||||
poly = poly[np.array(index), :]
|
||||
return poly
|
||||
|
||||
def gen_min_area_quad_from_poly(self, poly):
|
||||
"""
|
||||
Generate min area quad from poly.
|
||||
"""
|
||||
point_num = poly.shape[0]
|
||||
min_area_quad = np.zeros((4, 2), dtype=np.float32)
|
||||
if point_num == 4:
|
||||
min_area_quad = poly
|
||||
center_point = np.sum(poly, axis=0) / 4
|
||||
else:
|
||||
rect = cv2.minAreaRect(
|
||||
poly.astype(np.int32)
|
||||
) # (center (x,y), (width, height), angle of rotation)
|
||||
center_point = rect[0]
|
||||
box = np.array(cv2.boxPoints(rect))
|
||||
|
||||
first_point_idx = 0
|
||||
min_dist = 1e4
|
||||
for i in range(4):
|
||||
dist = (
|
||||
np.linalg.norm(box[(i + 0) % 4] - poly[0])
|
||||
+ np.linalg.norm(box[(i + 1) % 4] - poly[point_num // 2 - 1])
|
||||
+ np.linalg.norm(box[(i + 2) % 4] - poly[point_num // 2])
|
||||
+ np.linalg.norm(box[(i + 3) % 4] - poly[-1])
|
||||
)
|
||||
if dist < min_dist:
|
||||
min_dist = dist
|
||||
first_point_idx = i
|
||||
|
||||
for i in range(4):
|
||||
min_area_quad[i] = box[(first_point_idx + i) % 4]
|
||||
|
||||
return min_area_quad, center_point
|
||||
|
||||
def shrink_quad_along_width(self, quad, begin_width_ratio=0.0, end_width_ratio=1.0):
|
||||
"""
|
||||
Generate shrink_quad_along_width.
|
||||
"""
|
||||
ratio_pair = np.array(
|
||||
[[begin_width_ratio], [end_width_ratio]], dtype=np.float32
|
||||
)
|
||||
p0_1 = quad[0] + (quad[1] - quad[0]) * ratio_pair
|
||||
p3_2 = quad[3] + (quad[2] - quad[3]) * ratio_pair
|
||||
return np.array([p0_1[0], p0_1[1], p3_2[1], p3_2[0]])
|
||||
|
||||
def shrink_poly_along_width(
|
||||
self, quads, shrink_ratio_of_width, expand_height_ratio=1.0
|
||||
):
|
||||
"""
|
||||
shrink poly with given length.
|
||||
"""
|
||||
upper_edge_list = []
|
||||
|
||||
def get_cut_info(edge_len_list, cut_len):
|
||||
for idx, edge_len in enumerate(edge_len_list):
|
||||
cut_len -= edge_len
|
||||
if cut_len <= 0.000001:
|
||||
ratio = (cut_len + edge_len_list[idx]) / edge_len_list[idx]
|
||||
return idx, ratio
|
||||
|
||||
for quad in quads:
|
||||
upper_edge_len = np.linalg.norm(quad[0] - quad[1])
|
||||
upper_edge_list.append(upper_edge_len)
|
||||
|
||||
# length of left edge and right edge.
|
||||
left_length = np.linalg.norm(quads[0][0] - quads[0][3]) * expand_height_ratio
|
||||
right_length = np.linalg.norm(quads[-1][1] - quads[-1][2]) * expand_height_ratio
|
||||
|
||||
shrink_length = (
|
||||
min(left_length, right_length, sum(upper_edge_list)) * shrink_ratio_of_width
|
||||
)
|
||||
# shrinking length
|
||||
upper_len_left = shrink_length
|
||||
upper_len_right = sum(upper_edge_list) - shrink_length
|
||||
|
||||
left_idx, left_ratio = get_cut_info(upper_edge_list, upper_len_left)
|
||||
left_quad = self.shrink_quad_along_width(
|
||||
quads[left_idx], begin_width_ratio=left_ratio, end_width_ratio=1
|
||||
)
|
||||
right_idx, right_ratio = get_cut_info(upper_edge_list, upper_len_right)
|
||||
right_quad = self.shrink_quad_along_width(
|
||||
quads[right_idx], begin_width_ratio=0, end_width_ratio=right_ratio
|
||||
)
|
||||
|
||||
out_quad_list = []
|
||||
if left_idx == right_idx:
|
||||
out_quad_list.append(
|
||||
[left_quad[0], right_quad[1], right_quad[2], left_quad[3]]
|
||||
)
|
||||
else:
|
||||
out_quad_list.append(left_quad)
|
||||
for idx in range(left_idx + 1, right_idx):
|
||||
out_quad_list.append(quads[idx])
|
||||
out_quad_list.append(right_quad)
|
||||
|
||||
return np.array(out_quad_list), list(range(left_idx, right_idx + 1))
|
||||
|
||||
def vector_angle(self, A, B):
|
||||
"""
|
||||
Calculate the angle between vector AB and x-axis positive direction.
|
||||
"""
|
||||
AB = np.array([B[1] - A[1], B[0] - A[0]])
|
||||
return np.arctan2(*AB)
|
||||
|
||||
def theta_line_cross_point(self, theta, point):
|
||||
"""
|
||||
Calculate the line through given point and angle in ax + by + c =0 form.
|
||||
"""
|
||||
x, y = point
|
||||
cos = np.cos(theta)
|
||||
sin = np.sin(theta)
|
||||
return [sin, -cos, cos * y - sin * x]
|
||||
|
||||
def line_cross_two_point(self, A, B):
|
||||
"""
|
||||
Calculate the line through given point A and B in ax + by + c =0 form.
|
||||
"""
|
||||
angle = self.vector_angle(A, B)
|
||||
return self.theta_line_cross_point(angle, A)
|
||||
|
||||
def average_angle(self, poly):
|
||||
"""
|
||||
Calculate the average angle between left and right edge in given poly.
|
||||
"""
|
||||
p0, p1, p2, p3 = poly
|
||||
angle30 = self.vector_angle(p3, p0)
|
||||
angle21 = self.vector_angle(p2, p1)
|
||||
return (angle30 + angle21) / 2
|
||||
|
||||
def line_cross_point(self, line1, line2):
|
||||
"""
|
||||
line1 and line2 in 0=ax+by+c form, compute the cross point of line1 and line2
|
||||
"""
|
||||
a1, b1, c1 = line1
|
||||
a2, b2, c2 = line2
|
||||
d = a1 * b2 - a2 * b1
|
||||
|
||||
if d == 0:
|
||||
# print("line1", line1)
|
||||
# print("line2", line2)
|
||||
print("Cross point does not exist")
|
||||
return np.array([0, 0], dtype=np.float32)
|
||||
else:
|
||||
x = (b1 * c2 - b2 * c1) / d
|
||||
y = (a2 * c1 - a1 * c2) / d
|
||||
|
||||
return np.array([x, y], dtype=np.float32)
|
||||
|
||||
def quad2tcl(self, poly, ratio):
|
||||
"""
|
||||
Generate center line by poly clock-wise point. (4, 2)
|
||||
"""
|
||||
ratio_pair = np.array([[0.5 - ratio / 2], [0.5 + ratio / 2]], dtype=np.float32)
|
||||
p0_3 = poly[0] + (poly[3] - poly[0]) * ratio_pair
|
||||
p1_2 = poly[1] + (poly[2] - poly[1]) * ratio_pair
|
||||
return np.array([p0_3[0], p1_2[0], p1_2[1], p0_3[1]])
|
||||
|
||||
def poly2tcl(self, poly, ratio):
|
||||
"""
|
||||
Generate center line by poly clock-wise point.
|
||||
"""
|
||||
ratio_pair = np.array([[0.5 - ratio / 2], [0.5 + ratio / 2]], dtype=np.float32)
|
||||
tcl_poly = np.zeros_like(poly)
|
||||
point_num = poly.shape[0]
|
||||
|
||||
for idx in range(point_num // 2):
|
||||
point_pair = (
|
||||
poly[idx] + (poly[point_num - 1 - idx] - poly[idx]) * ratio_pair
|
||||
)
|
||||
tcl_poly[idx] = point_pair[0]
|
||||
tcl_poly[point_num - 1 - idx] = point_pair[1]
|
||||
return tcl_poly
|
||||
|
||||
def gen_quad_tbo(self, quad, tcl_mask, tbo_map):
|
||||
"""
|
||||
Generate tbo_map for give quad.
|
||||
"""
|
||||
# upper and lower line function: ax + by + c = 0;
|
||||
up_line = self.line_cross_two_point(quad[0], quad[1])
|
||||
lower_line = self.line_cross_two_point(quad[3], quad[2])
|
||||
|
||||
quad_h = 0.5 * (
|
||||
np.linalg.norm(quad[0] - quad[3]) + np.linalg.norm(quad[1] - quad[2])
|
||||
)
|
||||
quad_w = 0.5 * (
|
||||
np.linalg.norm(quad[0] - quad[1]) + np.linalg.norm(quad[2] - quad[3])
|
||||
)
|
||||
|
||||
# average angle of left and right line.
|
||||
angle = self.average_angle(quad)
|
||||
|
||||
xy_in_poly = np.argwhere(tcl_mask == 1)
|
||||
for y, x in xy_in_poly:
|
||||
point = (x, y)
|
||||
line = self.theta_line_cross_point(angle, point)
|
||||
cross_point_upper = self.line_cross_point(up_line, line)
|
||||
cross_point_lower = self.line_cross_point(lower_line, line)
|
||||
##FIX, offset reverse
|
||||
upper_offset_x, upper_offset_y = cross_point_upper - point
|
||||
lower_offset_x, lower_offset_y = cross_point_lower - point
|
||||
tbo_map[y, x, 0] = upper_offset_y
|
||||
tbo_map[y, x, 1] = upper_offset_x
|
||||
tbo_map[y, x, 2] = lower_offset_y
|
||||
tbo_map[y, x, 3] = lower_offset_x
|
||||
tbo_map[y, x, 4] = 1.0 / max(min(quad_h, quad_w), 1.0) * 2
|
||||
return tbo_map
|
||||
|
||||
def poly2quads(self, poly):
|
||||
"""
|
||||
Split poly into quads.
|
||||
"""
|
||||
quad_list = []
|
||||
point_num = poly.shape[0]
|
||||
|
||||
# point pair
|
||||
point_pair_list = []
|
||||
for idx in range(point_num // 2):
|
||||
point_pair = [poly[idx], poly[point_num - 1 - idx]]
|
||||
point_pair_list.append(point_pair)
|
||||
|
||||
quad_num = point_num // 2 - 1
|
||||
for idx in range(quad_num):
|
||||
# reshape and adjust to clock-wise
|
||||
quad_list.append(
|
||||
(np.array(point_pair_list)[[idx, idx + 1]]).reshape(4, 2)[[0, 2, 3, 1]]
|
||||
)
|
||||
|
||||
return np.array(quad_list)
|
||||
|
||||
def __call__(self, data):
|
||||
im = data["image"]
|
||||
text_polys = data["polys"]
|
||||
text_tags = data["ignore_tags"]
|
||||
if im is None:
|
||||
return None
|
||||
if text_polys.shape[0] == 0:
|
||||
return None
|
||||
|
||||
h, w, _ = im.shape
|
||||
text_polys, text_tags, hv_tags = self.check_and_validate_polys(
|
||||
text_polys, text_tags, (h, w)
|
||||
)
|
||||
|
||||
if text_polys.shape[0] == 0:
|
||||
return None
|
||||
|
||||
# set aspect ratio and keep area fix
|
||||
asp_scales = np.arange(1.0, 1.55, 0.1)
|
||||
asp_scale = np.random.choice(asp_scales)
|
||||
|
||||
if np.random.rand() < 0.5:
|
||||
asp_scale = 1.0 / asp_scale
|
||||
asp_scale = math.sqrt(asp_scale)
|
||||
|
||||
asp_wx = asp_scale
|
||||
asp_hy = 1.0 / asp_scale
|
||||
im = cv2.resize(im, dsize=None, fx=asp_wx, fy=asp_hy)
|
||||
text_polys[:, :, 0] *= asp_wx
|
||||
text_polys[:, :, 1] *= asp_hy
|
||||
|
||||
h, w, _ = im.shape
|
||||
if max(h, w) > 2048:
|
||||
rd_scale = 2048.0 / max(h, w)
|
||||
im = cv2.resize(im, dsize=None, fx=rd_scale, fy=rd_scale)
|
||||
text_polys *= rd_scale
|
||||
h, w, _ = im.shape
|
||||
if min(h, w) < 16:
|
||||
return None
|
||||
|
||||
# no background
|
||||
im, text_polys, text_tags, hv_tags = self.crop_area(
|
||||
im, text_polys, text_tags, hv_tags, crop_background=False
|
||||
)
|
||||
|
||||
if text_polys.shape[0] == 0:
|
||||
return None
|
||||
# continue for all ignore case
|
||||
if np.sum((text_tags * 1.0)) >= text_tags.size:
|
||||
return None
|
||||
new_h, new_w, _ = im.shape
|
||||
if (new_h is None) or (new_w is None):
|
||||
return None
|
||||
# resize image
|
||||
std_ratio = float(self.input_size) / max(new_w, new_h)
|
||||
rand_scales = np.array(
|
||||
[0.25, 0.375, 0.5, 0.625, 0.75, 0.875, 1.0, 1.0, 1.0, 1.0, 1.0]
|
||||
)
|
||||
rz_scale = std_ratio * np.random.choice(rand_scales)
|
||||
im = cv2.resize(im, dsize=None, fx=rz_scale, fy=rz_scale)
|
||||
text_polys[:, :, 0] *= rz_scale
|
||||
text_polys[:, :, 1] *= rz_scale
|
||||
|
||||
# add gaussian blur
|
||||
if np.random.rand() < 0.1 * 0.5:
|
||||
ks = np.random.permutation(5)[0] + 1
|
||||
ks = int(ks / 2) * 2 + 1
|
||||
im = cv2.GaussianBlur(im, ksize=(ks, ks), sigmaX=0, sigmaY=0)
|
||||
# add brighter
|
||||
if np.random.rand() < 0.1 * 0.5:
|
||||
im = im * (1.0 + np.random.rand() * 0.5)
|
||||
im = np.clip(im, 0.0, 255.0)
|
||||
# add darker
|
||||
if np.random.rand() < 0.1 * 0.5:
|
||||
im = im * (1.0 - np.random.rand() * 0.5)
|
||||
im = np.clip(im, 0.0, 255.0)
|
||||
|
||||
# Padding the im to [input_size, input_size]
|
||||
new_h, new_w, _ = im.shape
|
||||
if min(new_w, new_h) < self.input_size * 0.5:
|
||||
return None
|
||||
|
||||
im_padded = np.ones((self.input_size, self.input_size, 3), dtype=np.float32)
|
||||
im_padded[:, :, 2] = 0.485 * 255
|
||||
im_padded[:, :, 1] = 0.456 * 255
|
||||
im_padded[:, :, 0] = 0.406 * 255
|
||||
|
||||
# Random the start position
|
||||
del_h = self.input_size - new_h
|
||||
del_w = self.input_size - new_w
|
||||
sh, sw = 0, 0
|
||||
if del_h > 1:
|
||||
sh = int(np.random.rand() * del_h)
|
||||
if del_w > 1:
|
||||
sw = int(np.random.rand() * del_w)
|
||||
|
||||
# Padding
|
||||
im_padded[sh : sh + new_h, sw : sw + new_w, :] = im.copy()
|
||||
text_polys[:, :, 0] += sw
|
||||
text_polys[:, :, 1] += sh
|
||||
|
||||
score_map, border_map, training_mask = self.generate_tcl_label(
|
||||
(self.input_size, self.input_size), text_polys, text_tags, 0.25
|
||||
)
|
||||
|
||||
# SAST head
|
||||
tvo_map, tco_map = self.generate_tvo_and_tco(
|
||||
(self.input_size, self.input_size),
|
||||
text_polys,
|
||||
text_tags,
|
||||
tcl_ratio=0.3,
|
||||
ds_ratio=0.25,
|
||||
)
|
||||
# print("test--------tvo_map shape:", tvo_map.shape)
|
||||
|
||||
im_padded[:, :, 2] -= 0.485 * 255
|
||||
im_padded[:, :, 1] -= 0.456 * 255
|
||||
im_padded[:, :, 0] -= 0.406 * 255
|
||||
im_padded[:, :, 2] /= 255.0 * 0.229
|
||||
im_padded[:, :, 1] /= 255.0 * 0.224
|
||||
im_padded[:, :, 0] /= 255.0 * 0.225
|
||||
im_padded = im_padded.transpose((2, 0, 1))
|
||||
|
||||
data["image"] = im_padded[::-1, :, :]
|
||||
data["score_map"] = score_map[np.newaxis, :, :]
|
||||
data["border_map"] = border_map.transpose((2, 0, 1))
|
||||
data["training_mask"] = training_mask[np.newaxis, :, :]
|
||||
data["tvo_map"] = tvo_map.transpose((2, 0, 1))
|
||||
data["tco_map"] = tco_map.transpose((2, 0, 1))
|
||||
return data
|
||||
55
ppocr/data/imaug/ssl_img_aug.py
Normal file
55
ppocr/data/imaug/ssl_img_aug.py
Normal file
@@ -0,0 +1,55 @@
|
||||
# copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import math
|
||||
import cv2
|
||||
import numpy as np
|
||||
import random
|
||||
from PIL import Image
|
||||
|
||||
from .rec_img_aug import resize_norm_img
|
||||
|
||||
|
||||
class SSLRotateResize(object):
|
||||
def __init__(
|
||||
self, image_shape, padding=False, select_all=True, mode="train", **kwargs
|
||||
):
|
||||
self.image_shape = image_shape
|
||||
self.padding = padding
|
||||
self.select_all = select_all
|
||||
self.mode = mode
|
||||
|
||||
def __call__(self, data):
|
||||
img = data["image"]
|
||||
|
||||
data["image_r90"] = cv2.rotate(img, cv2.ROTATE_90_CLOCKWISE)
|
||||
data["image_r180"] = cv2.rotate(data["image_r90"], cv2.ROTATE_90_CLOCKWISE)
|
||||
data["image_r270"] = cv2.rotate(data["image_r180"], cv2.ROTATE_90_CLOCKWISE)
|
||||
|
||||
images = []
|
||||
for key in ["image", "image_r90", "image_r180", "image_r270"]:
|
||||
images.append(
|
||||
resize_norm_img(
|
||||
data.pop(key), image_shape=self.image_shape, padding=self.padding
|
||||
)[0]
|
||||
)
|
||||
data["image"] = np.stack(images, axis=0)
|
||||
data["label"] = np.array(list(range(4)))
|
||||
if not self.select_all:
|
||||
data["image"] = data["image"][0::2] # just choose 0 and 180
|
||||
data["label"] = data["label"][0:2] # label needs to be continuous
|
||||
if self.mode == "test":
|
||||
data["image"] = data["image"][0]
|
||||
data["label"] = data["label"][0]
|
||||
return data
|
||||
232
ppocr/data/imaug/table_ops.py
Normal file
232
ppocr/data/imaug/table_ops.py
Normal file
@@ -0,0 +1,232 @@
|
||||
"""
|
||||
# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
"""
|
||||
|
||||
from __future__ import absolute_import
|
||||
from __future__ import division
|
||||
from __future__ import print_function
|
||||
from __future__ import unicode_literals
|
||||
|
||||
import sys
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
|
||||
class GenTableMask(object):
|
||||
"""gen table mask"""
|
||||
|
||||
def __init__(self, shrink_h_max, shrink_w_max, mask_type=0, **kwargs):
|
||||
self.shrink_h_max = 5
|
||||
self.shrink_w_max = 5
|
||||
self.mask_type = mask_type
|
||||
|
||||
def projection(self, erosion, h, w, spilt_threshold=0):
|
||||
# 水平投影
|
||||
projection_map = np.ones_like(erosion)
|
||||
project_val_array = [0 for _ in range(0, h)]
|
||||
|
||||
for j in range(0, h):
|
||||
for i in range(0, w):
|
||||
if erosion[j, i] == 255:
|
||||
project_val_array[j] += 1
|
||||
# 根据数组，获取切割点
|
||||
start_idx = 0 # 记录进入字符区的索引
|
||||
end_idx = 0 # 记录进入空白区域的索引
|
||||
in_text = False # 是否遍历到了字符区内
|
||||
box_list = []
|
||||
for i in range(len(project_val_array)):
|
||||
if (
|
||||
in_text == False and project_val_array[i] > spilt_threshold
|
||||
): # 进入字符区了
|
||||
in_text = True
|
||||
start_idx = i
|
||||
elif (
|
||||
project_val_array[i] <= spilt_threshold and in_text == True
|
||||
): # 进入空白区了
|
||||
end_idx = i
|
||||
in_text = False
|
||||
if end_idx - start_idx <= 2:
|
||||
continue
|
||||
box_list.append((start_idx, end_idx + 1))
|
||||
|
||||
if in_text:
|
||||
box_list.append((start_idx, h - 1))
|
||||
# 绘制投影直方图
|
||||
for j in range(0, h):
|
||||
for i in range(0, project_val_array[j]):
|
||||
projection_map[j, i] = 0
|
||||
return box_list, projection_map
|
||||
|
||||
def projection_cx(self, box_img):
|
||||
box_gray_img = cv2.cvtColor(box_img, cv2.COLOR_BGR2GRAY)
|
||||
h, w = box_gray_img.shape
|
||||
# 灰度图片进行二值化处理
|
||||
ret, thresh1 = cv2.threshold(box_gray_img, 200, 255, cv2.THRESH_BINARY_INV)
|
||||
# 纵向腐蚀
|
||||
if h < w:
|
||||
kernel = np.ones((2, 1), np.uint8)
|
||||
erode = cv2.erode(thresh1, kernel, iterations=1)
|
||||
else:
|
||||
erode = thresh1
|
||||
# 水平膨胀
|
||||
kernel = np.ones((1, 5), np.uint8)
|
||||
erosion = cv2.dilate(erode, kernel, iterations=1)
|
||||
# 水平投影
|
||||
projection_map = np.ones_like(erosion)
|
||||
project_val_array = [0 for _ in range(0, h)]
|
||||
|
||||
for j in range(0, h):
|
||||
for i in range(0, w):
|
||||
if erosion[j, i] == 255:
|
||||
project_val_array[j] += 1
|
||||
# 根据数组，获取切割点
|
||||
start_idx = 0 # 记录进入字符区的索引
|
||||
end_idx = 0 # 记录进入空白区域的索引
|
||||
in_text = False # 是否遍历到了字符区内
|
||||
box_list = []
|
||||
spilt_threshold = 0
|
||||
for i in range(len(project_val_array)):
|
||||
if (
|
||||
in_text == False and project_val_array[i] > spilt_threshold
|
||||
): # 进入字符区了
|
||||
in_text = True
|
||||
start_idx = i
|
||||
elif (
|
||||
project_val_array[i] <= spilt_threshold and in_text == True
|
||||
): # 进入空白区了
|
||||
end_idx = i
|
||||
in_text = False
|
||||
if end_idx - start_idx <= 2:
|
||||
continue
|
||||
box_list.append((start_idx, end_idx + 1))
|
||||
|
||||
if in_text:
|
||||
box_list.append((start_idx, h - 1))
|
||||
# 绘制投影直方图
|
||||
for j in range(0, h):
|
||||
for i in range(0, project_val_array[j]):
|
||||
projection_map[j, i] = 0
|
||||
split_bbox_list = []
|
||||
if len(box_list) > 1:
|
||||
for i, (h_start, h_end) in enumerate(box_list):
|
||||
if i == 0:
|
||||
h_start = 0
|
||||
if i == len(box_list):
|
||||
h_end = h
|
||||
word_img = erosion[h_start : h_end + 1, :]
|
||||
word_h, word_w = word_img.shape
|
||||
w_split_list, w_projection_map = self.projection(
|
||||
word_img.T, word_w, word_h
|
||||
)
|
||||
w_start, w_end = w_split_list[0][0], w_split_list[-1][1]
|
||||
if h_start > 0:
|
||||
h_start -= 1
|
||||
h_end += 1
|
||||
word_img = box_img[h_start : h_end + 1 :, w_start : w_end + 1, :]
|
||||
split_bbox_list.append([w_start, h_start, w_end, h_end])
|
||||
else:
|
||||
split_bbox_list.append([0, 0, w, h])
|
||||
return split_bbox_list
|
||||
|
||||
def shrink_bbox(self, bbox):
|
||||
left, top, right, bottom = bbox
|
||||
sh_h = min(max(int((bottom - top) * 0.1), 1), self.shrink_h_max)
|
||||
sh_w = min(max(int((right - left) * 0.1), 1), self.shrink_w_max)
|
||||
left_new = left + sh_w
|
||||
right_new = right - sh_w
|
||||
top_new = top + sh_h
|
||||
bottom_new = bottom - sh_h
|
||||
if left_new >= right_new:
|
||||
left_new = left
|
||||
right_new = right
|
||||
if top_new >= bottom_new:
|
||||
top_new = top
|
||||
bottom_new = bottom
|
||||
return [left_new, top_new, right_new, bottom_new]
|
||||
|
||||
def __call__(self, data):
|
||||
img = data["image"]
|
||||
cells = data["cells"]
|
||||
height, width = img.shape[0:2]
|
||||
if self.mask_type == 1:
|
||||
mask_img = np.zeros((height, width), dtype=np.float32)
|
||||
else:
|
||||
mask_img = np.zeros((height, width, 3), dtype=np.float32)
|
||||
cell_num = len(cells)
|
||||
for cno in range(cell_num):
|
||||
if "bbox" in cells[cno]:
|
||||
bbox = cells[cno]["bbox"]
|
||||
left, top, right, bottom = bbox
|
||||
box_img = img[top:bottom, left:right, :].copy()
|
||||
split_bbox_list = self.projection_cx(box_img)
|
||||
for sno in range(len(split_bbox_list)):
|
||||
split_bbox_list[sno][0] += left
|
||||
split_bbox_list[sno][1] += top
|
||||
split_bbox_list[sno][2] += left
|
||||
split_bbox_list[sno][3] += top
|
||||
|
||||
for sno in range(len(split_bbox_list)):
|
||||
left, top, right, bottom = split_bbox_list[sno]
|
||||
left, top, right, bottom = self.shrink_bbox(
|
||||
[left, top, right, bottom]
|
||||
)
|
||||
if self.mask_type == 1:
|
||||
mask_img[top:bottom, left:right] = 1.0
|
||||
data["mask_img"] = mask_img
|
||||
else:
|
||||
mask_img[top:bottom, left:right, :] = (255, 255, 255)
|
||||
data["image"] = mask_img
|
||||
return data
|
||||
|
||||
|
||||
class ResizeTableImage(object):
|
||||
def __init__(self, max_len, resize_bboxes=False, infer_mode=False, **kwargs):
|
||||
super(ResizeTableImage, self).__init__()
|
||||
self.max_len = max_len
|
||||
self.resize_bboxes = resize_bboxes
|
||||
self.infer_mode = infer_mode
|
||||
|
||||
def __call__(self, data):
|
||||
img = data["image"]
|
||||
height, width = img.shape[0:2]
|
||||
ratio = self.max_len / (max(height, width) * 1.0)
|
||||
resize_h = int(height * ratio)
|
||||
resize_w = int(width * ratio)
|
||||
resize_img = cv2.resize(img, (resize_w, resize_h))
|
||||
if self.resize_bboxes and not self.infer_mode:
|
||||
data["bboxes"] = data["bboxes"] * ratio
|
||||
data["image"] = resize_img
|
||||
data["src_img"] = img
|
||||
data["shape"] = np.array([height, width, ratio, ratio])
|
||||
data["max_len"] = self.max_len
|
||||
return data
|
||||
|
||||
|
||||
class PaddingTableImage(object):
|
||||
def __init__(self, size, **kwargs):
|
||||
super(PaddingTableImage, self).__init__()
|
||||
self.size = size
|
||||
|
||||
def __call__(self, data):
|
||||
img = data["image"]
|
||||
pad_h, pad_w = self.size
|
||||
padding_img = np.zeros((pad_h, pad_w, 3), dtype=np.float32)
|
||||
height, width = img.shape[0:2]
|
||||
padding_img[0:height, 0:width, :] = img.copy()
|
||||
data["image"] = padding_img
|
||||
shape = data["shape"].tolist()
|
||||
shape.extend([pad_h, pad_w])
|
||||
data["shape"] = np.array(shape)
|
||||
return data
|
||||
17
ppocr/data/imaug/text_image_aug/__init__.py
Normal file
17
ppocr/data/imaug/text_image_aug/__init__.py
Normal file
@@ -0,0 +1,17 @@
|
||||
# copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from .augment import tia_perspective, tia_distort, tia_stretch
|
||||
|
||||
__all__ = ["tia_distort", "tia_stretch", "tia_perspective"]
|
||||
123
ppocr/data/imaug/text_image_aug/augment.py
Normal file
123
ppocr/data/imaug/text_image_aug/augment.py
Normal file
@@ -0,0 +1,123 @@
|
||||
# copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
"""
|
||||
This code is refer from:
|
||||
https://github.com/RubanSeven/Text-Image-Augmentation-python/blob/master/augment.py
|
||||
"""
|
||||
|
||||
import numpy as np
|
||||
from .warp_mls import WarpMLS
|
||||
|
||||
|
||||
def tia_distort(src, segment=4):
|
||||
img_h, img_w = src.shape[:2]
|
||||
|
||||
cut = img_w // segment
|
||||
thresh = cut // 3
|
||||
|
||||
src_pts = list()
|
||||
dst_pts = list()
|
||||
|
||||
src_pts.append([0, 0])
|
||||
src_pts.append([img_w, 0])
|
||||
src_pts.append([img_w, img_h])
|
||||
src_pts.append([0, img_h])
|
||||
|
||||
dst_pts.append([np.random.randint(thresh), np.random.randint(thresh)])
|
||||
dst_pts.append([img_w - np.random.randint(thresh), np.random.randint(thresh)])
|
||||
dst_pts.append(
|
||||
[img_w - np.random.randint(thresh), img_h - np.random.randint(thresh)]
|
||||
)
|
||||
dst_pts.append([np.random.randint(thresh), img_h - np.random.randint(thresh)])
|
||||
|
||||
half_thresh = thresh * 0.5
|
||||
|
||||
for cut_idx in np.arange(1, segment, 1):
|
||||
src_pts.append([cut * cut_idx, 0])
|
||||
src_pts.append([cut * cut_idx, img_h])
|
||||
dst_pts.append(
|
||||
[
|
||||
cut * cut_idx + np.random.randint(thresh) - half_thresh,
|
||||
np.random.randint(thresh) - half_thresh,
|
||||
]
|
||||
)
|
||||
dst_pts.append(
|
||||
[
|
||||
cut * cut_idx + np.random.randint(thresh) - half_thresh,
|
||||
img_h + np.random.randint(thresh) - half_thresh,
|
||||
]
|
||||
)
|
||||
|
||||
trans = WarpMLS(src, src_pts, dst_pts, img_w, img_h)
|
||||
dst = trans.generate()
|
||||
|
||||
return dst
|
||||
|
||||
|
||||
def tia_stretch(src, segment=4):
|
||||
img_h, img_w = src.shape[:2]
|
||||
|
||||
cut = img_w // segment
|
||||
thresh = cut * 4 // 5
|
||||
|
||||
src_pts = list()
|
||||
dst_pts = list()
|
||||
|
||||
src_pts.append([0, 0])
|
||||
src_pts.append([img_w, 0])
|
||||
src_pts.append([img_w, img_h])
|
||||
src_pts.append([0, img_h])
|
||||
|
||||
dst_pts.append([0, 0])
|
||||
dst_pts.append([img_w, 0])
|
||||
dst_pts.append([img_w, img_h])
|
||||
dst_pts.append([0, img_h])
|
||||
|
||||
half_thresh = thresh * 0.5
|
||||
|
||||
for cut_idx in np.arange(1, segment, 1):
|
||||
move = np.random.randint(thresh) - half_thresh
|
||||
src_pts.append([cut * cut_idx, 0])
|
||||
src_pts.append([cut * cut_idx, img_h])
|
||||
dst_pts.append([cut * cut_idx + move, 0])
|
||||
dst_pts.append([cut * cut_idx + move, img_h])
|
||||
|
||||
trans = WarpMLS(src, src_pts, dst_pts, img_w, img_h)
|
||||
dst = trans.generate()
|
||||
|
||||
return dst
|
||||
|
||||
|
||||
def tia_perspective(src):
|
||||
img_h, img_w = src.shape[:2]
|
||||
|
||||
thresh = img_h // 2
|
||||
|
||||
src_pts = list()
|
||||
dst_pts = list()
|
||||
|
||||
src_pts.append([0, 0])
|
||||
src_pts.append([img_w, 0])
|
||||
src_pts.append([img_w, img_h])
|
||||
src_pts.append([0, img_h])
|
||||
|
||||
dst_pts.append([0, np.random.randint(thresh)])
|
||||
dst_pts.append([img_w, np.random.randint(thresh)])
|
||||
dst_pts.append([img_w, img_h - np.random.randint(thresh)])
|
||||
dst_pts.append([0, img_h - np.random.randint(thresh)])
|
||||
|
||||
trans = WarpMLS(src, src_pts, dst_pts, img_w, img_h)
|
||||
dst = trans.generate()
|
||||
|
||||
return dst
|
||||
187
ppocr/data/imaug/text_image_aug/warp_mls.py
Normal file
187
ppocr/data/imaug/text_image_aug/warp_mls.py
Normal file
@@ -0,0 +1,187 @@
|
||||
# copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
"""
|
||||
This code is refer from:
|
||||
https://github.com/RubanSeven/Text-Image-Augmentation-python/blob/master/warp_mls.py
|
||||
"""
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
class WarpMLS:
|
||||
def __init__(self, src, src_pts, dst_pts, dst_w, dst_h, trans_ratio=1.0):
|
||||
self.src = src
|
||||
self.src_pts = src_pts
|
||||
self.dst_pts = dst_pts
|
||||
self.pt_count = len(self.dst_pts)
|
||||
self.dst_w = dst_w
|
||||
self.dst_h = dst_h
|
||||
self.trans_ratio = trans_ratio
|
||||
self.grid_size = 100
|
||||
self.rdx = np.zeros((self.dst_h, self.dst_w))
|
||||
self.rdy = np.zeros((self.dst_h, self.dst_w))
|
||||
|
||||
@staticmethod
|
||||
def __bilinear_interp(x, y, v11, v12, v21, v22):
|
||||
return (v11 * (1 - y) + v12 * y) * (1 - x) + (v21 * (1 - y) + v22 * y) * x
|
||||
|
||||
def generate(self):
|
||||
self.calc_delta()
|
||||
return self.gen_img()
|
||||
|
||||
def calc_delta(self):
|
||||
w = np.zeros(self.pt_count, dtype=np.float32)
|
||||
|
||||
if self.pt_count < 2:
|
||||
return
|
||||
|
||||
i = 0
|
||||
while 1:
|
||||
if self.dst_w <= i < self.dst_w + self.grid_size - 1:
|
||||
i = self.dst_w - 1
|
||||
elif i >= self.dst_w:
|
||||
break
|
||||
|
||||
j = 0
|
||||
while 1:
|
||||
if self.dst_h <= j < self.dst_h + self.grid_size - 1:
|
||||
j = self.dst_h - 1
|
||||
elif j >= self.dst_h:
|
||||
break
|
||||
|
||||
sw = 0
|
||||
swp = np.zeros(2, dtype=np.float32)
|
||||
swq = np.zeros(2, dtype=np.float32)
|
||||
new_pt = np.zeros(2, dtype=np.float32)
|
||||
cur_pt = np.array([i, j], dtype=np.float32)
|
||||
|
||||
k = 0
|
||||
for k in range(self.pt_count):
|
||||
if i == self.dst_pts[k][0] and j == self.dst_pts[k][1]:
|
||||
break
|
||||
|
||||
w[k] = 1.0 / (
|
||||
(i - self.dst_pts[k][0]) * (i - self.dst_pts[k][0])
|
||||
+ (j - self.dst_pts[k][1]) * (j - self.dst_pts[k][1])
|
||||
)
|
||||
|
||||
sw += w[k]
|
||||
swp = swp + w[k] * np.array(self.dst_pts[k])
|
||||
swq = swq + w[k] * np.array(self.src_pts[k])
|
||||
|
||||
if k == self.pt_count - 1:
|
||||
pstar = 1 / sw * swp
|
||||
qstar = 1 / sw * swq
|
||||
|
||||
miu_s = 0
|
||||
for k in range(self.pt_count):
|
||||
if i == self.dst_pts[k][0] and j == self.dst_pts[k][1]:
|
||||
continue
|
||||
pt_i = self.dst_pts[k] - pstar
|
||||
miu_s += w[k] * np.sum(pt_i * pt_i)
|
||||
|
||||
cur_pt -= pstar
|
||||
cur_pt_j = np.array([-cur_pt[1], cur_pt[0]])
|
||||
|
||||
for k in range(self.pt_count):
|
||||
if i == self.dst_pts[k][0] and j == self.dst_pts[k][1]:
|
||||
continue
|
||||
|
||||
pt_i = self.dst_pts[k] - pstar
|
||||
pt_j = np.array([-pt_i[1], pt_i[0]])
|
||||
|
||||
tmp_pt = np.zeros(2, dtype=np.float32)
|
||||
tmp_pt[0] = (
|
||||
np.sum(pt_i * cur_pt) * self.src_pts[k][0]
|
||||
- np.sum(pt_j * cur_pt) * self.src_pts[k][1]
|
||||
)
|
||||
tmp_pt[1] = (
|
||||
-np.sum(pt_i * cur_pt_j) * self.src_pts[k][0]
|
||||
+ np.sum(pt_j * cur_pt_j) * self.src_pts[k][1]
|
||||
)
|
||||
tmp_pt *= w[k] / miu_s
|
||||
new_pt += tmp_pt
|
||||
|
||||
new_pt += qstar
|
||||
else:
|
||||
new_pt = self.src_pts[k]
|
||||
|
||||
self.rdx[j, i] = new_pt[0] - i
|
||||
self.rdy[j, i] = new_pt[1] - j
|
||||
|
||||
j += self.grid_size
|
||||
i += self.grid_size
|
||||
|
||||
def gen_img(self):
|
||||
src_h, src_w = self.src.shape[:2]
|
||||
dst = np.zeros_like(self.src, dtype=np.float32)
|
||||
|
||||
for i in np.arange(0, self.dst_h, self.grid_size):
|
||||
for j in np.arange(0, self.dst_w, self.grid_size):
|
||||
ni = i + self.grid_size
|
||||
nj = j + self.grid_size
|
||||
w = h = self.grid_size
|
||||
if ni >= self.dst_h:
|
||||
ni = self.dst_h - 1
|
||||
h = ni - i + 1
|
||||
if nj >= self.dst_w:
|
||||
nj = self.dst_w - 1
|
||||
w = nj - j + 1
|
||||
|
||||
di = np.reshape(np.arange(h), (-1, 1))
|
||||
dj = np.reshape(np.arange(w), (1, -1))
|
||||
delta_x = self.__bilinear_interp(
|
||||
di / h,
|
||||
dj / w,
|
||||
self.rdx[i, j],
|
||||
self.rdx[i, nj],
|
||||
self.rdx[ni, j],
|
||||
self.rdx[ni, nj],
|
||||
)
|
||||
delta_y = self.__bilinear_interp(
|
||||
di / h,
|
||||
dj / w,
|
||||
self.rdy[i, j],
|
||||
self.rdy[i, nj],
|
||||
self.rdy[ni, j],
|
||||
self.rdy[ni, nj],
|
||||
)
|
||||
nx = j + dj + delta_x * self.trans_ratio
|
||||
ny = i + di + delta_y * self.trans_ratio
|
||||
nx = np.clip(nx, 0, src_w - 1)
|
||||
ny = np.clip(ny, 0, src_h - 1)
|
||||
nxi = np.array(np.floor(nx), dtype=np.int32)
|
||||
nyi = np.array(np.floor(ny), dtype=np.int32)
|
||||
nxi1 = np.array(np.ceil(nx), dtype=np.int32)
|
||||
nyi1 = np.array(np.ceil(ny), dtype=np.int32)
|
||||
|
||||
if len(self.src.shape) == 3:
|
||||
x = np.tile(np.expand_dims(ny - nyi, axis=-1), (1, 1, 3))
|
||||
y = np.tile(np.expand_dims(nx - nxi, axis=-1), (1, 1, 3))
|
||||
else:
|
||||
x = ny - nyi
|
||||
y = nx - nxi
|
||||
dst[i : i + h, j : j + w] = self.__bilinear_interp(
|
||||
x,
|
||||
y,
|
||||
self.src[nyi, nxi],
|
||||
self.src[nyi, nxi1],
|
||||
self.src[nyi1, nxi],
|
||||
self.src[nyi1, nxi1],
|
||||
)
|
||||
|
||||
dst = np.clip(dst, 0, 255)
|
||||
dst = np.array(dst, dtype=np.uint8)
|
||||
|
||||
return dst
|
||||
807
ppocr/data/imaug/unimernet_aug.py
Normal file
807
ppocr/data/imaug/unimernet_aug.py
Normal file
@@ -0,0 +1,807 @@
|
||||
# copyright (c) 2024 PaddlePaddle Authors. All Rights Reserve.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from __future__ import absolute_import
|
||||
from __future__ import division
|
||||
from __future__ import print_function
|
||||
from __future__ import unicode_literals
|
||||
|
||||
import os
|
||||
|
||||
os.environ["NO_ALBUMENTATIONS_UPDATE"] = "1"
|
||||
|
||||
import cv2
|
||||
import math
|
||||
import numpy as np
|
||||
from io import BytesIO
|
||||
import albumentations as A
|
||||
from PIL import Image, ImageOps, ImageDraw
|
||||
from scipy.ndimage import zoom as scizoom
|
||||
|
||||
|
||||
class Erosion(A.ImageOnlyTransform):
|
||||
def __init__(self, scale, always_apply=False, p=0.5):
|
||||
super().__init__(always_apply=always_apply, p=p)
|
||||
if type(scale) is tuple or type(scale) is list:
|
||||
assert len(scale) == 2
|
||||
self.scale = scale
|
||||
else:
|
||||
self.scale = (scale, scale)
|
||||
|
||||
def apply(self, img, **params):
|
||||
kernel = cv2.getStructuringElement(
|
||||
cv2.MORPH_ELLIPSE, tuple(np.random.randint(self.scale[0], self.scale[1], 2))
|
||||
)
|
||||
img = cv2.erode(img, kernel, iterations=1)
|
||||
return img
|
||||
|
||||
|
||||
class Dilation(A.ImageOnlyTransform):
|
||||
def __init__(self, scale, always_apply=False, p=0.5):
|
||||
super().__init__(always_apply=always_apply, p=p)
|
||||
if type(scale) is tuple or type(scale) is list:
|
||||
assert len(scale) == 2
|
||||
self.scale = scale
|
||||
else:
|
||||
self.scale = (scale, scale)
|
||||
|
||||
def apply(self, img, **params):
|
||||
kernel = cv2.getStructuringElement(
|
||||
cv2.MORPH_ELLIPSE, tuple(np.random.randint(self.scale[0], self.scale[1], 2))
|
||||
)
|
||||
img = cv2.dilate(img, kernel, iterations=1)
|
||||
return img
|
||||
|
||||
|
||||
class Bitmap(A.ImageOnlyTransform):
|
||||
|
||||
def __init__(self, value=0, lower=200, always_apply=False, p=0.5):
|
||||
super().__init__(always_apply=always_apply, p=p)
|
||||
self.lower = lower
|
||||
self.value = value
|
||||
|
||||
def apply(self, img, **params):
|
||||
img = img.copy()
|
||||
img[img < self.lower] = self.value
|
||||
return img
|
||||
|
||||
|
||||
def clipped_zoom(img, zoom_factor):
|
||||
h = img.shape[1]
|
||||
ch = int(np.ceil(h / float(zoom_factor)))
|
||||
top = (h - ch) // 2
|
||||
img = scizoom(
|
||||
img[top : top + ch, top : top + ch], (zoom_factor, zoom_factor, 1), order=1
|
||||
)
|
||||
trim_top = (img.shape[0] - h) // 2
|
||||
|
||||
return img[trim_top : trim_top + h, trim_top : trim_top + h]
|
||||
|
||||
|
||||
def disk(radius, alias_blur=0.1, dtype=np.float32):
|
||||
if radius <= 8:
|
||||
coords = np.arange(-8, 8 + 1)
|
||||
ksize = (3, 3)
|
||||
else:
|
||||
coords = np.arange(-radius, radius + 1)
|
||||
ksize = (5, 5)
|
||||
x, y = np.meshgrid(coords, coords)
|
||||
aliased_disk = np.asarray((x**2 + y**2) <= radius**2, dtype=dtype)
|
||||
aliased_disk /= np.sum(aliased_disk)
|
||||
|
||||
return cv2.GaussianBlur(aliased_disk, ksize=ksize, sigmaX=alias_blur)
|
||||
|
||||
|
||||
def plasma_fractal(mapsize=256, wibbledecay=3, rng=None):
|
||||
"""
|
||||
Generate a heightmap using diamond-square algorithm.
|
||||
Return square 2d array, side length 'mapsize', of floats in range 0-255.
|
||||
'mapsize' must be a power of two.
|
||||
"""
|
||||
assert mapsize & (mapsize - 1) == 0
|
||||
maparray = np.empty((mapsize, mapsize), dtype=np.float_)
|
||||
maparray[0, 0] = 0
|
||||
stepsize = mapsize
|
||||
wibble = 100
|
||||
if rng is None:
|
||||
rng = np.random.default_rng()
|
||||
|
||||
def wibbledmean(array):
|
||||
return array / 4 + wibble * rng.uniform(-wibble, wibble, array.shape)
|
||||
|
||||
def fillsquares():
|
||||
"""For each square of points stepsize apart,
|
||||
calculate middle value as mean of points + wibble"""
|
||||
cornerref = maparray[0:mapsize:stepsize, 0:mapsize:stepsize]
|
||||
squareaccum = cornerref + np.roll(cornerref, shift=-1, axis=0)
|
||||
squareaccum += np.roll(squareaccum, shift=-1, axis=1)
|
||||
maparray[
|
||||
stepsize // 2 : mapsize : stepsize, stepsize // 2 : mapsize : stepsize
|
||||
] = wibbledmean(squareaccum)
|
||||
|
||||
def filldiamonds():
|
||||
"""For each diamond of points stepsize apart,
|
||||
calculate middle value as mean of points + wibble"""
|
||||
drgrid = maparray[
|
||||
stepsize // 2 : mapsize : stepsize, stepsize // 2 : mapsize : stepsize
|
||||
]
|
||||
ulgrid = maparray[0:mapsize:stepsize, 0:mapsize:stepsize]
|
||||
ldrsum = drgrid + np.roll(drgrid, 1, axis=0)
|
||||
lulsum = ulgrid + np.roll(ulgrid, -1, axis=1)
|
||||
ltsum = ldrsum + lulsum
|
||||
maparray[0:mapsize:stepsize, stepsize // 2 : mapsize : stepsize] = wibbledmean(
|
||||
ltsum
|
||||
)
|
||||
tdrsum = drgrid + np.roll(drgrid, 1, axis=1)
|
||||
tulsum = ulgrid + np.roll(ulgrid, -1, axis=0)
|
||||
ttsum = tdrsum + tulsum
|
||||
maparray[stepsize // 2 : mapsize : stepsize, 0:mapsize:stepsize] = wibbledmean(
|
||||
ttsum
|
||||
)
|
||||
|
||||
while stepsize >= 2:
|
||||
fillsquares()
|
||||
filldiamonds()
|
||||
stepsize //= 2
|
||||
wibble /= wibbledecay
|
||||
|
||||
maparray -= maparray.min()
|
||||
return maparray / maparray.max()
|
||||
|
||||
|
||||
class Fog(A.ImageOnlyTransform):
|
||||
def __init__(self, mag=-1, always_apply=False, p=1.0):
|
||||
super().__init__(always_apply=always_apply, p=p)
|
||||
self.rng = np.random.default_rng()
|
||||
self.mag = mag
|
||||
|
||||
def apply(self, img, **params):
|
||||
img = Image.fromarray(img.astype(np.uint8))
|
||||
w, h = img.size
|
||||
c = [(1.5, 2), (2.0, 2), (2.5, 1.7)]
|
||||
if self.mag < 0 or self.mag >= len(c):
|
||||
index = self.rng.integers(0, len(c))
|
||||
else:
|
||||
index = self.mag
|
||||
c = c[index]
|
||||
|
||||
n_channels = len(img.getbands())
|
||||
isgray = n_channels == 1
|
||||
|
||||
img = np.asarray(img) / 255.0
|
||||
max_val = img.max()
|
||||
max_size = 2 ** math.ceil(math.log2(max(w, h)) + 1)
|
||||
fog = (
|
||||
c[0]
|
||||
* plasma_fractal(mapsize=max_size, wibbledecay=c[1], rng=self.rng)[:h, :w][
|
||||
..., np.newaxis
|
||||
]
|
||||
)
|
||||
if isgray:
|
||||
fog = np.squeeze(fog)
|
||||
else:
|
||||
fog = np.repeat(fog, 3, axis=2)
|
||||
|
||||
img += fog
|
||||
img = np.clip(img * max_val / (max_val + c[0]), 0, 1) * 255
|
||||
return img.astype(np.uint8)
|
||||
|
||||
|
||||
class Frost(A.ImageOnlyTransform):
|
||||
def __init__(self, mag=-1, always_apply=False, p=1.0):
|
||||
super().__init__(always_apply=always_apply, p=p)
|
||||
self.rng = np.random.default_rng()
|
||||
self.mag = mag
|
||||
|
||||
def apply(self, img, **params):
|
||||
img = Image.fromarray(img.astype(np.uint8))
|
||||
w, h = img.size
|
||||
c = [(0.78, 0.22), (0.64, 0.36), (0.5, 0.5)]
|
||||
if self.mag < 0 or self.mag >= len(c):
|
||||
index = self.rng.integers(0, len(c))
|
||||
else:
|
||||
index = self.mag
|
||||
c = c[index]
|
||||
|
||||
file_dir = os.path.dirname(__file__)
|
||||
filename = [
|
||||
os.path.join(file_dir, "frost_img", "frost1.jpg"),
|
||||
os.path.join(file_dir, "frost_img", "frost2.png"),
|
||||
os.path.join(file_dir, "frost_img", "frost3.png"),
|
||||
os.path.join(file_dir, "frost_img", "frost4.jpg"),
|
||||
os.path.join(file_dir, "frost_img", "frost5.jpg"),
|
||||
os.path.join(file_dir, "frost_img", "frost6.jpg"),
|
||||
]
|
||||
index = self.rng.integers(0, len(filename))
|
||||
filename = filename[index]
|
||||
frost = Image.open(filename).convert("RGB")
|
||||
|
||||
f_w, f_h = frost.size
|
||||
if w / h > f_w / f_h:
|
||||
f_h = round(f_h * w / f_w)
|
||||
f_w = w
|
||||
else:
|
||||
f_w = round(f_w * h / f_h)
|
||||
f_h = h
|
||||
frost = np.asarray(frost.resize((f_w, f_h)))
|
||||
|
||||
# randomly crop
|
||||
y_start, x_start = self.rng.integers(0, f_h - h + 1), self.rng.integers(
|
||||
0, f_w - w + 1
|
||||
)
|
||||
frost = frost[y_start : y_start + h, x_start : x_start + w]
|
||||
|
||||
n_channels = len(img.getbands())
|
||||
isgray = n_channels == 1
|
||||
|
||||
img = np.asarray(img)
|
||||
|
||||
if isgray:
|
||||
img = np.expand_dims(img, axis=2)
|
||||
img = np.repeat(img, 3, axis=2)
|
||||
|
||||
img = np.clip(np.round(c[0] * img + c[1] * frost), 0, 255)
|
||||
img = img.astype(np.uint8)
|
||||
if isgray:
|
||||
img = np.squeeze(img)
|
||||
return img
|
||||
|
||||
|
||||
class Snow(A.ImageOnlyTransform):
|
||||
def __init__(self, mag=-1, always_apply=False, p=1.0):
|
||||
super().__init__(always_apply=always_apply, p=p)
|
||||
self.rng = np.random.default_rng()
|
||||
self.mag = mag
|
||||
|
||||
def apply(self, img, **params):
|
||||
from wand.image import Image as WandImage
|
||||
|
||||
img = Image.fromarray(img.astype(np.uint8))
|
||||
w, h = img.size
|
||||
c = [
|
||||
(0.1, 0.3, 3, 0.5, 10, 4, 0.8),
|
||||
(0.2, 0.3, 2, 0.5, 12, 4, 0.7),
|
||||
(0.55, 0.3, 4, 0.9, 12, 8, 0.7),
|
||||
]
|
||||
if self.mag < 0 or self.mag >= len(c):
|
||||
index = self.rng.integers(0, len(c))
|
||||
else:
|
||||
index = self.mag
|
||||
c = c[index]
|
||||
|
||||
n_channels = len(img.getbands())
|
||||
isgray = n_channels == 1
|
||||
|
||||
img = np.asarray(img, dtype=np.float32) / 255.0
|
||||
if isgray:
|
||||
img = np.expand_dims(img, axis=2)
|
||||
img = np.repeat(img, 3, axis=2)
|
||||
|
||||
snow_layer = self.rng.normal(size=img.shape[:2], loc=c[0], scale=c[1])
|
||||
|
||||
snow_layer[snow_layer < c[3]] = 0
|
||||
|
||||
snow_layer = Image.fromarray(
|
||||
(np.clip(snow_layer.squeeze(), 0, 1) * 255).astype(np.uint8), mode="L"
|
||||
)
|
||||
output = BytesIO()
|
||||
snow_layer.save(output, format="PNG")
|
||||
snow_layer = WandImage(blob=output.getvalue())
|
||||
|
||||
snow_layer.motion_blur(
|
||||
radius=c[4], sigma=c[5], angle=self.rng.uniform(-135, -45)
|
||||
)
|
||||
|
||||
snow_layer = (
|
||||
cv2.imdecode(
|
||||
np.frombuffer(snow_layer.make_blob(), np.uint8), cv2.IMREAD_UNCHANGED
|
||||
)
|
||||
/ 255.0
|
||||
)
|
||||
|
||||
snow_layer = snow_layer[..., np.newaxis]
|
||||
|
||||
img = c[6] * img
|
||||
gray_img = (1 - c[6]) * np.maximum(
|
||||
img, cv2.cvtColor(img, cv2.COLOR_RGB2GRAY).reshape(h, w, 1) * 1.5 + 0.5
|
||||
)
|
||||
img += gray_img
|
||||
img = np.clip(img + snow_layer + np.rot90(snow_layer, k=2), 0, 1) * 255
|
||||
img = img.astype(np.uint8)
|
||||
if isgray:
|
||||
img = np.squeeze(img)
|
||||
return img
|
||||
|
||||
|
||||
class Rain(A.ImageOnlyTransform):
|
||||
def __init__(self, mag=-1, always_apply=False, p=1.0):
|
||||
super().__init__(always_apply=always_apply, p=p)
|
||||
self.rng = np.random.default_rng()
|
||||
self.mag = mag
|
||||
|
||||
def apply(self, img, **params):
|
||||
img = Image.fromarray(img.astype(np.uint8))
|
||||
img = img.copy()
|
||||
w, h = img.size
|
||||
n_channels = len(img.getbands())
|
||||
isgray = n_channels == 1
|
||||
line_width = self.rng.integers(1, 2)
|
||||
|
||||
c = [50, 70, 90]
|
||||
if self.mag < 0 or self.mag >= len(c):
|
||||
index = 0
|
||||
else:
|
||||
index = self.mag
|
||||
c = c[index]
|
||||
|
||||
n_rains = self.rng.integers(c, c + 20)
|
||||
slant = self.rng.integers(-60, 60)
|
||||
fillcolor = 200 if isgray else (200, 200, 200)
|
||||
|
||||
draw = ImageDraw.Draw(img)
|
||||
max_length = min(w, h, 10)
|
||||
for i in range(1, n_rains):
|
||||
length = self.rng.integers(5, max_length)
|
||||
x1 = self.rng.integers(0, w - length)
|
||||
y1 = self.rng.integers(0, h - length)
|
||||
x2 = x1 + length * math.sin(slant * math.pi / 180.0)
|
||||
y2 = y1 + length * math.cos(slant * math.pi / 180.0)
|
||||
x2 = int(x2)
|
||||
y2 = int(y2)
|
||||
draw.line([(x1, y1), (x2, y2)], width=line_width, fill=fillcolor)
|
||||
img = np.asarray(img).astype(np.uint8)
|
||||
return img
|
||||
|
||||
|
||||
class Shadow(A.ImageOnlyTransform):
|
||||
def __init__(self, mag=-1, always_apply=False, p=1.0):
|
||||
super().__init__(always_apply=always_apply, p=p)
|
||||
self.rng = np.random.default_rng()
|
||||
self.mag = mag
|
||||
|
||||
def apply(self, img, **params):
|
||||
img = Image.fromarray(img.astype(np.uint8))
|
||||
w, h = img.size
|
||||
n_channels = len(img.getbands())
|
||||
isgray = n_channels == 1
|
||||
|
||||
c = [64, 96, 128]
|
||||
if self.mag < 0 or self.mag >= len(c):
|
||||
index = 0
|
||||
else:
|
||||
index = self.mag
|
||||
c = c[index]
|
||||
|
||||
img = img.convert("RGBA")
|
||||
overlay = Image.new("RGBA", img.size, (255, 255, 255, 0))
|
||||
draw = ImageDraw.Draw(overlay)
|
||||
transparency = self.rng.integers(c, c + 32)
|
||||
x1 = self.rng.integers(0, w // 2)
|
||||
y1 = 0
|
||||
|
||||
x2 = self.rng.integers(w // 2, w)
|
||||
y2 = 0
|
||||
|
||||
x3 = self.rng.integers(w // 2, w)
|
||||
y3 = h - 1
|
||||
|
||||
x4 = self.rng.integers(0, w // 2)
|
||||
y4 = h - 1
|
||||
|
||||
draw.polygon(
|
||||
[(x1, y1), (x2, y2), (x3, y3), (x4, y4)], fill=(0, 0, 0, transparency)
|
||||
)
|
||||
|
||||
img = Image.alpha_composite(img, overlay)
|
||||
img = img.convert("RGB")
|
||||
if isgray:
|
||||
img = ImageOps.grayscale(img)
|
||||
img = np.asarray(img).astype(np.uint8)
|
||||
return img
|
||||
|
||||
|
||||
class UniMERNetTrainTransform:
|
||||
def __init__(self, bitmap_prob=0.04, **kwargs):
|
||||
self.bitmap_prob = bitmap_prob
|
||||
self.train_transform = A.Compose(
|
||||
[
|
||||
A.Compose(
|
||||
[
|
||||
Bitmap(p=0.05),
|
||||
A.OneOf([Fog(), Frost(), Snow(), Rain(), Shadow()], p=0.2),
|
||||
A.OneOf([Erosion((2, 3)), Dilation((2, 3))], p=0.2),
|
||||
A.ShiftScaleRotate(
|
||||
shift_limit=0,
|
||||
scale_limit=(-0.15, 0),
|
||||
rotate_limit=1,
|
||||
border_mode=0,
|
||||
interpolation=3,
|
||||
value=[255, 255, 255],
|
||||
p=1,
|
||||
),
|
||||
A.GridDistortion(
|
||||
distort_limit=0.1,
|
||||
border_mode=0,
|
||||
interpolation=3,
|
||||
value=[255, 255, 255],
|
||||
p=0.5,
|
||||
),
|
||||
],
|
||||
p=0.15,
|
||||
),
|
||||
A.RGBShift(r_shift_limit=15, g_shift_limit=15, b_shift_limit=15, p=0.3),
|
||||
A.GaussNoise(10, p=0.2),
|
||||
A.RandomBrightnessContrast(0.05, (-0.2, 0), True, p=0.2),
|
||||
A.ImageCompression(95, p=0.3),
|
||||
A.ToGray(always_apply=True),
|
||||
A.Normalize((0.7931, 0.7931, 0.7931), (0.1738, 0.1738, 0.1738)),
|
||||
]
|
||||
)
|
||||
|
||||
def __call__(self, data):
|
||||
img = data["image"]
|
||||
if np.random.random() < self.bitmap_prob:
|
||||
img[img != 255] = 0
|
||||
img = self.train_transform(image=img)["image"]
|
||||
data["image"] = img
|
||||
return data
|
||||
|
||||
|
||||
class UniMERNetTestTransform:
|
||||
def __init__(self, **kwargs):
|
||||
self.test_transform = A.Compose(
|
||||
[
|
||||
A.ToGray(always_apply=True),
|
||||
A.Normalize((0.7931, 0.7931, 0.7931), (0.1738, 0.1738, 0.1738)),
|
||||
]
|
||||
)
|
||||
|
||||
def __call__(self, data):
|
||||
img = data["image"]
|
||||
img = self.test_transform(image=img)["image"]
|
||||
data["image"] = img
|
||||
return data
|
||||
|
||||
|
||||
class GoTImgDecode:
|
||||
def __init__(self, input_size, random_padding=False, **kwargs):
|
||||
self.input_size = input_size
|
||||
self.random_padding = random_padding
|
||||
|
||||
def crop_margin(self, img):
|
||||
data = np.array(img.convert("L"))
|
||||
data = data.astype(np.uint8)
|
||||
max_val = data.max()
|
||||
min_val = data.min()
|
||||
if max_val == min_val:
|
||||
return img
|
||||
data = (data - min_val) / (max_val - min_val) * 255
|
||||
gray = 255 * (data < 200).astype(np.uint8)
|
||||
coords = cv2.findNonZero(gray) # Find all non-zero points (text)
|
||||
a, b, w, h = cv2.boundingRect(coords) # Find minimum spanning bounding box
|
||||
return img.crop((a, b, w + a, h + b))
|
||||
|
||||
def get_dimensions(self, img):
|
||||
if hasattr(img, "getbands"):
|
||||
channels = len(img.getbands())
|
||||
else:
|
||||
channels = img.channels
|
||||
width, height = img.size
|
||||
return [channels, height, width]
|
||||
|
||||
def _compute_resized_output_size(self, image_size, size, max_size=None):
|
||||
if len(size) == 1: # specified size only for the smallest edge
|
||||
h, w = image_size
|
||||
short, long = (w, h) if w <= h else (h, w)
|
||||
requested_new_short = size if isinstance(size, int) else size[0]
|
||||
|
||||
new_short, new_long = requested_new_short, int(
|
||||
requested_new_short * long / short
|
||||
)
|
||||
|
||||
if max_size is not None:
|
||||
if max_size <= requested_new_short:
|
||||
raise ValueError(
|
||||
f"max_size = {max_size} must be strictly greater than the requested "
|
||||
f"size for the smaller edge size = {size}"
|
||||
)
|
||||
if new_long > max_size:
|
||||
new_short, new_long = int(max_size * new_short / new_long), max_size
|
||||
|
||||
new_w, new_h = (new_short, new_long) if w <= h else (new_long, new_short)
|
||||
else: # specified both h and w
|
||||
new_w, new_h = size[1], size[0]
|
||||
return [new_h, new_w]
|
||||
|
||||
def resize(self, img, size):
|
||||
_, image_height, image_width = self.get_dimensions(img)
|
||||
if isinstance(size, int):
|
||||
size = [size]
|
||||
max_size = None
|
||||
output_size = self._compute_resized_output_size(
|
||||
(image_height, image_width), size, max_size
|
||||
)
|
||||
img = img.resize(tuple(output_size[::-1]), resample=2)
|
||||
return img
|
||||
|
||||
def __call__(self, data):
|
||||
filename = data["filename"]
|
||||
img = Image.open(filename)
|
||||
try:
|
||||
img = self.crop_margin(img.convert("RGB"))
|
||||
except OSError:
|
||||
return
|
||||
if img.height == 0 or img.width == 0:
|
||||
return
|
||||
img = self.resize(img, min(self.input_size))
|
||||
img.thumbnail((self.input_size[1], self.input_size[0]))
|
||||
delta_width = self.input_size[1] - img.width
|
||||
delta_height = self.input_size[0] - img.height
|
||||
if self.random_padding:
|
||||
pad_width = np.random.randint(low=0, high=delta_width + 1)
|
||||
pad_height = np.random.randint(low=0, high=delta_height + 1)
|
||||
else:
|
||||
pad_width = delta_width // 2
|
||||
pad_height = delta_height // 2
|
||||
padding = (
|
||||
pad_width,
|
||||
pad_height,
|
||||
delta_width - pad_width,
|
||||
delta_height - pad_height,
|
||||
)
|
||||
|
||||
data["image"] = np.array(ImageOps.expand(img, padding))
|
||||
return data
|
||||
|
||||
|
||||
class UniMERNetImgDecode:
|
||||
def __init__(
|
||||
self,
|
||||
input_size,
|
||||
random_padding=False,
|
||||
random_resize=False,
|
||||
random_crop=False,
|
||||
**kwargs,
|
||||
):
|
||||
self.input_size = input_size
|
||||
self.is_random_padding = random_padding
|
||||
self.is_random_resize = random_resize
|
||||
self.is_random_crop = random_crop
|
||||
|
||||
def crop_margin(self, img):
|
||||
data = np.array(img.convert("L"))
|
||||
data = data.astype(np.uint8)
|
||||
max_val = data.max()
|
||||
min_val = data.min()
|
||||
if max_val == min_val:
|
||||
return img
|
||||
data = (data - min_val) / (max_val - min_val) * 255
|
||||
gray = 255 * (data < 200).astype(np.uint8)
|
||||
coords = cv2.findNonZero(gray) # Find all non-zero points (text)
|
||||
a, b, w, h = cv2.boundingRect(coords) # Find minimum spanning bounding box
|
||||
return img.crop((a, b, w + a, h + b))
|
||||
|
||||
def get_dimensions(self, img):
|
||||
if hasattr(img, "getbands"):
|
||||
channels = len(img.getbands())
|
||||
else:
|
||||
channels = img.channels
|
||||
width, height = img.size
|
||||
return [channels, height, width]
|
||||
|
||||
def _compute_resized_output_size(self, image_size, size, max_size=None):
|
||||
if len(size) == 1: # specified size only for the smallest edge
|
||||
h, w = image_size
|
||||
short, long = (w, h) if w <= h else (h, w)
|
||||
requested_new_short = size if isinstance(size, int) else size[0]
|
||||
|
||||
new_short, new_long = requested_new_short, int(
|
||||
requested_new_short * long / short
|
||||
)
|
||||
|
||||
if max_size is not None:
|
||||
if max_size <= requested_new_short:
|
||||
raise ValueError(
|
||||
f"max_size = {max_size} must be strictly greater than the requested "
|
||||
f"size for the smaller edge size = {size}"
|
||||
)
|
||||
if new_long > max_size:
|
||||
new_short, new_long = int(max_size * new_short / new_long), max_size
|
||||
|
||||
new_w, new_h = (new_short, new_long) if w <= h else (new_long, new_short)
|
||||
else: # specified both h and w
|
||||
new_w, new_h = size[1], size[0]
|
||||
return [new_h, new_w]
|
||||
|
||||
def resize(self, img, size):
|
||||
_, image_height, image_width = self.get_dimensions(img)
|
||||
if isinstance(size, int):
|
||||
size = [size]
|
||||
max_size = None
|
||||
output_size = self._compute_resized_output_size(
|
||||
(image_height, image_width), size, max_size
|
||||
)
|
||||
img = img.resize(tuple(output_size[::-1]), resample=2)
|
||||
return img
|
||||
|
||||
def random_resize(self, img):
|
||||
scale = np.random.uniform(0.5, 1)
|
||||
img = img.resize([int(scale * s) for s in img.size])
|
||||
return img
|
||||
|
||||
def random_crop(self, img, crop_ratio):
|
||||
width, height = img.width, img.height
|
||||
max_crop_pixel = min(width, height) * crop_ratio
|
||||
crop_left = np.random.uniform(0, max_crop_pixel)
|
||||
crop_right = np.random.uniform(0, max_crop_pixel)
|
||||
crop_top = np.random.uniform(0, max_crop_pixel)
|
||||
crop_bottom = np.random.uniform(0, max_crop_pixel)
|
||||
# 计算裁剪后的边界
|
||||
left = crop_left
|
||||
top = crop_top
|
||||
right = width - crop_right
|
||||
bottom = height - crop_bottom
|
||||
# 裁剪图像
|
||||
img = img.crop((left, top, right, bottom))
|
||||
|
||||
return img
|
||||
|
||||
def __call__(self, data):
|
||||
filename = data["filename"]
|
||||
img = Image.open(filename)
|
||||
try:
|
||||
if self.is_random_resize:
|
||||
img = self.random_resize(img)
|
||||
img = self.crop_margin(img.convert("RGB"))
|
||||
if "label" in data and self.is_random_crop:
|
||||
label = data["label"]
|
||||
equation_length = len(label)
|
||||
if equation_length < 256:
|
||||
img = self.random_crop(img, crop_ratio=0.1)
|
||||
elif 256 < equation_length <= 512:
|
||||
img = self.random_crop(img, crop_ratio=0.05)
|
||||
else:
|
||||
img = self.random_crop(img, crop_ratio=0.03)
|
||||
except OSError:
|
||||
return
|
||||
if img.height == 0 or img.width == 0:
|
||||
return
|
||||
img = self.resize(img, min(self.input_size))
|
||||
img.thumbnail((self.input_size[1], self.input_size[0]))
|
||||
delta_width = self.input_size[1] - img.width
|
||||
delta_height = self.input_size[0] - img.height
|
||||
if self.is_random_padding:
|
||||
pad_width = np.random.randint(low=0, high=delta_width + 1)
|
||||
pad_height = np.random.randint(low=0, high=delta_height + 1)
|
||||
else:
|
||||
pad_width = delta_width // 2
|
||||
pad_height = delta_height // 2
|
||||
padding = (
|
||||
pad_width,
|
||||
pad_height,
|
||||
delta_width - pad_width,
|
||||
delta_height - pad_height,
|
||||
)
|
||||
|
||||
data["image"] = np.array(ImageOps.expand(img, padding))
|
||||
return data
|
||||
|
||||
|
||||
class UniMERNetResize:
|
||||
def __init__(self, input_size, random_padding=False, **kwargs):
|
||||
self.input_size = input_size
|
||||
self.random_padding = random_padding
|
||||
|
||||
def crop_margin(self, img):
|
||||
data = np.array(img.convert("L"))
|
||||
data = data.astype(np.uint8)
|
||||
max_val = data.max()
|
||||
min_val = data.min()
|
||||
if max_val == min_val:
|
||||
return img
|
||||
data = (data - min_val) / (max_val - min_val) * 255
|
||||
gray = 255 * (data < 200).astype(np.uint8)
|
||||
|
||||
coords = cv2.findNonZero(gray) # Find all non-zero points (text)
|
||||
a, b, w, h = cv2.boundingRect(coords) # Find minimum spanning bounding box
|
||||
return img.crop((a, b, w + a, h + b))
|
||||
|
||||
def get_dimensions(self, img):
|
||||
if hasattr(img, "getbands"):
|
||||
channels = len(img.getbands())
|
||||
else:
|
||||
channels = img.channels
|
||||
width, height = img.size
|
||||
return [channels, height, width]
|
||||
|
||||
def _compute_resized_output_size(self, image_size, size, max_size=None):
|
||||
if len(size) == 1: # specified size only for the smallest edge
|
||||
h, w = image_size
|
||||
short, long = (w, h) if w <= h else (h, w)
|
||||
requested_new_short = size if isinstance(size, int) else size[0]
|
||||
|
||||
new_short, new_long = requested_new_short, int(
|
||||
requested_new_short * long / short
|
||||
)
|
||||
|
||||
if max_size is not None:
|
||||
if max_size <= requested_new_short:
|
||||
raise ValueError(
|
||||
f"max_size = {max_size} must be strictly greater than the requested "
|
||||
f"size for the smaller edge size = {size}"
|
||||
)
|
||||
if new_long > max_size:
|
||||
new_short, new_long = int(max_size * new_short / new_long), max_size
|
||||
|
||||
new_w, new_h = (new_short, new_long) if w <= h else (new_long, new_short)
|
||||
else: # specified both h and w
|
||||
new_w, new_h = size[1], size[0]
|
||||
return [new_h, new_w]
|
||||
|
||||
def resize(self, img, size):
|
||||
_, image_height, image_width = self.get_dimensions(img)
|
||||
if isinstance(size, int):
|
||||
size = [size]
|
||||
max_size = None
|
||||
output_size = self._compute_resized_output_size(
|
||||
(image_height, image_width), size, max_size
|
||||
)
|
||||
img.resize(tuple(output_size[::-1]), resample=2)
|
||||
return img
|
||||
|
||||
def __call__(self, data):
|
||||
img = data["image"]
|
||||
img = Image.fromarray(img)
|
||||
try:
|
||||
img = self.crop_margin(img)
|
||||
except OSError:
|
||||
return
|
||||
if img.height == 0 or img.width == 0:
|
||||
return
|
||||
img = self.resize(img, min(self.input_size))
|
||||
img.thumbnail((self.input_size[1], self.input_size[0]))
|
||||
delta_width = self.input_size[1] - img.width
|
||||
delta_height = self.input_size[0] - img.height
|
||||
if self.random_padding:
|
||||
pad_width = np.random.randint(low=0, high=delta_width + 1)
|
||||
pad_height = np.random.randint(low=0, high=delta_height + 1)
|
||||
else:
|
||||
pad_width = delta_width // 2
|
||||
pad_height = delta_height // 2
|
||||
padding = (
|
||||
pad_width,
|
||||
pad_height,
|
||||
delta_width - pad_width,
|
||||
delta_height - pad_height,
|
||||
)
|
||||
data["image"] = np.array(ImageOps.expand(img, padding))
|
||||
return data
|
||||
|
||||
|
||||
class UniMERNetImageFormat:
|
||||
def __init__(self, **kwargs):
|
||||
pass
|
||||
|
||||
def __call__(self, data):
|
||||
img = data["image"]
|
||||
im_h, im_w = img.shape[:2]
|
||||
divide_h = math.ceil(im_h / 32) * 32
|
||||
divide_w = math.ceil(im_w / 32) * 32
|
||||
img = img[:, :, 0]
|
||||
img = np.pad(
|
||||
img, ((0, divide_h - im_h), (0, divide_w - im_w)), constant_values=(1, 1)
|
||||
)
|
||||
img_expanded = img[:, :, np.newaxis].transpose(2, 0, 1)
|
||||
data["image"] = img_expanded
|
||||
return data
|
||||
29
ppocr/data/imaug/vqa/__init__.py
Normal file
29
ppocr/data/imaug/vqa/__init__.py
Normal file
@@ -0,0 +1,29 @@
|
||||
# copyright (c) 2021 PaddlePaddle Authors. All Rights Reserve.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from .token import (
|
||||
VQATokenPad,
|
||||
VQASerTokenChunk,
|
||||
VQAReTokenChunk,
|
||||
VQAReTokenRelation,
|
||||
TensorizeEntitiesRelations,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"VQATokenPad",
|
||||
"VQASerTokenChunk",
|
||||
"VQAReTokenChunk",
|
||||
"VQAReTokenRelation",
|
||||
"TensorizeEntitiesRelations",
|
||||
]
|
||||
34
ppocr/data/imaug/vqa/augment.py
Normal file
34
ppocr/data/imaug/vqa/augment.py
Normal file
@@ -0,0 +1,34 @@
|
||||
# copyright (c) 2022 PaddlePaddle Authors. All Rights Reserve.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import os
|
||||
import sys
|
||||
import numpy as np
|
||||
import random
|
||||
from copy import deepcopy
|
||||
|
||||
|
||||
def order_by_tbyx(ocr_info):
|
||||
res = sorted(ocr_info, key=lambda r: (r["bbox"][1], r["bbox"][0]))
|
||||
for i in range(len(res) - 1):
|
||||
for j in range(i, 0, -1):
|
||||
if abs(res[j + 1]["bbox"][1] - res[j]["bbox"][1]) < 20 and (
|
||||
res[j + 1]["bbox"][0] < res[j]["bbox"][0]
|
||||
):
|
||||
tmp = deepcopy(res[j])
|
||||
res[j] = deepcopy(res[j + 1])
|
||||
res[j + 1] = deepcopy(tmp)
|
||||
else:
|
||||
break
|
||||
return res
|
||||
18
ppocr/data/imaug/vqa/token/__init__.py
Normal file
18
ppocr/data/imaug/vqa/token/__init__.py
Normal file
@@ -0,0 +1,18 @@
|
||||
# copyright (c) 2021 PaddlePaddle Authors. All Rights Reserve.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from .vqa_token_chunk import VQASerTokenChunk, VQAReTokenChunk
|
||||
from .vqa_token_pad import VQATokenPad
|
||||
from .vqa_token_relation import VQAReTokenRelation
|
||||
from .vqa_re_convert import TensorizeEntitiesRelations
|
||||
49
ppocr/data/imaug/vqa/token/vqa_re_convert.py
Normal file
49
ppocr/data/imaug/vqa/token/vqa_re_convert.py
Normal file
@@ -0,0 +1,49 @@
|
||||
# copyright (c) 2022 PaddlePaddle Authors. All Rights Reserve.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
class TensorizeEntitiesRelations(object):
|
||||
def __init__(self, max_seq_len=512, infer_mode=False, **kwargs):
|
||||
self.max_seq_len = max_seq_len
|
||||
self.infer_mode = infer_mode
|
||||
|
||||
def __call__(self, data):
|
||||
entities = data["entities"]
|
||||
relations = data["relations"]
|
||||
|
||||
entities_new = np.full(
|
||||
shape=[self.max_seq_len + 1, 3], fill_value=-1, dtype="int64"
|
||||
)
|
||||
entities_new[0, 0] = len(entities["start"])
|
||||
entities_new[0, 1] = len(entities["end"])
|
||||
entities_new[0, 2] = len(entities["label"])
|
||||
entities_new[1 : len(entities["start"]) + 1, 0] = np.array(entities["start"])
|
||||
entities_new[1 : len(entities["end"]) + 1, 1] = np.array(entities["end"])
|
||||
entities_new[1 : len(entities["label"]) + 1, 2] = np.array(entities["label"])
|
||||
|
||||
relations_new = np.full(
|
||||
shape=[self.max_seq_len * self.max_seq_len + 1, 2],
|
||||
fill_value=-1,
|
||||
dtype="int64",
|
||||
)
|
||||
relations_new[0, 0] = len(relations["head"])
|
||||
relations_new[0, 1] = len(relations["tail"])
|
||||
relations_new[1 : len(relations["head"]) + 1, 0] = np.array(relations["head"])
|
||||
relations_new[1 : len(relations["tail"]) + 1, 1] = np.array(relations["tail"])
|
||||
|
||||
data["entities"] = entities_new
|
||||
data["relations"] = relations_new
|
||||
return data
|
||||
134
ppocr/data/imaug/vqa/token/vqa_token_chunk.py
Normal file
134
ppocr/data/imaug/vqa/token/vqa_token_chunk.py
Normal file
@@ -0,0 +1,134 @@
|
||||
# copyright (c) 2021 PaddlePaddle Authors. All Rights Reserve.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from collections import defaultdict
|
||||
|
||||
|
||||
class VQASerTokenChunk(object):
|
||||
def __init__(self, max_seq_len=512, infer_mode=False, **kwargs):
|
||||
self.max_seq_len = max_seq_len
|
||||
self.infer_mode = infer_mode
|
||||
|
||||
def __call__(self, data):
|
||||
encoded_inputs_all = []
|
||||
seq_len = len(data["input_ids"])
|
||||
for index in range(0, seq_len, self.max_seq_len):
|
||||
chunk_beg = index
|
||||
chunk_end = min(index + self.max_seq_len, seq_len)
|
||||
encoded_inputs_example = {}
|
||||
for key in data:
|
||||
if key in [
|
||||
"label",
|
||||
"input_ids",
|
||||
"labels",
|
||||
"token_type_ids",
|
||||
"bbox",
|
||||
"attention_mask",
|
||||
]:
|
||||
if self.infer_mode and key == "labels":
|
||||
encoded_inputs_example[key] = data[key]
|
||||
else:
|
||||
encoded_inputs_example[key] = data[key][chunk_beg:chunk_end]
|
||||
else:
|
||||
encoded_inputs_example[key] = data[key]
|
||||
|
||||
encoded_inputs_all.append(encoded_inputs_example)
|
||||
if len(encoded_inputs_all) == 0:
|
||||
return None
|
||||
return encoded_inputs_all[0]
|
||||
|
||||
|
||||
class VQAReTokenChunk(object):
|
||||
def __init__(
|
||||
self, max_seq_len=512, entities_labels=None, infer_mode=False, **kwargs
|
||||
):
|
||||
self.max_seq_len = max_seq_len
|
||||
self.entities_labels = (
|
||||
{"HEADER": 0, "QUESTION": 1, "ANSWER": 2}
|
||||
if entities_labels is None
|
||||
else entities_labels
|
||||
)
|
||||
self.infer_mode = infer_mode
|
||||
|
||||
def __call__(self, data):
|
||||
# prepare data
|
||||
entities = data.pop("entities")
|
||||
relations = data.pop("relations")
|
||||
encoded_inputs_all = []
|
||||
for index in range(0, len(data["input_ids"]), self.max_seq_len):
|
||||
item = {}
|
||||
for key in data:
|
||||
if key in [
|
||||
"label",
|
||||
"input_ids",
|
||||
"labels",
|
||||
"token_type_ids",
|
||||
"bbox",
|
||||
"attention_mask",
|
||||
]:
|
||||
if self.infer_mode and key == "labels":
|
||||
item[key] = data[key]
|
||||
else:
|
||||
item[key] = data[key][index : index + self.max_seq_len]
|
||||
else:
|
||||
item[key] = data[key]
|
||||
# select entity in current chunk
|
||||
entities_in_this_span = []
|
||||
global_to_local_map = {} #
|
||||
for entity_id, entity in enumerate(entities):
|
||||
if (
|
||||
index <= entity["start"] < index + self.max_seq_len
|
||||
and index <= entity["end"] < index + self.max_seq_len
|
||||
):
|
||||
entity["start"] = entity["start"] - index
|
||||
entity["end"] = entity["end"] - index
|
||||
global_to_local_map[entity_id] = len(entities_in_this_span)
|
||||
entities_in_this_span.append(entity)
|
||||
|
||||
# select relations in current chunk
|
||||
relations_in_this_span = []
|
||||
for relation in relations:
|
||||
if (
|
||||
index <= relation["start_index"] < index + self.max_seq_len
|
||||
and index <= relation["end_index"] < index + self.max_seq_len
|
||||
):
|
||||
relations_in_this_span.append(
|
||||
{
|
||||
"head": global_to_local_map[relation["head"]],
|
||||
"tail": global_to_local_map[relation["tail"]],
|
||||
"start_index": relation["start_index"] - index,
|
||||
"end_index": relation["end_index"] - index,
|
||||
}
|
||||
)
|
||||
item.update(
|
||||
{
|
||||
"entities": self.reformat(entities_in_this_span),
|
||||
"relations": self.reformat(relations_in_this_span),
|
||||
}
|
||||
)
|
||||
if len(item["entities"]) > 0:
|
||||
item["entities"]["label"] = [
|
||||
self.entities_labels[x] for x in item["entities"]["label"]
|
||||
]
|
||||
encoded_inputs_all.append(item)
|
||||
if len(encoded_inputs_all) == 0:
|
||||
return None
|
||||
return encoded_inputs_all[0]
|
||||
|
||||
def reformat(self, data):
|
||||
new_data = defaultdict(list)
|
||||
for item in data:
|
||||
for k, v in item.items():
|
||||
new_data[k].append(v)
|
||||
return new_data
|
||||
117
ppocr/data/imaug/vqa/token/vqa_token_pad.py
Normal file
117
ppocr/data/imaug/vqa/token/vqa_token_pad.py
Normal file
@@ -0,0 +1,117 @@
|
||||
# copyright (c) 2021 PaddlePaddle Authors. All Rights Reserve.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
import paddle
|
||||
import numpy as np
|
||||
|
||||
|
||||
class VQATokenPad(object):
|
||||
def __init__(
|
||||
self,
|
||||
max_seq_len=512,
|
||||
pad_to_max_seq_len=True,
|
||||
return_attention_mask=True,
|
||||
return_token_type_ids=True,
|
||||
truncation_strategy="longest_first",
|
||||
return_overflowing_tokens=False,
|
||||
return_special_tokens_mask=False,
|
||||
infer_mode=False,
|
||||
**kwargs,
|
||||
):
|
||||
self.max_seq_len = max_seq_len
|
||||
self.pad_to_max_seq_len = max_seq_len
|
||||
self.return_attention_mask = return_attention_mask
|
||||
self.return_token_type_ids = return_token_type_ids
|
||||
self.truncation_strategy = truncation_strategy
|
||||
self.return_overflowing_tokens = return_overflowing_tokens
|
||||
self.return_special_tokens_mask = return_special_tokens_mask
|
||||
self.pad_token_label_id = paddle.nn.CrossEntropyLoss().ignore_index
|
||||
self.infer_mode = infer_mode
|
||||
|
||||
def __call__(self, data):
|
||||
needs_to_be_padded = (
|
||||
self.pad_to_max_seq_len and len(data["input_ids"]) < self.max_seq_len
|
||||
)
|
||||
|
||||
if needs_to_be_padded:
|
||||
if "tokenizer_params" in data:
|
||||
tokenizer_params = data.pop("tokenizer_params")
|
||||
else:
|
||||
tokenizer_params = dict(
|
||||
padding_side="right", pad_token_type_id=0, pad_token_id=1
|
||||
)
|
||||
|
||||
difference = self.max_seq_len - len(data["input_ids"])
|
||||
if tokenizer_params["padding_side"] == "right":
|
||||
if self.return_attention_mask:
|
||||
data["attention_mask"] = [1] * len(data["input_ids"]) + [
|
||||
0
|
||||
] * difference
|
||||
if self.return_token_type_ids:
|
||||
data["token_type_ids"] = (
|
||||
data["token_type_ids"]
|
||||
+ [tokenizer_params["pad_token_type_id"]] * difference
|
||||
)
|
||||
if self.return_special_tokens_mask:
|
||||
data["special_tokens_mask"] = (
|
||||
data["special_tokens_mask"] + [1] * difference
|
||||
)
|
||||
data["input_ids"] = (
|
||||
data["input_ids"] + [tokenizer_params["pad_token_id"]] * difference
|
||||
)
|
||||
if not self.infer_mode:
|
||||
data["labels"] = (
|
||||
data["labels"] + [self.pad_token_label_id] * difference
|
||||
)
|
||||
data["bbox"] = data["bbox"] + [[0, 0, 0, 0]] * difference
|
||||
elif tokenizer_params["padding_side"] == "left":
|
||||
if self.return_attention_mask:
|
||||
data["attention_mask"] = [0] * difference + [1] * len(
|
||||
data["input_ids"]
|
||||
)
|
||||
if self.return_token_type_ids:
|
||||
data["token_type_ids"] = [
|
||||
tokenizer_params["pad_token_type_id"]
|
||||
] * difference + data["token_type_ids"]
|
||||
if self.return_special_tokens_mask:
|
||||
data["special_tokens_mask"] = [1] * difference + data[
|
||||
"special_tokens_mask"
|
||||
]
|
||||
data["input_ids"] = [
|
||||
tokenizer_params["pad_token_id"]
|
||||
] * difference + data["input_ids"]
|
||||
if not self.infer_mode:
|
||||
data["labels"] = [self.pad_token_label_id] * difference + data[
|
||||
"labels"
|
||||
]
|
||||
data["bbox"] = [[0, 0, 0, 0]] * difference + data["bbox"]
|
||||
else:
|
||||
if self.return_attention_mask:
|
||||
data["attention_mask"] = [1] * len(data["input_ids"])
|
||||
|
||||
for key in data:
|
||||
if key in [
|
||||
"input_ids",
|
||||
"labels",
|
||||
"token_type_ids",
|
||||
"bbox",
|
||||
"attention_mask",
|
||||
]:
|
||||
if self.infer_mode:
|
||||
if key != "labels":
|
||||
length = min(len(data[key]), self.max_seq_len)
|
||||
data[key] = data[key][:length]
|
||||
else:
|
||||
continue
|
||||
data[key] = np.array(data[key], dtype="int64")
|
||||
return data
|
||||
76
ppocr/data/imaug/vqa/token/vqa_token_relation.py
Normal file
76
ppocr/data/imaug/vqa/token/vqa_token_relation.py
Normal file
@@ -0,0 +1,76 @@
|
||||
# copyright (c) 2021 PaddlePaddle Authors. All Rights Reserve.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
|
||||
class VQAReTokenRelation(object):
|
||||
def __init__(self, **kwargs):
|
||||
pass
|
||||
|
||||
def __call__(self, data):
|
||||
"""
|
||||
build relations
|
||||
"""
|
||||
entities = data["entities"]
|
||||
relations = data["relations"]
|
||||
id2label = data.pop("id2label")
|
||||
empty_entity = data.pop("empty_entity")
|
||||
entity_id_to_index_map = data.pop("entity_id_to_index_map")
|
||||
|
||||
relations = list(set(relations))
|
||||
relations = [
|
||||
rel
|
||||
for rel in relations
|
||||
if rel[0] not in empty_entity and rel[1] not in empty_entity
|
||||
]
|
||||
kv_relations = []
|
||||
for rel in relations:
|
||||
pair = [id2label[rel[0]], id2label[rel[1]]]
|
||||
if pair == ["question", "answer"]:
|
||||
kv_relations.append(
|
||||
{
|
||||
"head": entity_id_to_index_map[rel[0]],
|
||||
"tail": entity_id_to_index_map[rel[1]],
|
||||
}
|
||||
)
|
||||
elif pair == ["answer", "question"]:
|
||||
kv_relations.append(
|
||||
{
|
||||
"head": entity_id_to_index_map[rel[1]],
|
||||
"tail": entity_id_to_index_map[rel[0]],
|
||||
}
|
||||
)
|
||||
else:
|
||||
continue
|
||||
relations = sorted(
|
||||
[
|
||||
{
|
||||
"head": rel["head"],
|
||||
"tail": rel["tail"],
|
||||
"start_index": self.get_relation_span(rel, entities)[0],
|
||||
"end_index": self.get_relation_span(rel, entities)[1],
|
||||
}
|
||||
for rel in kv_relations
|
||||
],
|
||||
key=lambda x: x["head"],
|
||||
)
|
||||
|
||||
data["relations"] = relations
|
||||
return data
|
||||
|
||||
def get_relation_span(self, rel, entities):
|
||||
bound = []
|
||||
for entity_index in [rel["head"], rel["tail"]]:
|
||||
bound.append(entities[entity_index]["start"])
|
||||
bound.append(entities[entity_index]["end"])
|
||||
return min(bound), max(bound)
|
||||
Reference in New Issue
Block a user