This commit is contained in:
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
|
||||
Reference in New Issue
Block a user