Browse Source

feat: Enhance OCR data parsing by introducing automatic tool detection and refining analysis method selection

zhch158_admin 6 hours ago
parent
commit
03f776143b
1 changed files with 38 additions and 8 deletions
  1. 38 8
      table_line_generator/table_line_generator.py

+ 38 - 8
table_line_generator/table_line_generator.py

@@ -62,7 +62,7 @@ class TableLineGenerator:
 
 
 
 
     @staticmethod
     @staticmethod
-    def parse_ocr_data(ocr_result: Dict, tool: str = "ppstructv3") -> Tuple[List[int], Dict]:
+    def parse_ocr_data(ocr_result: Dict, tool: str = "auto") -> Tuple[List[int], Dict]:
         """
         """
         统一的 OCR 数据解析接口(第一步:仅读取数据)
         统一的 OCR 数据解析接口(第一步:仅读取数据)
         
         
@@ -73,14 +73,29 @@ class TableLineGenerator:
         Returns:
         Returns:
             (table_bbox, ocr_data): 表格边界框和文本框列表
             (table_bbox, ocr_data): 表格边界框和文本框列表
         """
         """
-        if tool.lower() == "mineru":
-            return TableLineGenerator._parse_mineru_data(ocr_result)
-        elif tool.lower() in ["ppstructv3", "ppstructure"]:
+        if tool.lower() == "auto":
+            tool_type = TableLineGenerator.detect_ocr_tool_type(ocr_result)
+        else:
+            tool_type = tool.lower() if tool else None
+        
+        if tool_type == "ppstructure":
             return TableLineGenerator._parse_ppstructure_data(ocr_result)
             return TableLineGenerator._parse_ppstructure_data(ocr_result)
+        elif tool_type == "mineru":
+            return TableLineGenerator._parse_mineru_data(ocr_result)
         else:
         else:
             raise ValueError(f"不支持的工具类型: {tool}")
             raise ValueError(f"不支持的工具类型: {tool}")
     
     
     @staticmethod
     @staticmethod
+    def detect_ocr_tool_type(ocr_result: Dict) -> str:
+        """
+        检测 OCR 工具类型
+        """
+        if 'parsing_res_list' in ocr_result and 'overall_ocr_res' in ocr_result:
+            return "ppstructure"
+        else:
+            return "mineru"
+
+    @staticmethod
     def _parse_mineru_data(mineru_result: Union[Dict, List]) -> Tuple[List[int], Dict]:
     def _parse_mineru_data(mineru_result: Union[Dict, List]) -> Tuple[List[int], Dict]:
         """
         """
         解析 MinerU 格式数据(仅提取数据,不分析结构)
         解析 MinerU 格式数据(仅提取数据,不分析结构)
@@ -205,6 +220,14 @@ class TableLineGenerator:
         return table_bbox, ocr_data
         return table_bbox, ocr_data
     
     
     # ==================== 统一接口:第二步 - 分析结构 ====================
     # ==================== 统一接口:第二步 - 分析结构 ====================
+    def detect_analysis_method(self) -> str:
+        """
+        检测分析方法
+        """
+        if 'text_boxes' in self.ocr_data:
+            return "mineru" if any('row' in item and 'col' in item for item in self.ocr_data['text_boxes']) else "cluster"
+        else:
+            return "cluster"
     
     
     def analyze_table_structure(self, 
     def analyze_table_structure(self, 
                                y_tolerance: int = 5,
                                y_tolerance: int = 5,
@@ -231,8 +254,7 @@ class TableLineGenerator:
         # 🔑 自动选择方法
         # 🔑 自动选择方法
         if method == "auto":
         if method == "auto":
             # 根据数据特征自动选择
             # 根据数据特征自动选择
-            has_cell_index = any('row' in item and 'col' in item for item in self.ocr_data.get('text_boxes', []))
-            method = "mineru" if has_cell_index else "cluster"
+            method = self.detect_analysis_method()
             print(f"🤖 自动选择分析方法: {method}")
             print(f"🤖 自动选择分析方法: {method}")
         
         
         # 🔑 根据方法选择算法
         # 🔑 根据方法选择算法
@@ -327,11 +349,18 @@ class TableLineGenerator:
         # 3. 执行倾斜校正 (skew_angle)
         # 3. 执行倾斜校正 (skew_angle)
         if abs(skew_angle) > 0.1:
         if abs(skew_angle) > 0.1:
             print(f"   📐 执行倾斜校正: {skew_angle:.2f}°")
             print(f"   📐 执行倾斜校正: {skew_angle:.2f}°")
-            # 图片逆时针歪了 skew_angle 度,需要顺时针转 skew_angle 度校正
+            # 逻辑说明:
+            # skew_angle 表示图片内容的倾斜角度。
+            # 正值(+) 表示内容逆时针倾斜(左高右低,或者坐标系定义下的逆时针)。
+            # 为了校正,我们需要向相反方向旋转,即顺时针旋转。
+            # PIL.Image.rotate(angle) 中,正值是逆时针旋转,负值是顺时针旋转。
+            # 所以 correction_angle = -skew_angle
             correction_angle = -skew_angle
             correction_angle = -skew_angle
             current_image = current_image.rotate(correction_angle, expand=False, fillcolor='white')
             current_image = current_image.rotate(correction_angle, expand=False, fillcolor='white')
             
             
             # 更新 bbox 坐标
             # 更新 bbox 坐标
+            # 注意: BBoxExtractor.rotate_point 已修正为符合 PIL 的正=逆时针逻辑
+            # 所以这里传入 correction_angle 即可保持一致
             working_text_boxes = BBoxExtractor.correct_boxes_skew(
             working_text_boxes = BBoxExtractor.correct_boxes_skew(
                 working_text_boxes, 
                 working_text_boxes, 
                 correction_angle, 
                 correction_angle, 
@@ -911,7 +940,8 @@ def _parse_table_body_structure(table_body: str) -> Tuple[int, int]:
         
         
         num_rows = len(rows)
         num_rows = len(rows)
         first_row = rows[0]
         first_row = rows[0]
-        num_cols = len(first_row.find_all(['td', 'th']))
+        # 寻找最大列数,避免某些行缺失列
+        num_cols = max(len(row.find_all(['td', 'th'])) for row in rows)
         
         
         return num_rows, num_cols
         return num_rows, num_cols