Browse Source

feat: 增强智能表格处理功能,改进结果提取和边界框计算

zhch158_admin 3 months ago
parent
commit
33a0ed6b69
1 changed files with 117 additions and 7 deletions
  1. 117 7
      zhch/table_mode_selector.py

+ 117 - 7
zhch/table_mode_selector.py

@@ -494,10 +494,63 @@ class IntelligentTableProcessor:
                 'tables': []
                 'tables': []
             }
             }
             
             
-            # 提取表格结果
+            # 提取表格结果 - 根据实际的PP-StructureV3输出结构
             for item in raw_result:
             for item in raw_result:
-                if hasattr(item, 'table_recognition_res') or 'table_recognition_res' in item:
-                    table_res = item.get('table_recognition_res', item.table_recognition_res)
+                print(f"处理结果项: {type(item)}")
+                
+                # 检查是否有table_res_list字段(PP-StructureV3的实际结构)
+                if hasattr(item, 'table_res_list') or 'table_res_list' in item:
+                    table_res_list = item.get('table_res_list', getattr(item, 'table_res_list', []))
+                    
+                    if table_res_list and len(table_res_list) > 0:
+                        formatted_result['table_count'] = len(table_res_list)
+                        
+                        for i, table_item in enumerate(table_res_list):
+                            # 提取HTML内容
+                            html_content = table_item.get('pred_html', 'HTML不可用')
+                            
+                            # 提取表格区域信息
+                            table_region_id = table_item.get('table_region_id', i)
+                            
+                            # 尝试从cell_box_list获取bbox信息
+                            bbox = [0, 0, 100, 100]  # 默认值
+                            if 'cell_box_list' in table_item and table_item['cell_box_list']:
+                                # 从单元格列表计算整体边界框
+                                bbox = self.calculate_table_bbox_from_cells(table_item['cell_box_list'])
+                            
+                            formatted_result['tables'].append({
+                                'table_id': i,
+                                'table_region_id': table_region_id,
+                                'html': html_content,
+                                'bbox': bbox,
+                                'cell_count': len(table_item.get('cell_box_list', [])),
+                                'neighbor_texts': table_item.get('neighbor_texts', '')
+                            })
+                            
+                            print(f"提取表格 {i}: region_id={table_region_id}, cells={len(table_item.get('cell_box_list', []))}")
+                
+                # 检查parsing_res_list(可能包含额外的表格信息)
+                elif hasattr(item, 'parsing_res_list') or 'parsing_res_list' in item:
+                    parsing_res_list = item.get('parsing_res_list', getattr(item, 'parsing_res_list', []))
+                    
+                    for parsing_item in parsing_res_list:
+                        if hasattr(parsing_item, 'label') and parsing_item.label == 'table':
+                            # 这是一个表格解析结果
+                            formatted_result['table_count'] += 1
+                            
+                            html_content = getattr(parsing_item, 'html', 'HTML不可用')
+                            bbox = getattr(parsing_item, 'bbox', [0, 0, 100, 100])
+                            
+                            formatted_result['tables'].append({
+                                'table_id': len(formatted_result['tables']),
+                                'html': html_content,
+                                'bbox': bbox,
+                                'source': 'parsing_res'
+                            })
+                
+                # 兼容旧版本的table_recognition_res结构
+                elif hasattr(item, 'table_recognition_res') or 'table_recognition_res' in item:
+                    table_res = item.get('table_recognition_res', getattr(item, 'table_recognition_res', None))
                     if table_res and len(table_res) > 0:
                     if table_res and len(table_res) > 0:
                         formatted_result['table_count'] = len(table_res)
                         formatted_result['table_count'] = len(table_res)
                         for i, table in enumerate(table_res):
                         for i, table in enumerate(table_res):
@@ -511,7 +564,46 @@ class IntelligentTableProcessor:
             
             
         except Exception as e:
         except Exception as e:
             print(f"结果格式化失败: {e}")
             print(f"结果格式化失败: {e}")
+            import traceback
+            traceback.print_exc()
             return self.create_mock_result(mode)
             return self.create_mock_result(mode)
+
+    def calculate_table_bbox_from_cells(self, cell_box_list):
+        """从单元格列表计算表格的整体边界框"""
+        try:
+            if not cell_box_list:
+                return [0, 0, 100, 100]
+            
+            min_x = float('inf')
+            min_y = float('inf')
+            max_x = float('-inf')
+            max_y = float('-inf')
+            
+            for cell in cell_box_list:
+                # cell格式可能是 [x1, y1, x2, y2] 或其他格式
+                if isinstance(cell, (list, tuple)) and len(cell) >= 4:
+                    x1, y1, x2, y2 = cell[:4]
+                    min_x = min(min_x, x1, x2)
+                    min_y = min(min_y, y1, y2)
+                    max_x = max(max_x, x1, x2)
+                    max_y = max(max_y, y1, y2)
+                elif hasattr(cell, 'bbox'):
+                    bbox = cell.bbox
+                    if len(bbox) >= 4:
+                        x1, y1, x2, y2 = bbox[:4]
+                        min_x = min(min_x, x1, x2)
+                        min_y = min(min_y, y1, y2)
+                        max_x = max(max_x, x1, x2)
+                        max_y = max(max_y, y1, y2)
+            
+            if min_x == float('inf'):
+                return [0, 0, 100, 100]
+            
+            return [int(min_x), int(min_y), int(max_x), int(max_y)]
+            
+        except Exception as e:
+            print(f"计算表格边界框失败: {e}")
+            return [0, 0, 100, 100]
     
     
     def create_mock_result(self, mode):
     def create_mock_result(self, mode):
         """创建模拟结果(用于测试和错误回退)"""
         """创建模拟结果(用于测试和错误回退)"""
@@ -649,7 +741,8 @@ class IntelligentTableProcessor:
             # print(f"优化配置: {optimized_config}")
             # print(f"优化配置: {optimized_config}")
             
             
             # 5. 执行处理
             # 5. 执行处理
-            result = self.execute_with_mode(image_path, best_mode, optimized_config=None)
+            # result = self.execute_with_mode(image_path, best_mode, optimized_config=None)
+            result = self.execute_with_mode(table_image, best_mode, optimized_config=None)
             
             
             return {
             return {
                 'result': result,
                 'result': result,
@@ -715,9 +808,26 @@ def demo_intelligent_table_processing():
             
             
             if process_result.get('tables'):
             if process_result.get('tables'):
                 for i, table in enumerate(process_result['tables']):
                 for i, table in enumerate(process_result['tables']):
-                    print(f"  表格 {i}: bbox={table.get('bbox', 'N/A')}")
-                    html_preview = table.get('html', '')[:100]
-                    print(f"  HTML预览: {html_preview}...")
+                    print(f"\n  表格 {i}:")
+                    print(f"    bbox: {table.get('bbox', 'N/A')}")
+                    print(f"    单元格数量: {table.get('cell_count', 'N/A')}")
+                    print(f"    区域ID: {table.get('table_region_id', 'N/A')}")
+                    
+                    html_content = table.get('html', '')
+                    if len(html_content) > 200:
+                        html_preview = html_content[:200] + "..."
+                    else:
+                        html_preview = html_content
+                    print(f"    HTML预览: {html_preview}")
+                    
+                    # 保存完整HTML到文件
+                    html_filename = f"./table_{i}_result.html"
+                    try:
+                        with open(html_filename, 'w', encoding='utf-8') as f:
+                            f.write(html_content)
+                        print(f"    完整HTML已保存到: {html_filename}")
+                    except Exception as e:
+                        print(f"    保存HTML失败: {e}")
         
         
         # 根据置信度给出建议
         # 根据置信度给出建议
         if result['confidence_score'] > 0.8:
         if result['confidence_score'] > 0.8: