Spaces:
Building
Building
| # 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/make_border_map.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 | |
| 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 | |
| if 'total_epoch' in kwargs and 'epoch' in kwargs and kwargs[ | |
| 'epoch'] != "None": | |
| self.shrink_ratio = self.shrink_ratio + 0.2 * kwargs[ | |
| 'epoch'] / float(kwargs['total_epoch']) | |
| 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 | |