# -*- 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 from shapely.geometry import Polygon import pyclipper __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