Spaces:
Runtime error
Runtime error
from __future__ import absolute_import, division, print_function, unicode_literals | |
import cv2 | |
import numpy as np | |
import pyclipper | |
from shapely.geometry import Polygon | |
__all__ = ["MakeShrinkMap"] | |
class MakeShrinkMap(object): | |
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, **kwargs): | |
self.min_text_size = min_text_size | |
self.shrink_ratio = shrink_ratio | |
def __call__(self, data): | |
image = data["image"] | |
text_polys = data["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: | |
polygon_shape = Polygon(polygon) | |
subject = [tuple(l) for l in polygon] | |
padding = pyclipper.PyclipperOffset() | |
padding.AddPath(subject, pyclipper.JT_ROUND, pyclipper.ET_CLOSEDPOLYGON) | |
shrinked = [] | |
# Increase the shrink ratio every time we get multiple polygon returned back | |
possible_ratios = np.arange(self.shrink_ratio, 1, self.shrink_ratio) | |
np.append(possible_ratios, 1) | |
# print(possible_ratios) | |
for ratio in possible_ratios: | |
# print(f"Change shrink ratio to {ratio}") | |
distance = ( | |
polygon_shape.area | |
* (1 - np.power(ratio, 2)) | |
/ polygon_shape.length | |
) | |
shrinked = padding.Execute(-distance) | |
if len(shrinked) == 1: | |
break | |
if shrinked == []: | |
cv2.fillPoly(mask, polygon.astype(np.int32)[np.newaxis, :, :], 0) | |
ignore_tags[i] = True | |
continue | |
for each_shirnk in shrinked: | |
shirnk = np.array(each_shirnk).reshape(-1, 2) | |
cv2.fillPoly(gt, [shirnk.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): | |
""" | |
compute polygon area | |
""" | |
area = 0 | |
q = polygon[-1] | |
for p in polygon: | |
area += p[0] * q[1] - p[1] * q[0] | |
q = p | |
return area / 2.0 | |