|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
""" |
|
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 |
|
|
|
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] |
|
|
|
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 |
|
|