Переглянути джерело

feat: add coordinate calculation methods including IoU and overlap ratio, and implement polygon to bbox conversion utilities for enhanced spatial analysis

zhch158_admin 6 днів тому
батько
коміт
fdc80948b0
1 змінених файлів з 172 додано та 1 видалено
  1. 172 1
      zhch/universal_doc_parser/core/coordinate_utils.py

+ 172 - 1
zhch/universal_doc_parser/core/coordinate_utils.py

@@ -2,6 +2,8 @@
 坐标转换工具模块
 
 提供各种坐标转换功能:
+- 底层坐标计算(IoU、重叠比例)
+- 多边形/bbox 格式转换
 - 相对坐标 → 绝对坐标转换
 - OCR 格式转换
 - 旋转坐标逆变换
@@ -9,7 +11,7 @@
 """
 import re
 import json
-from typing import Dict, List, Any, Optional, Tuple
+from typing import Dict, List, Any, Optional, Tuple, Union
 import numpy as np
 from loguru import logger
 
@@ -21,10 +23,179 @@ except ImportError:
     MERGER_AVAILABLE = False
     BBoxExtractor = None
 
+# 导入 MinerU 组件(用于 IoU 计算)
+try:
+    from mineru.utils.boxbase import calculate_iou as mineru_calculate_iou
+    from mineru.utils.boxbase import calculate_overlap_area_2_minbox_area_ratio
+    MINERU_BOXBASE_AVAILABLE = True
+except ImportError:
+    MINERU_BOXBASE_AVAILABLE = False
+    mineru_calculate_iou = None
+    calculate_overlap_area_2_minbox_area_ratio = None
+
 
 class CoordinateUtils:
     """坐标转换工具类"""
     
+    # ==================== 底层坐标计算方法 ====================
+    
+    @staticmethod
+    def calculate_iou(bbox1: List[float], bbox2: List[float]) -> float:
+        """
+        计算两个 bbox 的 IoU(交并比)
+        
+        Args:
+            bbox1: 第一个 bbox [x1, y1, x2, y2]
+            bbox2: 第二个 bbox [x1, y1, x2, y2]
+            
+        Returns:
+            IoU 值
+        """
+        if MINERU_BOXBASE_AVAILABLE and mineru_calculate_iou is not None:
+            return mineru_calculate_iou(bbox1, bbox2)
+        
+        # 备用实现
+        x_left = max(bbox1[0], bbox2[0])
+        y_top = max(bbox1[1], bbox2[1])
+        x_right = min(bbox1[2], bbox2[2])
+        y_bottom = min(bbox1[3], bbox2[3])
+        
+        if x_right < x_left or y_bottom < y_top:
+            return 0.0
+        
+        intersection_area = (x_right - x_left) * (y_bottom - y_top)
+        bbox1_area = (bbox1[2] - bbox1[0]) * (bbox1[3] - bbox1[1])
+        bbox2_area = (bbox2[2] - bbox2[0]) * (bbox2[3] - bbox2[1])
+        
+        if bbox1_area == 0 or bbox2_area == 0:
+            return 0.0
+        
+        return intersection_area / float(bbox1_area + bbox2_area - intersection_area)
+    
+    @staticmethod
+    def calculate_overlap_ratio(bbox1: List[float], bbox2: List[float]) -> float:
+        """
+        计算重叠面积占小框面积的比例
+        
+        Args:
+            bbox1: 第一个 bbox [x1, y1, x2, y2]
+            bbox2: 第二个 bbox [x1, y1, x2, y2]
+            
+        Returns:
+            重叠比例
+        """
+        if MINERU_BOXBASE_AVAILABLE and calculate_overlap_area_2_minbox_area_ratio is not None:
+            return calculate_overlap_area_2_minbox_area_ratio(bbox1, bbox2)
+        
+        # 备用实现
+        x_left = max(bbox1[0], bbox2[0])
+        y_top = max(bbox1[1], bbox2[1])
+        x_right = min(bbox1[2], bbox2[2])
+        y_bottom = min(bbox1[3], bbox2[3])
+        
+        if x_right < x_left or y_bottom < y_top:
+            return 0.0
+        
+        intersection_area = (x_right - x_left) * (y_bottom - y_top)
+        area1 = (bbox1[2] - bbox1[0]) * (bbox1[3] - bbox1[1])
+        area2 = (bbox2[2] - bbox2[0]) * (bbox2[3] - bbox2[1])
+        min_area = min(area1, area2)
+        
+        if min_area == 0:
+            return 0.0
+        
+        return intersection_area / min_area
+    
+    @staticmethod
+    def calculate_overlap_in_bbox1_ratio(
+        bbox1: List[float], 
+        bbox2: List[float]
+    ) -> float:
+        """
+        计算 bbox1 被 bbox2 覆盖的面积比例
+        
+        Args:
+            bbox1: 第一个 bbox [x1, y1, x2, y2]
+            bbox2: 第二个 bbox [x1, y1, x2, y2]
+            
+        Returns:
+            bbox1 被覆盖的比例
+        """
+        x_left = max(bbox1[0], bbox2[0])
+        y_top = max(bbox1[1], bbox2[1])
+        x_right = min(bbox1[2], bbox2[2])
+        y_bottom = min(bbox1[3], bbox2[3])
+        
+        if x_right < x_left or y_bottom < y_top:
+            return 0.0
+        
+        intersection_area = (x_right - x_left) * (y_bottom - y_top)
+        bbox1_area = (bbox1[2] - bbox1[0]) * (bbox1[3] - bbox1[1])
+        
+        if bbox1_area == 0:
+            return 0.0
+        
+        return intersection_area / bbox1_area
+    
+    @staticmethod
+    def poly_to_bbox(poly: Union[List, None]) -> List[float]:
+        """
+        将多边形坐标转换为 bbox 格式
+        
+        Args:
+            poly: 多边形坐标,支持以下格式:
+                - [[x1,y1], [x2,y1], [x2,y2], [x1,y2]] (4个点)
+                - [x1, y1, x2, y1, x2, y2, x1, y2] (8个值)
+                - [x1, y1, x2, y2] (4个值,已是bbox)
+                
+        Returns:
+            bbox [x1, y1, x2, y2]
+        """
+        if not poly:
+            return [0, 0, 0, 0]
+        
+        # 处理嵌套列表格式 [[x1,y1], [x2,y1], ...]
+        if isinstance(poly[0], (list, tuple)):
+            xs = [p[0] for p in poly]
+            ys = [p[1] for p in poly]
+            return [min(xs), min(ys), max(xs), max(ys)]
+        
+        # 处理平面列表格式
+        if len(poly) == 4:
+            # 已经是 bbox 格式
+            return list(poly)
+        elif len(poly) >= 8:
+            # 8点格式:[x1, y1, x2, y1, x2, y2, x1, y2]
+            xs = [poly[i] for i in range(0, len(poly), 2)]
+            ys = [poly[i] for i in range(1, len(poly), 2)]
+            return [min(xs), min(ys), max(xs), max(ys)]
+        
+        return [0, 0, 0, 0]
+    
+    @staticmethod
+    def bbox_to_poly(bbox: List[float]) -> List[List[float]]:
+        """
+        将 bbox 转换为多边形坐标
+        
+        Args:
+            bbox: [x1, y1, x2, y2]
+            
+        Returns:
+            [[x1,y1], [x2,y1], [x2,y2], [x1,y2]]
+        """
+        if not bbox or len(bbox) < 4:
+            return [[0, 0], [0, 0], [0, 0], [0, 0]]
+        
+        x1, y1, x2, y2 = bbox[:4]
+        return [
+            [float(x1), float(y1)],
+            [float(x2), float(y1)],
+            [float(x2), float(y2)],
+            [float(x1), float(y2)]
+        ]
+    
+    # ==================== 图像裁剪 ====================
+    
     @staticmethod
     def crop_region(image: np.ndarray, bbox: List[float]) -> np.ndarray:
         """