Update utils/ocr_utils.py

utils/ocr_utils.py

@@ -9,48 +9,71 @@ from shapely.geometry import box as shapely_box
 from shapely.geometry import Polygon
 from shapely.ops import unary_union
 import networkx as nx
+from shapely.ops import unary_union
 
 ocr_model = PaddleOCR(use_textline_orientation=True, lang='ch')
 
-def inflate_bbox(polygon, inflation_ratio=0.05):
-    xs, ys = zip(*polygon)
-    x_min, x_max = min(xs), max(xs)
-    y_min, y_max = min(ys), max(ys)
-
-    width = x_max - x_min
-    height = y_max - y_min
 
-    x_min -= width * inflation_ratio
-    x_max += width * inflation_ratio
-    y_min -= height * inflation_ratio
-    y_max += height * inflation_ratio
+def inflate_polygon(polygon_points, percent=0.05):
+    poly = Polygon(polygon_points)
+    if not poly.is_valid:
+        poly = poly.convex_hull
+    minx, miny, maxx, maxy = poly.bounds
+    diagonal = ((maxx - minx)**2 + (maxy - miny)**2)**0.5
+    inflate_dist = diagonal * percent
+    return poly.buffer(inflate_dist)
 
-    return shapely_box(x_min, y_min, x_max, y_max)
+def group_nearby_boxes(lines, inflation_percent=0.05):
+    from collections import defaultdict
 
-def group_nearby_boxes(lines, inflation_ratio=0.05):
     n = len(lines)
-    inflated_boxes = [inflate_bbox(poly, inflation_ratio) for poly, _ in lines]
-
-    # Build graph where edges represent overlapping boxes
-    G = nx.Graph()
-    G.add_nodes_from(range(n))
-
+    inflated_polys = []
+    original_polys = []
+    texts = []
+
+    for poly_pts, text in lines:
+        inflated = inflate_polygon(poly_pts, percent=inflation_percent)
+        original = Polygon(poly_pts)
+        inflated_polys.append(inflated)
+        original_polys.append(original)
+        texts.append(text)
+
+    # Build connectivity graph
+    adjacency = defaultdict(set)
     for i in range(n):
         for j in range(i + 1, n):
-            if inflated_boxes[i].intersects(inflated_boxes[j]):
-                G.add_edge(i, j)
+            if inflated_polys[i].intersects(inflated_polys[j]):
+                adjacency[i].add(j)
+                adjacency[j].add(i)
 
-    # Extract connected components as groups
+    # DFS to find connected components
+    visited = [False] * n
     groups = []
-    for component in nx.connected_components(G):
-        group_polys = [lines[i][0] for i in component]
-        group_texts = [lines[i][1] for i in component]
-        groups.append({
+
+    def dfs(i, group):
+        visited[i] = True
+        group.append(i)
+        for neighbor in adjacency[i]:
+            if not visited[neighbor]:
+                dfs(neighbor, group)
+
+    for i in range(n):
+        if not visited[i]:
+            group = []
+            dfs(i, group)
+            groups.append(group)
+
+    # Construct output groups
+    grouped = []
+    for group in groups:
+        group_polys = [list(original_polys[i].exterior.coords) for i in group]
+        group_texts = [texts[i] for i in group]
+        grouped.append({
             "polygons": group_polys,
             "texts": group_texts
         })
 
-    return groups
+    return grouped
 
 def extract_and_translate_chunk(image: Image.Image):
     np_img = np.array(image)