# -*- coding:utf-8 -*- 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 def __call__(self, data: dict) -> dict: 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