Update utils/ocr_utils.py

utils/ocr_utils.py

@@ -6,72 +6,59 @@ ocr_model = PaddleOCR(use_textline_orientation=True, lang='ch')
 
 def group_nearby_boxes(lines, max_y_gap=50):
     groups = []
-    used = set()
-
-    def is_valid_polygon(poly):
-        return isinstance(poly, (list, tuple)) and all(
-            isinstance(p, (list, tuple)) and len(p) == 2 for p in poly
-        )
+    used = [False] * len(lines)
 
     for i, (poly_i, text_i) in enumerate(lines):
-        if i in used or not is_valid_polygon(poly_i):
+        if used[i]:
             continue
 
-        group = [(poly_i, text_i)]
-        used.add(i)
+        group_polys = [poly_i]
+        group_texts = [text_i]
+        used[i] = True
 
-        xi_min, yi_min = min(pt[1] for pt in poly_i), min(pt[0] for pt in poly_i)
-        xi_max, yi_max = max(pt[1] for pt in poly_i), max(pt[0] for pt in poly_i)
+        xi_center = np.mean([p[0] for p in poly_i])
+        yi_center = np.mean([p[1] for p in poly_i])
 
         for j, (poly_j, text_j) in enumerate(lines):
-            if j in used or not is_valid_polygon(poly_j):
+            if used[j]:
                 continue
-
-            xj_min, yj_min = min(pt[1] for pt in poly_j), min(pt[0] for pt in poly_j)
-            xj_max, yj_max = max(pt[1] for pt in poly_j), max(pt[0] for pt in poly_j)
-
-            if abs(yj_min - yi_min) < max_y_gap or abs(yj_max - yi_max) < max_y_gap:
-                group.append((poly_j, text_j))
-                used.add(j)
-
-        groups.append(group)
-
+            xj_center = np.mean([p[0] for p in poly_j])
+            yj_center = np.mean([p[1] for p in poly_j])
+
+            # If Y-distance is small, consider them in the same group (horizontal alignment)
+            if abs(yj_center - yi_center) < max_y_gap:
+                group_polys.append(poly_j)
+                group_texts.append(text_j)
+                used[j] = True
+
+        groups.append({
+            "polygons": group_polys,
+            "texts": group_texts
+        })
     return groups
 
+
 def extract_and_translate_chunk(image: Image.Image):
     np_img = np.array(image)
     results = ocr_model.ocr(np_img)
 
-    if not results or not isinstance(results[0], dict):
-        print("⚠️ No OCR results or unexpected format")
+    # Handle both PaddleOCR and fallback
+    if not results or not results[0]:
         return []
 
-    ocr_data = results[0]
-
-    rec_texts = ocr_data.get("rec_texts", [])
-    rec_polys = ocr_data.get("rec_polys", [])
-
-    if not rec_texts or not rec_polys:
-        print("⚠️ Missing OCR text or polygons")
-        return []
-
-    # Step 1: Build list of (polygon, text)
-    lines = []
-    for poly, text in zip(rec_polys, rec_texts):
-        polygon = [(int(x), int(y)) for x, y in poly]
-        lines.append((polygon, text.strip()))
+    lines = results[0]  # [(polygon, text), ...]
 
     print("🔍 OCR Raw Output:", lines)
 
-    # Step 2: Group nearby text boxes
+    # Group nearby lines into combined regions
     grouped = group_nearby_boxes(lines)
-
-    # Step 3: Translate and compute convex hulls
     translations = []
+
     for group in grouped:
         polygons = group["polygons"]
-        merged_text = "".join(group["texts"]).strip()
+        texts = group["texts"]
 
+        merged_text = "".join(texts).strip()
         if not merged_text:
             continue
 
@@ -81,9 +68,10 @@ def extract_and_translate_chunk(image: Image.Image):
             print("⚠️ Translation failed:", e)
             translated = ""
 
+        # Compute convex hull from all polygon points
         all_points = np.array([pt for polygon in polygons for pt in polygon])
         if len(all_points) < 3:
-            continue  # convex hull needs ≥ 3
+            continue
 
         hull_indices = ConvexHull(all_points).vertices
         hull = [tuple(map(int, all_points[i])) for i in hull_indices]