e904b667c6
Build/Publish Develop Docs / deploy (push) Failing after 1s
PaddleOCR Code Style Check / check-code-style (push) Failing after 1s
PaddleOCR PR Tests GPU / detect-changes (push) Failing after 1s
PaddleOCR PR Tests / detect-changes (push) Failing after 1s
PaddleOCR PR Tests GPU / test-pr-gpu (push) Has been cancelled
PaddleOCR PR Tests / test-pr (push) Has been cancelled
PaddleOCR PR Tests GPU / test-pr-gpu-impl (push) Has been cancelled
PaddleOCR PR Tests / test-pr-python (3.13) (push) Has been cancelled
PaddleOCR PR Tests / test-pr-python (3.8) (push) Has been cancelled
PaddleOCR PR Tests / test-pr-python (3.9) (push) Has been cancelled
578 lines
20 KiB
Python
578 lines
20 KiB
Python
# 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
|
|
|
|
from paddle import get_device
|
|
|
|
from shapely.geometry import Polygon, box as shapely_box
|
|
from shapely import intersection
|
|
|
|
|
|
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 get_min_rotated_rect_side(poly):
|
|
"""
|
|
Compute the minimum side length of the minimum bounding rotated rectangle of a polygon.
|
|
"""
|
|
poly = np.array(poly).astype(np.float32)
|
|
if len(poly) < 3:
|
|
return 0
|
|
rect = cv2.minAreaRect(poly)
|
|
width, height = rect[1]
|
|
return min(width, height)
|
|
|
|
|
|
def get_min_quad_side(quad):
|
|
"""
|
|
Compute the minimum side length of a quadrilateral.
|
|
"""
|
|
if len(quad) != 4:
|
|
return 0
|
|
quad = np.array(quad)
|
|
sides = []
|
|
for i in range(4):
|
|
side = np.linalg.norm(quad[i] - quad[(i + 1) % 4])
|
|
sides.append(side)
|
|
return min(sides) if sides else 0
|
|
|
|
|
|
def clip_poly_to_rect(poly, x, y, w, h):
|
|
"""
|
|
Clip a polygon to a rectangular region and return the clipped quadrilateral.
|
|
|
|
Args:
|
|
poly: Original polygon vertices [[x1, y1], [x2, y2], ...]
|
|
x, y, w, h: Position and size of the clipping rectangle
|
|
|
|
Returns:
|
|
Clipped quadrilateral vertices, or None if the result is invalid.
|
|
"""
|
|
try:
|
|
# Create polygon and clipping rectangle
|
|
poly_shape = Polygon(poly)
|
|
crop_rect = shapely_box(x, y, x + w, y + h)
|
|
|
|
# Compute intersection
|
|
clipped = intersection(poly_shape, crop_rect)
|
|
|
|
# No intersection or empty
|
|
if clipped.is_empty:
|
|
return None
|
|
|
|
# Get intersection coordinates
|
|
if clipped.geom_type == "Polygon":
|
|
coords = list(clipped.exterior.coords[:-1]) # Remove duplicate last point
|
|
elif clipped.geom_type == "MultiPolygon":
|
|
# Multiple polygons, select the largest by area
|
|
largest = max(clipped.geoms, key=lambda p: p.area)
|
|
coords = list(largest.exterior.coords[:-1])
|
|
elif clipped.geom_type == "GeometryCollection":
|
|
# Extract polygons from geometry collection
|
|
polygons = [g for g in clipped.geoms if g.geom_type == "Polygon"]
|
|
if not polygons:
|
|
return None
|
|
largest = max(polygons, key=lambda p: p.area)
|
|
coords = list(largest.exterior.coords[:-1])
|
|
else:
|
|
return None
|
|
|
|
# Less than 3 points, invalid
|
|
if len(coords) <= 3:
|
|
return None
|
|
|
|
# Convert to numpy array
|
|
coords = np.array(coords)
|
|
|
|
# Exactly 4 points, return directly
|
|
if len(coords) == 4:
|
|
return coords
|
|
|
|
# More than 4 points, use Douglas-Peucker to simplify to quadrilateral
|
|
# Output points are a subset of original coords (no out-of-bounds), IoU~0.99
|
|
if len(coords) > 4:
|
|
poly_cv = coords.reshape(-1, 1, 2).astype(np.float32)
|
|
peri = cv2.arcLength(poly_cv, True)
|
|
if peri < 1e-6:
|
|
return None
|
|
lo, hi = 0.0, 0.5
|
|
best = None
|
|
for _ in range(50):
|
|
mid = (lo + hi) / 2
|
|
approx = cv2.approxPolyDP(poly_cv, mid * peri, True)
|
|
if len(approx) <= 4:
|
|
best = approx
|
|
hi = mid
|
|
else:
|
|
lo = mid
|
|
if best is not None and len(best) >= 3:
|
|
return best.reshape(-1, 2)
|
|
return None
|
|
|
|
return coords
|
|
except Exception as e:
|
|
return None
|
|
|
|
|
|
class RandomCrop(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"]
|
|
|
|
# Separate care and ignore text boxes
|
|
care_indices = [i for i, tag in enumerate(ignore_tags) if not tag]
|
|
all_care_polys = [text_polys[i] for i in care_indices]
|
|
|
|
h, w, _ = img.shape
|
|
|
|
# If no valid text boxes, still need to resize and pad the image
|
|
if len(all_care_polys) == 0:
|
|
# Use entire image as crop region, skip crop loop
|
|
crop_x, crop_y, crop_w, crop_h = 0, 0, w, h
|
|
valid_care_data = []
|
|
else:
|
|
# Pre-compute char heights (min side of min bounding rotated rect) for all care boxes
|
|
char_heights = np.array(
|
|
[get_min_rotated_rect_side(poly) for poly in all_care_polys]
|
|
)
|
|
|
|
# Try to find a suitable crop region
|
|
valid_care_data = []
|
|
for attempt in range(self.max_tries):
|
|
# Randomly determine crop region width and height
|
|
crop_w_min = min(int(w * self.min_crop_side_ratio), self.size[0])
|
|
crop_w_max = int(self.size[0] * 3)
|
|
crop_w = (
|
|
w
|
|
if crop_w_min >= crop_w_max
|
|
else min(random.randint(crop_w_min, crop_w_max), w)
|
|
)
|
|
|
|
crop_h_min = min(int(h * self.min_crop_side_ratio), self.size[1])
|
|
crop_h_max = int(self.size[1] * 3)
|
|
crop_h = (
|
|
h
|
|
if crop_h_min >= crop_h_max
|
|
else min(random.randint(crop_h_min, crop_h_max), h)
|
|
)
|
|
|
|
# Randomly determine crop region start position
|
|
crop_x = 0 if crop_w >= w else random.randint(0, w - crop_w)
|
|
crop_y = 0 if crop_h >= h else random.randint(0, h - crop_h)
|
|
|
|
# Check each care text box, clip and validate simultaneously (computed once)
|
|
valid_care_data = []
|
|
for care_idx, (poly, char_height) in enumerate(
|
|
zip(all_care_polys, char_heights)
|
|
):
|
|
# Quick check: completely outside, skip
|
|
if is_poly_outside_rect(poly, crop_x, crop_y, crop_w, crop_h):
|
|
continue
|
|
|
|
# Completely inside, no clipping needed
|
|
if is_poly_in_rect(poly, crop_x, crop_y, crop_w, crop_h):
|
|
valid_care_data.append(
|
|
(care_idx, None)
|
|
) # None means no clipping needed
|
|
continue
|
|
|
|
# Truncated box, clip and validate (executed once)
|
|
clipped_poly = clip_poly_to_rect(
|
|
poly, crop_x, crop_y, crop_w, crop_h
|
|
)
|
|
if clipped_poly is None:
|
|
continue
|
|
|
|
# Validate clipped polygon - area check
|
|
clipped_area = cv2.contourArea(clipped_poly.astype(np.float32))
|
|
if clipped_area < 80:
|
|
continue
|
|
|
|
# Validate - char height check
|
|
clipped_char_height = get_min_rotated_rect_side(clipped_poly)
|
|
if clipped_char_height < char_height * 0.35:
|
|
continue
|
|
|
|
# Validate - min side length check (quadrilaterals only)
|
|
if len(clipped_poly) == 4:
|
|
min_side = get_min_quad_side(clipped_poly)
|
|
if min_side < char_height * 0.35:
|
|
continue
|
|
|
|
# All validations passed, save the clipped polygon
|
|
valid_care_data.append((care_idx, clipped_poly))
|
|
|
|
# At least one valid text box, use this crop region
|
|
if len(valid_care_data) >= 1:
|
|
break
|
|
else:
|
|
# All attempts failed, use original region
|
|
crop_x, crop_y, crop_w, crop_h = 0, 0, w, h
|
|
valid_care_data = [(i, None) for i in range(len(all_care_polys))]
|
|
|
|
# Crop and scale image
|
|
# Only shrink when crop region is larger than target size, otherwise just pad
|
|
need_resize = crop_w > self.size[0] or crop_h > self.size[1]
|
|
|
|
if need_resize:
|
|
# Crop region larger than target, need to shrink
|
|
scale_w = self.size[0] / crop_w
|
|
scale_h = self.size[1] / crop_h
|
|
scale = min(scale_w, scale_h)
|
|
h_resized = int(crop_h * scale)
|
|
w_resized = int(crop_w * scale)
|
|
else:
|
|
# Crop region smaller than or equal to target, no upscaling
|
|
scale = 1.0
|
|
h_resized = crop_h
|
|
w_resized = crop_w
|
|
|
|
if self.keep_ratio:
|
|
# Random padding - compute padding size
|
|
pad_h = self.size[1] - h_resized
|
|
pad_w = self.size[0] - w_resized
|
|
|
|
# Randomly distribute padding to each side
|
|
pad_top = random.randint(0, pad_h) if pad_h > 0 else 0
|
|
pad_left = random.randint(0, pad_w) if pad_w > 0 else 0
|
|
|
|
# Resize cropped image (only when shrinking is needed)
|
|
cropped_img = img[crop_y : crop_y + crop_h, crop_x : crop_x + crop_w]
|
|
if need_resize:
|
|
resized_img = cv2.resize(cropped_img, (w_resized, h_resized))
|
|
else:
|
|
resized_img = cropped_img
|
|
|
|
# Create padded image
|
|
padimg = np.zeros((self.size[1], self.size[0], img.shape[2]), img.dtype)
|
|
padimg[pad_top : pad_top + h_resized, pad_left : pad_left + w_resized] = (
|
|
resized_img
|
|
)
|
|
img = padimg
|
|
else:
|
|
img = cv2.resize(
|
|
img[crop_y : crop_y + crop_h, crop_x : crop_x + crop_w],
|
|
tuple(self.size),
|
|
)
|
|
pad_left = 0
|
|
pad_top = 0
|
|
|
|
# Build fast lookup set of valid care indices
|
|
valid_care_indices_set = {care_idx for care_idx, _ in valid_care_data}
|
|
|
|
# Build mapping from care_idx to clipped polygon
|
|
care_idx_to_clipped = {
|
|
care_idx: clipped for care_idx, clipped in valid_care_data
|
|
}
|
|
|
|
# Build output text box list
|
|
text_polys_crop = []
|
|
ignore_tags_crop = []
|
|
texts_crop = []
|
|
|
|
for all_idx, (poly, text, tag) in enumerate(
|
|
zip(text_polys, texts, ignore_tags)
|
|
):
|
|
if tag:
|
|
# Ignore text box, simple processing
|
|
if not is_poly_outside_rect(poly, crop_x, crop_y, crop_w, crop_h):
|
|
adjusted_poly = (poly - (crop_x, crop_y)) * scale + (
|
|
pad_left,
|
|
pad_top,
|
|
)
|
|
adjusted_poly[:, 0] = np.clip(adjusted_poly[:, 0], 0, self.size[0])
|
|
adjusted_poly[:, 1] = np.clip(adjusted_poly[:, 1], 0, self.size[1])
|
|
text_polys_crop.append(adjusted_poly.tolist())
|
|
ignore_tags_crop.append(tag)
|
|
texts_crop.append(text)
|
|
else:
|
|
# Care text box, find corresponding care_idx
|
|
try:
|
|
care_idx = care_indices.index(all_idx)
|
|
except ValueError:
|
|
continue
|
|
|
|
# Check if this is a valid text box
|
|
if care_idx not in valid_care_indices_set:
|
|
continue
|
|
|
|
# Get clipped polygon (if any)
|
|
clipped_poly = care_idx_to_clipped[care_idx]
|
|
|
|
if clipped_poly is None:
|
|
# Completely inside, use original polygon
|
|
adjusted_poly = (poly - (crop_x, crop_y)) * scale + (
|
|
pad_left,
|
|
pad_top,
|
|
)
|
|
else:
|
|
# Use clipped polygon
|
|
adjusted_poly = (clipped_poly - (crop_x, crop_y)) * scale + (
|
|
pad_left,
|
|
pad_top,
|
|
)
|
|
|
|
text_polys_crop.append(adjusted_poly.tolist())
|
|
ignore_tags_crop.append(tag)
|
|
texts_crop.append(text)
|
|
data["image"] = img
|
|
# Pad polygons to uniform point count to avoid inhomogeneous array error
|
|
if text_polys_crop:
|
|
max_points = max(len(p) for p in text_polys_crop)
|
|
for i, poly in enumerate(text_polys_crop):
|
|
if len(poly) < max_points:
|
|
text_polys_crop[i] = poly + [poly[-1]] * (max_points - len(poly))
|
|
data["polys"] = np.array(text_polys_crop, dtype=np.float32)
|
|
data["ignore_tags"] = ignore_tags_crop
|
|
data["texts"] = texts_crop
|
|
return data
|
|
|
|
|
|
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)
|
|
if "iluvatar_gpu" in get_device():
|
|
data["polys"] = np.array(text_polys_crop).astype(np.float32)
|
|
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
|