Spaces:
Runtime error
Runtime error
| from __future__ import absolute_import, division, print_function, unicode_literals | |
| import cv2 | |
| import numpy as np | |
| np.seterr(divide="ignore", invalid="ignore") | |
| import warnings | |
| import pyclipper | |
| from shapely.geometry import Polygon | |
| 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): | |
| 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 | |