【深度学习】OCR中的Shrink反向扩展

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()

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134

原图：

扩展后：