添加DBnet代码

data_loader/modules/augment.py

+# -*- 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(im.dtype)
+        return data
+
+
+class RandomScale:
+    def __init__(self, scales, random_rate):
+        """
+        :param scales: 尺度
+        :param ramdon_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 ramdon_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 ramdon_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 VerticallFlip:
+    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

data_loader/modules/iaa_augment.py

+# -*- 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

data_loader/modules/make_border_map.py

+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

data_loader/modules/make_shrink_map.py

+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)
+    shrinked = padding.Execute(-distance)
+    if shrinked == []:
+        shrinked = np.array(shrinked)
+    else:
+        shrinked = np.array(shrinked[0]).reshape(-1, 2)
+    return shrinked
+
+
+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:
+                shrinked = self.shrink_func(polygon, self.shrink_ratio)
+                if shrinked.size == 0:
+                    cv2.fillPoly(mask, polygon.astype(np.int32)[np.newaxis, :, :], 0)
+                    ignore_tags[i] = True
+                    continue
+                cv2.fillPoly(gt, [shrinked.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)

data_loader/modules/random_crop_data.py

+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

environment.yml

+name: dbnet
+channels:
+  - pytorch
+  - conda-forge
+  - defaults
+dependencies:
+  - pytorch==1.4
+  - torchvision==0.5
+  - anyconfig==0.9.10
+  - future==0.18.2
+  - imgaug==0.4.0
+  - matplotlib==3.1.2
+  - numpy==1.17.4
+  - opencv
+  - pyclipper
+  - PyYAML==5.2
+  - scikit-image==0.16.2
+  - Shapely==1.6.4
+  - tensorboard=2
+  - tqdm==4.40.1
+  - ipython
+  - pip
+  - pytorch
+  - torchvision
+  - pip:
+    - polygon3

eval.sh

+CUDA_VISIBLE_DEVICES=0 python3 tools/eval.py --model_path ''
\ No newline at end of file

generate_lists.sh

+#Only use if your file names of the images and txts are identical
+rm ./datasets/train_img.txt
+rm ./datasets/train_gt.txt
+rm ./datasets/test_img.txt
+rm ./datasets/test_gt.txt
+rm ./datasets/train.txt
+rm ./datasets/test.txt
+ls ./datasets/train/img/*.jpg > ./datasets/train_img.txt
+ls ./datasets/train/gt/*.txt > ./datasets/train_gt.txt
+ls ./datasets/test/img/*.jpg > ./datasets/test_img.txt
+ls ./datasets/test/gt/*.txt > ./datasets/test_gt.txt
+paste ./datasets/train_img.txt ./datasets/train_gt.txt > ./datasets/train.txt
+paste ./datasets/test_img.txt ./datasets/test_gt.txt > ./datasets/test.txt
+rm ./datasets/train_img.txt
+rm ./datasets/train_gt.txt
+rm ./datasets/test_img.txt
+rm ./datasets/test_gt.txt

models/__init__.py

+# -*- coding: utf-8 -*-
+# @Time    : 2019/8/23 21:55
+# @Author  : zhoujun
+import copy
+from .model import Model
+from .losses import build_loss
+
+__all__ = ['build_loss', 'build_model']
+support_model = ['Model']
+
+
+def build_model(config):
+    """
+    get architecture model class
+    """
+    copy_config = copy.deepcopy(config)
+    arch_type = copy_config.pop('type')
+    assert arch_type in support_model, f'{arch_type} is not developed yet!, only {support_model} are support now'
+    arch_model = eval(arch_type)(copy_config)
+    return arch_model