Update utils/polygon_utils.py

utils/polygon_utils.py

@@ -1,177 +1,159 @@
 """
 Enhanced polygon utilities with bubble-based correction
 """
-from shapely.geometry import Polygon, MultiPoint, Point
 import os
 import cv2
 import numpy as np
-from PIL import Image, ImageDraw, ImageFont
 import textwrap
+from shapely.geometry import Polygon, MultiPoint, Point
+from PIL import Image, ImageDraw, ImageFont
 
-FONT_PATH = os.path.join(os.path.dirname(__file__), "..", "NotoSansSC-Regular.ttf")
+FONT_PATH = os.path.abspath(
+    os.path.join(os.path.dirname(__file__), "..", "NotoSansSC-Regular.ttf")
+)
 
 
+# ============================ Geometry Helpers ============================
+
 def calculate_iou(poly1, poly2):
     """Calculate Intersection over Union between two polygons"""
     try:
         p1 = Polygon(poly1)
         p2 = Polygon(poly2)
-        
+
         if not p1.is_valid:
             p1 = p1.buffer(0)
         if not p2.is_valid:
             p2 = p2.buffer(0)
-            
+
         intersection = p1.intersection(p2).area
         union = p1.union(p2).area
-        
-        return intersection / union if union > 0 else 0
-    except:
-        return 0
+
+        return intersection / union if union > 0 else 0.0
+    except Exception as e:
+        print(f"⚠️ calculate_iou failed: {e}")
+        return 0.0
 
 
 def calculate_polygon_overlap(ocr_poly, bubble_poly):
     """
-    Calculate what percentage of OCR polygon is inside bubble polygon
+    Calculate what percentage of OCR polygon is inside bubble polygon.
+    overlap = area(ocr ∩ bubble) / area(ocr)
     """
     try:
         ocr_shape = Polygon(ocr_poly)
         bubble_shape = Polygon(bubble_poly)
-        
+
         if not ocr_shape.is_valid:
             ocr_shape = ocr_shape.buffer(0)
         if not bubble_shape.is_valid:
             bubble_shape = bubble_shape.buffer(0)
-            
+
         intersection = ocr_shape.intersection(bubble_shape).area
         ocr_area = ocr_shape.area
-        
-        return intersection / ocr_area if ocr_area > 0 else 0
-    except:
-        return 0
+
+        return intersection / ocr_area if ocr_area > 0 else 0.0
+    except Exception as e:
+        print(f"⚠️ calculate_polygon_overlap failed: {e}")
+        return 0.0
 
 
-def find_matching_bubble(ocr_polygon, bubble_polygons, iou_threshold=0.3, overlap_threshold=0.5):
+def match_polygon_to_bubble_by_overlap(ocr_poly, bubble_polygons, min_overlap=0.15):
     """
-    Find the best matching bubble for an OCR polygon.
-    
-    Uses two criteria:
-    1. IoU (Intersection over Union) - measures geometric similarity
-    2. Overlap ratio - what % of OCR box is inside bubble
-    
-    Returns: (best_bubble_index, best_bubble_polygon) or (None, None)
+    Return index of bubble with the highest overlap ratio with OCR polygon.
+    overlap = area(ocr ∩ bubble) / area(ocr)
+
+    If best overlap < min_overlap → no match.
     """
     if not bubble_polygons:
-        return None, None
-    
-    ocr_center = np.mean(ocr_polygon, axis=0)
+        return None
+
     best_idx = None
-    best_score = 0
-    
-    for idx, bubble_poly in enumerate(bubble_polygons):
-        # Check if OCR center is inside bubble (strong indicator)
-        try:
-            bubble_shape = Polygon(bubble_poly)
-            if not bubble_shape.is_valid:
-                bubble_shape = bubble_shape.buffer(0)
-            center_inside = bubble_shape.contains(Point(ocr_center))
-        except:
-            center_inside = False
-        
-        # Calculate overlap and IoU
-        overlap = calculate_polygon_overlap(ocr_polygon, bubble_poly)
-        iou = calculate_iou(ocr_polygon, bubble_poly)
-        
-        # Scoring: prioritize center inside + high overlap
-        score = 0
-        if center_inside:
-            score += 0.5
-        score += overlap * 0.3
-        score += iou * 0.2
-        
-        if score > best_score and (overlap > overlap_threshold or iou > iou_threshold or center_inside):
-            best_score = score
+    best_overlap = 0.0
+
+    for idx, bp in enumerate(bubble_polygons):
+        overlap = calculate_polygon_overlap(ocr_poly, bp)
+        if overlap > best_overlap:
+            best_overlap = overlap
             best_idx = idx
-    
-    if best_idx is not None:
-        return best_idx, bubble_polygons[best_idx]
-    
-    return None, None
 
+    if best_idx is not None and best_overlap >= min_overlap:
+        return best_idx
 
-def correct_polygon_with_bubble(ocr_polygon, bubble_polygon, strategy="intersect"):
+    return None
+
+
+# ====================== Polygon Correction with Bubbles ===================
+
+def correct_polygon_with_bubble(ocr_polygon, bubble_polygon, strategy="hybrid"):
     """
     Correct OCR polygon using bubble polygon.
-    
+
     Strategies:
-    - "bubble": Use bubble polygon directly
-    - "intersect": Use intersection of OCR and bubble
-    - "expand": Expand OCR polygon to fit bubble better
-    - "hybrid": Smart blend based on relative sizes
+    - "bubble":     Use bubble polygon directly
+    - "intersect":  Use intersection of OCR and bubble
+    - "expand":     Slightly expand OCR region inside bubble
+    - "hybrid":     Choose based on relative sizes & intersection
     """
     try:
         ocr_shape = Polygon(ocr_polygon)
         bubble_shape = Polygon(bubble_polygon)
-        
+
         if not ocr_shape.is_valid:
             ocr_shape = ocr_shape.buffer(0)
         if not bubble_shape.is_valid:
             bubble_shape = bubble_shape.buffer(0)
-        
+
+        # ---- Strategy: use bubble fully ----
         if strategy == "bubble":
-            # Use bubble directly (safest for ensuring text stays in bounds)
-            return list(bubble_shape.exterior.coords)[:-1]
-        
-        elif strategy == "intersect":
-            # Use intersection region
-            intersection = ocr_shape.intersection(bubble_shape)
-            if intersection.is_empty or intersection.area < ocr_shape.area * 0.3:
-                # Intersection too small, fall back to bubble
-                return list(bubble_shape.exterior.coords)[:-1]
-            
-            if intersection.geom_type == 'Polygon':
-                return list(intersection.exterior.coords)[:-1]
-            else:
-                # MultiPolygon or other - use largest piece
-                polys = list(intersection.geoms) if hasattr(intersection, 'geoms') else [intersection]
-                largest = max(polys, key=lambda p: p.area if hasattr(p, 'area') else 0)
-                return list(largest.exterior.coords)[:-1]
-        
-        elif strategy == "expand":
-            # Expand OCR box slightly toward bubble edges
-            ocr_center = np.array(ocr_shape.centroid.coords[0])
-            bubble_center = np.array(bubble_shape.centroid.coords[0])
-            
-            # If OCR is much smaller than bubble, expand it
-            if bubble_shape.area > ocr_shape.area * 2:
-                expanded = ocr_shape.buffer(10)  # expand by 10 pixels
-                clipped = expanded.intersection(bubble_shape)
-                if clipped.geom_type == 'Polygon':
-                    return list(clipped.exterior.coords)[:-1]
-            
-            return list(bubble_shape.exterior.coords)[:-1]
-        
-        elif strategy == "hybrid":
-            # Smart strategy: choose based on size ratio
-            size_ratio = bubble_shape.area / ocr_shape.area if ocr_shape.area > 0 else 999
-            
+            return [(int(x), int(y)) for x, y in bubble_shape.exterior.coords[:-1]]
+
+        # ---- Strategy: intersection region ----
+        if strategy == "intersect":
+            inter = ocr_shape.intersection(bubble_shape)
+            if inter.is_empty or inter.area < ocr_shape.area * 0.3:
+                # Intersection too small, bubble is safer
+                return [(int(x), int(y)) for x, y in bubble_shape.exterior.coords[:-1]]
+
+            if inter.geom_type == "Polygon":
+                return [(int(x), int(y)) for x, y in inter.exterior.coords[:-1]]
+
+            polys = list(inter.geoms) if hasattr(inter, "geoms") else [inter]
+            largest = max(polys, key=lambda p: p.area if hasattr(p, "area") else 0)
+            return [(int(x), int(y)) for x, y in largest.exterior.coords[:-1]]
+
+        # ---- Strategy: expand OCR slightly toward bubble ----
+        if strategy == "expand":
+            expanded = ocr_shape.buffer(10)  # ~10px expansion
+            clipped = expanded.intersection(bubble_shape)
+            if not clipped.is_empty and clipped.area > ocr_shape.area * 0.5:
+                if clipped.geom_type == "Polygon":
+                    return [(int(x), int(y)) for x, y in clipped.exterior.coords[:-1]]
+
+            return [(int(x), int(y)) for x, y in bubble_shape.exterior.coords[:-1]]
+
+        # ---- Strategy: hybrid ----
+        if strategy == "hybrid":
+            size_ratio = (
+                bubble_shape.area / ocr_shape.area if ocr_shape.area > 0 else 999
+            )
+
             if size_ratio > 3:
-                # Bubble much larger - likely OCR missed some text area
-                # Use slightly shrunk bubble
+                # Bubble is much larger than OCR region: likely multi-line speech
                 shrunk = bubble_shape.buffer(-5)
                 if shrunk.is_empty:
-                    return list(bubble_shape.exterior.coords)[:-1]
-                return list(shrunk.exterior.coords)[:-1]
+                    return [(int(x), int(y)) for x, y in bubble_shape.exterior.coords[:-1]]
+                return [(int(x), int(y)) for x, y in shrunk.exterior.coords[:-1]]
             elif size_ratio < 1.5:
-                # Similar sizes - use intersection
+                # Similar sizes: use intersection
                 return correct_polygon_with_bubble(ocr_polygon, bubble_polygon, "intersect")
             else:
-                # Moderate difference - use bubble
-                return list(bubble_shape.exterior.coords)[:-1]
-        
+                # Moderate difference → bubble is still safer
+                return [(int(x), int(y)) for x, y in bubble_shape.exterior.coords[:-1]]
+
         # Fallback
-        return list(bubble_shape.exterior.coords)[:-1]
-    
+        return [(int(x), int(y)) for x, y in bubble_shape.exterior.coords[:-1]]
+
     except Exception as e:
         print(f"⚠️ Polygon correction failed: {e}, using original OCR polygon")
         return ocr_polygon
@@ -180,58 +162,58 @@ def correct_polygon_with_bubble(ocr_polygon, bubble_polygon, strategy="intersect
 def correct_ocr_polygons_with_bubbles(translations, bubble_polygons, strategy="hybrid"):
     """
     Correct all OCR polygons using detected bubbles.
-    
-    Args:
-        translations: List of OCR results with 'polygon' key
-        bubble_polygons: List of detected bubble polygons
-        strategy: Correction strategy (see correct_polygon_with_bubble)
-    
+
+    Adds:
+      - "original_polygon"
+      - "matched_bubble_idx"
+
     Returns:
-        Updated translations with corrected polygons
+      updated translations list
     """
     corrected = []
-    unmatched_ocr = []
-    
+    unmatched = 0
+
     for t in translations:
         ocr_poly = t.get("polygon") or t.get("polygons")
         if not ocr_poly:
             corrected.append(t)
             continue
-        
-        bubble_idx, bubble_poly = find_matching_bubble(ocr_poly, bubble_polygons)
-        
-        if bubble_poly is not None:
-            # Correct the polygon
+
+        best_idx = match_polygon_to_bubble_by_overlap(ocr_poly, bubble_polygons)
+
+        t_copy = t.copy()
+        t_copy["original_polygon"] = ocr_poly
+
+        if best_idx is not None:
+            bubble_poly = bubble_polygons[best_idx]
             corrected_poly = correct_polygon_with_bubble(ocr_poly, bubble_poly, strategy)
-            
-            t_copy = t.copy()
             t_copy["polygon"] = corrected_poly
-            t_copy["original_polygon"] = ocr_poly  # keep for debugging
-            t_copy["matched_bubble_idx"] = bubble_idx
-            corrected.append(t_copy)
+            t_copy["matched_bubble_idx"] = best_idx
         else:
-            # No matching bubble - keep original
-            t_copy = t.copy()
+            # No match → keep original OCR polygon
             t_copy["matched_bubble_idx"] = None
-            corrected.append(t_copy)
-            unmatched_ocr.append(t)
-    
-    if unmatched_ocr:
-        print(f"ℹ️ {len(unmatched_ocr)}/{len(translations)} OCR regions had no matching bubble")
-    
+            t_copy["polygon"] = ocr_poly
+            unmatched += 1
+
+        corrected.append(t_copy)
+
+    if unmatched:
+        print(f"ℹ️ {unmatched}/{len(translations)} OCR regions had no matching bubble")
+
     return corrected
 
 
-# ========================================================================
-# Existing utility functions (kept for compatibility)
-# ========================================================================
+# ========================= Basic Polygon Utilities =======================
 
 def shrink_or_expand_polygon(polygon, shrink_ratio=0.9):
     """
     Resize a polygon around its centroid.
-    ratio < 1 → shrink
-    ratio > 1 → expand
+    shrink_ratio < 1 → shrink
+    shrink_ratio > 1 → expand
     """
+    if not polygon:
+        return polygon
+
     ratio = shrink_ratio
     cx = sum(x for x, _ in polygon) / len(polygon)
     cy = sum(y for _, y in polygon) / len(polygon)
@@ -249,12 +231,7 @@ def inpaint_polygon(img: Image.Image, polygon, mode="auto", fallback_color=(255,
     pts = np.array(polygon, np.int32).reshape((-1, 1, 2))
     cv2.fillPoly(mask, [pts], 255)
 
-    if mode == "fill":
-        img_copy = img.copy()
-        draw = ImageDraw.Draw(img_copy)
-        draw.polygon(polygon, fill=fallback_color)
-        return img_copy
-
+    # Could use cv2.inpaint for fancy filling; for manga bubbles simple fill is OK
     img_copy = img.copy()
     draw = ImageDraw.Draw(img_copy)
     draw.polygon(polygon, fill=fallback_color)
@@ -263,37 +240,65 @@ def inpaint_polygon(img: Image.Image, polygon, mode="auto", fallback_color=(255,
 
 def merge_polygons_to_convex_hull(polygons):
     points = [pt for poly in polygons for pt in poly]
+    if not points:
+        return []
     hull = MultiPoint(points).convex_hull
-    return list(hull.exterior.coords)
+    return [(int(x), int(y)) for x, y in hull.exterior.coords[:-1]]
 
 
-def render_translated_chunk(img: Image.Image, translations, font_path="NotoSansSC-Regular.ttf", font_scale=1.0):
-    img_copy = img.copy()
+# ======================== Rendering / Text Drawing =======================
 
+def render_translated_chunk(img: Image.Image, translations, font_path=None, font_scale=1.0):
+    """
+    Render list of translations (with 'polygon' and 'translated') onto image.
+    """
+    img_copy = img.copy()
     for entry in translations:
         polygon = entry.get("polygon") or entry.get("polygons")
         text = entry.get("translated", "")
 
         if polygon and text:
-            img_copy = draw_translated_text_convex(img_copy, polygon, text, font_path, font_scale)
-
+            img_copy = draw_translated_text_convex(
+                img_copy,
+                polygon,
+                text,
+                font_path=font_path or FONT_PATH,
+                font_scale=font_scale
+            )
     return img_copy
 
 
-def draw_translated_text_convex(img, polygon_coords, text, font_path="NotoSansSC-Regular.ttf", font_scale=1.0):
+def draw_translated_text_convex(img, polygon_coords, text, font_path=None, font_scale=1.0):
+    """
+    Inpaint inside polygon and draw wrapped text centered in it.
+    """
+    if font_path is None:
+        font_path = FONT_PATH
+
     font_polygon = polygon_coords
     render_polygon = shrink_or_expand_polygon(polygon_coords, shrink_ratio=0.9)
+
     img = inpaint_polygon(img, render_polygon, mode="auto", fallback_color=(255, 255, 255))
+
     debug_color = (180, 255, 180)
     draw = ImageDraw.Draw(img)
     draw.line(render_polygon + [render_polygon[0]], fill=debug_color, width=1)
-    draw_wrapped_text(img, render_polygon, text, font_path, polygon_for_size=font_polygon, font_scale=font_scale)
+
+    draw_wrapped_text(
+        img,
+        render_polygon,
+        text,
+        font_path,
+        polygon_for_size=font_polygon,
+        font_scale=font_scale,
+    )
     return img
 
 
 def draw_wrapped_text(img, polygon, text, font_path, polygon_for_size=None, font_scale=1.0):
-    from PIL import ImageDraw, ImageFont
-    import textwrap
+    """
+    Draw wrapped text centered in the polygon bounding box.
+    """
     polygon_for_size = polygon_for_size or polygon
     draw = ImageDraw.Draw(img)
     xs, ys = zip(*polygon_for_size)
@@ -301,15 +306,22 @@ def draw_wrapped_text(img, polygon, text, font_path, polygon_for_size=None, font
     y_min, y_max = min(ys), max(ys)
     box_width = x_max - x_min
     box_height = y_max - y_min
+
+    if box_width <= 0 or box_height <= 0:
+        return
+
     avg_char_width = 0.4
-    estimated_size = int(min(box_height / 1.2, box_width / (len(text) * avg_char_width)))
+    estimated_size = int(min(box_height / 1.2, box_width / (len(text) * avg_char_width + 1)))
     estimated_size = max(6, estimated_size)
     font_size = int(estimated_size * font_scale)
+
     font = ImageFont.truetype(font_path, font_size)
     max_chars = max(1, int(box_width / (font.getbbox("A")[2] + 1)))
+
     wrapped = textwrap.fill(text, width=max_chars)
     bbox = draw.textbbox((0, 0), wrapped, font=font)
     text_w, text_h = bbox[2] - bbox[0], bbox[3] - bbox[1]
     x = x_min + (box_width - text_w) / 2
     y = y_min + (box_height - text_h) / 2
-    draw.text((x, y), wrapped, font=font, fill="black", align="center")
+
+    draw.text((x, y), wrapped, font=font, fill="black", align="center")