chore: import upstream snapshot with attribution

benchmark/PaddleOCR_DBNet/data_loader/modules/__init__.py

@@ -0,0 +1,8 @@
+# -*- coding: utf-8 -*-
+# @Time    : 2019/12/4 10:53
+# @Author  : zhoujun
+from .iaa_augment import IaaAugment
+from .augment import *
+from .random_crop_data import EastRandomCropData, PSERandomCrop
+from .make_border_map import MakeBorderMap
+from .make_shrink_map import MakeShrinkMap

benchmark/PaddleOCR_DBNet/data_loader/modules/augment.py

@@ -0,0 +1,308 @@
+# -*- coding: utf-8 -*-
+# @Time    : 2019/8/23 21:52
+# @Author  : zhoujun
+
+import math
+import numbers
+import random
+
+import cv2
+import numpy as np
+from skimage.util import random_noise
+
+
+class RandomNoise:
+    def __init__(self, random_rate):
+        self.random_rate = random_rate
+
+    def __call__(self, data: dict):
+        """
+        对图片加噪声
+        :param data: {'img':,'text_polys':,'texts':,'ignore_tags':}
+        :return:
+        """
+        if random.random() > self.random_rate:
+            return data
+        data["img"] = (
+            random_noise(data["img"], mode="gaussian", clip=True) * 255
+        ).astype(data["img"].dtype)
+        return data
+
+
+class RandomScale:
+    def __init__(self, scales, random_rate):
+        """
+        :param scales: 尺度
+        :param random_rate: 随机系数
+        :return:
+        """
+        self.random_rate = random_rate
+        self.scales = scales
+
+    def __call__(self, data: dict) -> dict:
+        """
+        从scales中随机选择一个尺度，对图片和文本框进行缩放
+        :param data: {'img':,'text_polys':,'texts':,'ignore_tags':}
+        :return:
+        """
+        if random.random() > self.random_rate:
+            return data
+        im = data["img"]
+        text_polys = data["text_polys"]
+
+        tmp_text_polys = text_polys.copy()
+        rd_scale = float(np.random.choice(self.scales))
+        im = cv2.resize(im, dsize=None, fx=rd_scale, fy=rd_scale)
+        tmp_text_polys *= rd_scale
+
+        data["img"] = im
+        data["text_polys"] = tmp_text_polys
+        return data
+
+
+class RandomRotateImgBox:
+    def __init__(self, degrees, random_rate, same_size=False):
+        """
+        :param degrees: 角度，可以是一个数值或者list
+        :param random_rate: 随机系数
+        :param same_size: 是否保持和原图一样大
+        :return:
+        """
+        if isinstance(degrees, numbers.Number):
+            if degrees < 0:
+                raise ValueError("If degrees is a single number, it must be positive.")
+            degrees = (-degrees, degrees)
+        elif (
+            isinstance(degrees, list)
+            or isinstance(degrees, tuple)
+            or isinstance(degrees, np.ndarray)
+        ):
+            if len(degrees) != 2:
+                raise ValueError("If degrees is a sequence, it must be of len 2.")
+            degrees = degrees
+        else:
+            raise Exception("degrees must in Number or list or tuple or np.ndarray")
+        self.degrees = degrees
+        self.same_size = same_size
+        self.random_rate = random_rate
+
+    def __call__(self, data: dict) -> dict:
+        """
+        从scales中随机选择一个尺度，对图片和文本框进行缩放
+        :param data: {'img':,'text_polys':,'texts':,'ignore_tags':}
+        :return:
+        """
+        if random.random() > self.random_rate:
+            return data
+        im = data["img"]
+        text_polys = data["text_polys"]
+
+        # ---------------------- 旋转图像 ----------------------
+        w = im.shape[1]
+        h = im.shape[0]
+        angle = np.random.uniform(self.degrees[0], self.degrees[1])
+
+        if self.same_size:
+            nw = w
+            nh = h
+        else:
+            # 角度变弧度
+            rangle = np.deg2rad(angle)
+            # 计算旋转之后图像的w, h
+            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, 1)
+        # 计算原图中心点到新图中心点的偏移量
+        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]
+        # 仿射变换
+        rot_img = cv2.warpAffine(
+            im,
+            rot_mat,
+            (int(math.ceil(nw)), int(math.ceil(nh))),
+            flags=cv2.INTER_LANCZOS4,
+        )
+
+        # ---------------------- 矫正bbox坐标 ----------------------
+        # rot_mat是最终的旋转矩阵
+        # 获取原始bbox的四个中点，然后将这四个点转换到旋转后的坐标系下
+        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])
+        data["img"] = rot_img
+        data["text_polys"] = np.array(rot_text_polys)
+        return data
+
+
+class RandomResize:
+    def __init__(self, size, random_rate, keep_ratio=False):
+        """
+        :param input_size: resize尺寸,数字或者list的形式，如果为list形式，就是[w,h]
+        :param random_rate: 随机系数
+        :param keep_ratio: 是否保持长宽比
+        :return:
+        """
+        if isinstance(size, numbers.Number):
+            if size < 0:
+                raise ValueError(
+                    "If input_size is a single number, it must be positive."
+                )
+            size = (size, size)
+        elif (
+            isinstance(size, list)
+            or isinstance(size, tuple)
+            or isinstance(size, np.ndarray)
+        ):
+            if len(size) != 2:
+                raise ValueError("If input_size is a sequence, it must be of len 2.")
+            size = (size[0], size[1])
+        else:
+            raise Exception("input_size must in Number or list or tuple or np.ndarray")
+        self.size = size
+        self.keep_ratio = keep_ratio
+        self.random_rate = random_rate
+
+    def __call__(self, data: dict) -> dict:
+        """
+        从scales中随机选择一个尺度，对图片和文本框进行缩放
+        :param data: {'img':,'text_polys':,'texts':,'ignore_tags':}
+        :return:
+        """
+        if random.random() > self.random_rate:
+            return data
+        im = data["img"]
+        text_polys = data["text_polys"]
+
+        if self.keep_ratio:
+            # 将图片短边pad到和长边一样
+            h, w, c = im.shape
+            max_h = max(h, self.size[0])
+            max_w = max(w, self.size[1])
+            im_padded = np.zeros((max_h, max_w, c), dtype=np.uint8)
+            im_padded[:h, :w] = im.copy()
+            im = im_padded
+        text_polys = text_polys.astype(np.float32)
+        h, w, _ = im.shape
+        im = cv2.resize(im, self.size)
+        w_scale = self.size[0] / float(w)
+        h_scale = self.size[1] / float(h)
+        text_polys[:, :, 0] *= w_scale
+        text_polys[:, :, 1] *= h_scale
+
+        data["img"] = im
+        data["text_polys"] = text_polys
+        return data
+
+
+def resize_image(img, short_size):
+    height, width, _ = img.shape
+    if height < width:
+        new_height = short_size
+        new_width = new_height / height * width
+    else:
+        new_width = short_size
+        new_height = new_width / width * height
+    new_height = int(round(new_height / 32) * 32)
+    new_width = int(round(new_width / 32) * 32)
+    resized_img = cv2.resize(img, (new_width, new_height))
+    return resized_img, (new_width / width, new_height / height)
+
+
+class ResizeShortSize:
+    def __init__(self, short_size, resize_text_polys=True):
+        """
+        :param size: resize尺寸,数字或者list的形式，如果为list形式，就是[w,h]
+        :return:
+        """
+        self.short_size = short_size
+        self.resize_text_polys = resize_text_polys
+
+    def __call__(self, data: dict) -> dict:
+        """
+        对图片和文本框进行缩放
+        :param data: {'img':,'text_polys':,'texts':,'ignore_tags':}
+        :return:
+        """
+        im = data["img"]
+        text_polys = data["text_polys"]
+
+        h, w, _ = im.shape
+        short_edge = min(h, w)
+        if short_edge < self.short_size:
+            # 保证短边 >= short_size
+            scale = self.short_size / short_edge
+            im = cv2.resize(im, dsize=None, fx=scale, fy=scale)
+            scale = (scale, scale)
+            # im, scale = resize_image(im, self.short_size)
+            if self.resize_text_polys:
+                # text_polys *= scale
+                text_polys[:, 0] *= scale[0]
+                text_polys[:, 1] *= scale[1]
+
+        data["img"] = im
+        data["text_polys"] = text_polys
+        return data
+
+
+class HorizontalFlip:
+    def __init__(self, random_rate):
+        """
+
+        :param random_rate: 随机系数
+        """
+        self.random_rate = random_rate
+
+    def __call__(self, data: dict) -> dict:
+        """
+        从scales中随机选择一个尺度，对图片和文本框进行缩放
+        :param data: {'img':,'text_polys':,'texts':,'ignore_tags':}
+        :return:
+        """
+        if random.random() > self.random_rate:
+            return data
+        im = data["img"]
+        text_polys = data["text_polys"]
+
+        flip_text_polys = text_polys.copy()
+        flip_im = cv2.flip(im, 1)
+        h, w, _ = flip_im.shape
+        flip_text_polys[:, :, 0] = w - flip_text_polys[:, :, 0]
+
+        data["img"] = flip_im
+        data["text_polys"] = flip_text_polys
+        return data
+
+
+class VerticalFlip:
+    def __init__(self, random_rate):
+        """
+
+        :param random_rate: 随机系数
+        """
+        self.random_rate = random_rate
+
+    def __call__(self, data: dict) -> dict:
+        """
+        从scales中随机选择一个尺度，对图片和文本框进行缩放
+        :param data: {'img':,'text_polys':,'texts':,'ignore_tags':}
+        :return:
+        """
+        if random.random() > self.random_rate:
+            return data
+        im = data["img"]
+        text_polys = data["text_polys"]
+
+        flip_text_polys = text_polys.copy()
+        flip_im = cv2.flip(im, 0)
+        h, w, _ = flip_im.shape
+        flip_text_polys[:, :, 1] = h - flip_text_polys[:, :, 1]
+        data["img"] = flip_im
+        data["text_polys"] = flip_text_polys
+        return data

benchmark/PaddleOCR_DBNet/data_loader/modules/iaa_augment.py

@@ -0,0 +1,68 @@
+# -*- coding: utf-8 -*-
+# @Time    : 2019/12/4 18:06
+# @Author  : zhoujun
+import numpy as np
+import imgaug
+import imgaug.augmenters as iaa
+
+
+class AugmenterBuilder(object):
+    def __init__(self):
+        pass
+
+    def build(self, args, root=True):
+        if args is None or len(args) == 0:
+            return None
+        elif isinstance(args, list):
+            if root:
+                sequence = [self.build(value, root=False) for value in args]
+                return iaa.Sequential(sequence)
+            else:
+                return getattr(iaa, args[0])(
+                    *[self.to_tuple_if_list(a) for a in args[1:]]
+                )
+        elif isinstance(args, dict):
+            cls = getattr(iaa, args["type"])
+            return cls(**{k: self.to_tuple_if_list(v) for k, v in args["args"].items()})
+        else:
+            raise RuntimeError("unknown augmenter arg: " + str(args))
+
+    def to_tuple_if_list(self, obj):
+        if isinstance(obj, list):
+            return tuple(obj)
+        return obj
+
+
+class IaaAugment:
+    def __init__(self, augmenter_args):
+        self.augmenter_args = augmenter_args
+        self.augmenter = AugmenterBuilder().build(self.augmenter_args)
+
+    def __call__(self, data):
+        image = data["img"]
+        shape = image.shape
+
+        if self.augmenter:
+            aug = self.augmenter.to_deterministic()
+            data["img"] = aug.augment_image(image)
+            data = self.may_augment_annotation(aug, data, shape)
+        return data
+
+    def may_augment_annotation(self, aug, data, shape):
+        if aug is None:
+            return data
+
+        line_polys = []
+        for poly in data["text_polys"]:
+            new_poly = self.may_augment_poly(aug, shape, poly)
+            line_polys.append(new_poly)
+        data["text_polys"] = np.array(line_polys)
+        return data
+
+    def may_augment_poly(self, aug, img_shape, poly):
+        keypoints = [imgaug.Keypoint(p[0], p[1]) for p in poly]
+        keypoints = aug.augment_keypoints(
+            [imgaug.KeypointsOnImage(keypoints, shape=img_shape)]
+        )[0].keypoints
+        poly = [(p.x, p.y) for p in keypoints]
+        return poly

benchmark/PaddleOCR_DBNet/data_loader/modules/make_border_map.py

@@ -0,0 +1,159 @@
+import cv2
+import numpy as np
+
+np.seterr(divide="ignore", invalid="ignore")
+import pyclipper
+from shapely.geometry import Polygon
+
+
+class MakeBorderMap:
+    def __init__(self, shrink_ratio=0.4, thresh_min=0.3, thresh_max=0.7):
+        self.shrink_ratio = shrink_ratio
+        self.thresh_min = thresh_min
+        self.thresh_max = thresh_max
+
+    def __call__(self, data: dict) -> dict:
+        """
+        从scales中随机选择一个尺度，对图片和文本框进行缩放
+        :param data: {'img':,'text_polys':,'texts':,'ignore_tags':}
+        :return:
+        """
+        im = data["img"]
+        text_polys = data["text_polys"]
+        ignore_tags = data["ignore_tags"]
+
+        canvas = np.zeros(im.shape[:2], dtype=np.float32)
+        mask = np.zeros(im.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):
+        ex_point_1 = (
+            int(
+                round(point_1[0] + (point_1[0] - point_2[0]) * (1 + self.shrink_ratio))
+            ),
+            int(
+                round(point_1[1] + (point_1[1] - point_2[1]) * (1 + self.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 + self.shrink_ratio))
+            ),
+            int(
+                round(point_2[1] + (point_2[1] - point_1[1]) * (1 + self.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

benchmark/PaddleOCR_DBNet/data_loader/modules/make_shrink_map.py

@@ -0,0 +1,129 @@
+import numpy as np
+import cv2
+
+
+def shrink_polygon_py(polygon, shrink_ratio):
+    """
+    对框进行缩放，返回去的比例为1/shrink_ratio 即可
+    """
+    cx = polygon[:, 0].mean()
+    cy = polygon[:, 1].mean()
+    polygon[:, 0] = cx + (polygon[:, 0] - cx) * shrink_ratio
+    polygon[:, 1] = cy + (polygon[:, 1] - cy) * shrink_ratio
+    return polygon
+
+
+def shrink_polygon_pyclipper(polygon, shrink_ratio):
+    from shapely.geometry import Polygon
+    import pyclipper
+
+    polygon_shape = Polygon(polygon)
+    distance = (
+        polygon_shape.area * (1 - np.power(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)
+    shrunk = padding.Execute(-distance)
+    if shrunk == []:
+        shrunk = np.array(shrunk)
+    else:
+        shrunk = np.array(shrunk[0]).reshape(-1, 2)
+    return shrunk
+
+
+class MakeShrinkMap:
+    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, shrink_type="pyclipper"):
+        shrink_func_dict = {
+            "py": shrink_polygon_py,
+            "pyclipper": shrink_polygon_pyclipper,
+        }
+        self.shrink_func = shrink_func_dict[shrink_type]
+        self.min_text_size = min_text_size
+        self.shrink_ratio = shrink_ratio
+
+    def __call__(self, data: dict) -> dict:
+        """
+        从scales中随机选择一个尺度，对图片和文本框进行缩放
+        :param data: {'img':,'text_polys':,'texts':,'ignore_tags':}
+        :return:
+        """
+        image = data["img"]
+        text_polys = data["text_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:
+                shrunk = self.shrink_func(polygon, self.shrink_ratio)
+                if shrunk.size == 0:
+                    cv2.fillPoly(mask, polygon.astype(np.int32)[np.newaxis, :, :], 0)
+                    ignore_tags[i] = True
+                    continue
+                cv2.fillPoly(gt, [shrunk.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):
+        return cv2.contourArea(polygon)
+        # edge = 0
+        # for i in range(polygon.shape[0]):
+        #     next_index = (i + 1) % polygon.shape[0]
+        #     edge += (polygon[next_index, 0] - polygon[i, 0]) * (polygon[next_index, 1] - polygon[i, 1])
+        #
+        # return edge / 2.
+
+
+if __name__ == "__main__":
+    from shapely.geometry import Polygon
+    import pyclipper
+
+    polygon = np.array([[0, 0], [100, 10], [100, 100], [10, 90]])
+    a = shrink_polygon_py(polygon, 0.4)
+    print(a)
+    print(shrink_polygon_py(a, 1 / 0.4))
+    b = shrink_polygon_pyclipper(polygon, 0.4)
+    print(b)
+    poly = Polygon(b)
+    distance = poly.area * 1.5 / poly.length
+    offset = pyclipper.PyclipperOffset()
+    offset.AddPath(b, pyclipper.JT_ROUND, pyclipper.ET_CLOSEDPOLYGON)
+    expanded = np.array(offset.Execute(distance))
+    bounding_box = cv2.minAreaRect(expanded)
+    points = cv2.boxPoints(bounding_box)
+    print(points)

benchmark/PaddleOCR_DBNet/data_loader/modules/random_crop_data.py

@@ -0,0 +1,211 @@
+import random
+
+import cv2
+import numpy as np
+
+
+# random crop algorithm similar to https://github.com/argman/EAST
+class EastRandomCropData:
+    def __init__(
+        self,
+        size=(640, 640),
+        max_tries=50,
+        min_crop_side_ratio=0.1,
+        require_original_image=False,
+        keep_ratio=True,
+    ):
+        self.size = size
+        self.max_tries = max_tries
+        self.min_crop_side_ratio = min_crop_side_ratio
+        self.require_original_image = require_original_image
+        self.keep_ratio = keep_ratio
+
+    def __call__(self, data: dict) -> dict:
+        """
+        从scales中随机选择一个尺度，对图片和文本框进行缩放
+        :param data: {'img':,'text_polys':,'texts':,'ignore_tags':}
+        :return:
+        """
+        im = data["img"]
+        text_polys = data["text_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 = self.crop_area(im, all_care_polys)
+        # 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:
+            if len(im.shape) == 3:
+                padimg = np.zeros((self.size[1], self.size[0], im.shape[2]), im.dtype)
+            else:
+                padimg = np.zeros((self.size[1], self.size[0]), im.dtype)
+            padimg[:h, :w] = cv2.resize(
+                im[crop_y : crop_y + crop_h, crop_x : crop_x + crop_w], (w, h)
+            )
+            img = padimg
+        else:
+            img = cv2.resize(
+                im[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 self.is_poly_outside_rect(poly, 0, 0, w, h):
+                text_polys_crop.append(poly)
+                ignore_tags_crop.append(tag)
+                texts_crop.append(text)
+        data["img"] = img
+        data["text_polys"] = np.float32(text_polys_crop)
+        data["ignore_tags"] = ignore_tags_crop
+        data["texts"] = texts_crop
+        return data
+
+    def is_poly_in_rect(self, 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(self, 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(self, 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(self, 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(self, 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(self, im, text_polys):
+        h, w = im.shape[:2]
+        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 = self.split_regions(h_axis)
+        w_regions = self.split_regions(w_axis)
+
+        for i in range(self.max_tries):
+            if len(w_regions) > 1:
+                xmin, xmax = self.region_wise_random_select(w_regions, w)
+            else:
+                xmin, xmax = self.random_select(w_axis, w)
+            if len(h_regions) > 1:
+                ymin, ymax = self.region_wise_random_select(h_regions, h)
+            else:
+                ymin, ymax = self.random_select(h_axis, h)
+
+            if (
+                xmax - xmin < self.min_crop_side_ratio * w
+                or ymax - ymin < self.min_crop_side_ratio * h
+            ):
+                # area too small
+                continue
+            num_poly_in_rect = 0
+            for poly in text_polys:
+                if not self.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 PSERandomCrop:
+    def __init__(self, size):
+        self.size = size
+
+    def __call__(self, data):
+        imgs = data["imgs"]
+
+        h, w = imgs[0].shape[0:2]
+        th, tw = self.size
+        if w == tw and h == th:
+            return imgs
+
+        # label中存在文本实例，并且按照概率进行裁剪，使用threshold_label_map控制
+        if np.max(imgs[2]) > 0 and random.random() > 3 / 8:
+            # 文本实例的左上角点
+            tl = np.min(np.where(imgs[2] > 0), axis=1) - self.size
+            tl[tl < 0] = 0
+            # 文本实例的右下角点
+            br = np.max(np.where(imgs[2] > 0), axis=1) - self.size
+            br[br < 0] = 0
+            # 保证选到右下角点时，有足够的距离进行crop
+            br[0] = min(br[0], h - th)
+            br[1] = min(br[1], w - tw)
+
+            for _ in range(50000):
+                i = random.randint(tl[0], br[0])
+                j = random.randint(tl[1], br[1])
+                # 保证shrink_label_map有文本
+                if imgs[1][i : i + th, j : j + tw].sum() <= 0:
+                    continue
+                else:
+                    break
+        else:
+            i = random.randint(0, h - th)
+            j = random.randint(0, w - tw)
+
+        # return i, j, th, tw
+        for idx in range(len(imgs)):
+            if len(imgs[idx].shape) == 3:
+                imgs[idx] = imgs[idx][i : i + th, j : j + tw, :]
+            else:
+                imgs[idx] = imgs[idx][i : i + th, j : j + tw]
+        data["imgs"] = imgs
+        return data