import numpy as np
import cv2
import pyclipper
from shapely.geometry import Polygon
def expand_polygon_pyclipper(polygon, expand_ratio):
polygon_shape = Polygon(polygon)
distance = (
polygon_shape.area * (np.power(expand_ratio, 2) - 1) / polygon_shape.length
)
subject = [tuple(l) for l in polygon]
padding = pyclipper.PyclipperOffset()
padding.AddPath(subject, pyclipper.JT_ROUND, pyclipper.ET_CLOSEDPOLYGON)
expanded = padding.Execute(distance)
if expanded == []:
expanded = np.array(expanded)
else:
expanded = np.array(expanded[0]).reshape(-1, 2)
return expanded
class MakeExpandMap:
def __init__(self, min_text_size=8, expand_ratio=1.5, expand_type="pyclipper"):
expand_func_dict = {
"pyclipper": expand_polygon_pyclipper,
}
self.expand_func = expand_func_dict[expand_type]
self.min_text_size = min_text_size
self.expand_ratio = expand_ratio
def __call__(self, data: dict) -> dict:
image = data["img"]
text_polys = data["text_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)
expanded_polygons = []
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:
expanded = self.expand_func(polygon, self.expand_ratio)
expanded_polygons.append(expanded)
if expanded.size == 0:
cv2.fillPoly(mask, polygon.astype(np.int32)[np.newaxis, :, :], 0)
ignore_tags[i] = True
continue
cv2.fillPoly(gt, [expanded.astype(np.int32)], 1)
data["expand_map"] = gt
data["expand_mask"] = mask
data["expanded_polygons"] = expanded_polygons
return data
def validate_polygons(self, polygons, ignore_tags, h, w):
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):
return cv2.contourArea(polygon)
if __name__ == "__main__":
# 示例图像
image = np.ones((200, 200, 3), dtype=np.uint8) * 255
# 示例文本框多边形
text_polys = [
np.array([[50, 50], [150, 50], [150, 100], [50, 100]]),
np.array([[60, 120], [140, 120], [140, 160], [60, 160]])
]
# 示例忽略标志
ignore_tags = [False, False]
# 构建输入数据字典
data = {
"img": image,
"text_polys": text_polys,
"ignore_tags": ignore_tags
}
# 初始化 MakeExpandMap 类
make_expand_map = MakeExpandMap(min_text_size=8, expand_ratio=1.5, expand_type="pyclipper")
# 调用类处理数据
result = make_expand_map(data)
# 获取生成的expand_map和expand_mask
expand_map = result["expand_map"]
expand_mask = result["expand_mask"]
expanded_polygons = result["expanded_polygons"]
# 在原图上绘制expand前的多边形
original_image = image.copy()
for polygon in text_polys:
cv2.polylines(original_image, [polygon.astype(np.int32)], True, (0, 0, 255), 2)
# 在原图上绘制expand后的多边形
expanded_image = image.copy()
for polygon in expanded_polygons:
cv2.polylines(expanded_image, [polygon.astype(np.int32)], True, (0, 255, 0), 2)
# 保存结果图像
cv2.imwrite("original_image.png", original_image)
cv2.imwrite("expanded_image.png", expanded_image)
# cv2.imwrite("expand_map.png", expand_map * 255) # 将expand_map转换为图像
# cv2.imwrite("expand_mask.png", expand_mask * 255) # 将expand_mask转换为图像
# 显示结果
# cv2.imshow("Original Image", original_image)
# cv2.imshow("Expanded Image", expanded_image)
# cv2.imshow("Expand Map", expand_map)
# cv2.imshow("Expand Mask", expand_mask)
# cv2.waitKey(0)
# cv2.destroyAllWindows()
原图:
扩展后: