|
|
@@ -0,0 +1,404 @@
|
|
|
+# zhch/omnidocbench_eval_fixed.py
|
|
|
+import json
|
|
|
+import time
|
|
|
+from pathlib import Path
|
|
|
+from typing import List, Dict, Any, Tuple
|
|
|
+import cv2
|
|
|
+import numpy as np
|
|
|
+from paddlex import create_pipeline
|
|
|
+
|
|
|
+class OmniDocBenchEvaluator:
|
|
|
+ """OmniDocBench评估器（修正版），用于生成符合评测格式的结果"""
|
|
|
+
|
|
|
+ def __init__(self, pipeline_config_path: str = "./PP-StructureV3-zhch.yaml"):
|
|
|
+ """
|
|
|
+ 初始化评估器
|
|
|
+
|
|
|
+ Args:
|
|
|
+ pipeline_config_path: PaddleX pipeline配置文件路径
|
|
|
+ """
|
|
|
+ self.pipeline = create_pipeline(pipeline=pipeline_config_path)
|
|
|
+ self.category_mapping = self._get_category_mapping()
|
|
|
+
|
|
|
+ def _get_category_mapping(self) -> Dict[str, str]:
|
|
|
+ """获取PaddleX类别到OmniDocBench类别的映射"""
|
|
|
+ return {
|
|
|
+ # PaddleX -> OmniDocBench 类别映射
|
|
|
+ 'title': 'title',
|
|
|
+ 'text': 'text_block',
|
|
|
+ 'figure': 'figure',
|
|
|
+ 'figure_caption': 'figure_caption',
|
|
|
+ 'table': 'table',
|
|
|
+ 'table_caption': 'table_caption',
|
|
|
+ 'equation': 'equation_isolated',
|
|
|
+ 'header': 'header',
|
|
|
+ 'footer': 'footer',
|
|
|
+ 'reference': 'reference',
|
|
|
+ 'seal': 'abandon', # 印章通常作为舍弃类
|
|
|
+ 'number': 'page_number',
|
|
|
+ # 添加更多映射关系
|
|
|
+ }
|
|
|
+
|
|
|
+ def evaluate_single_image(self, image_path: str,
|
|
|
+ use_gpu: bool = True,
|
|
|
+ **kwargs) -> Dict[str, Any]:
|
|
|
+ """
|
|
|
+ 评估单张图像
|
|
|
+
|
|
|
+ Args:
|
|
|
+ image_path: 图像路径
|
|
|
+ use_gpu: 是否使用GPU
|
|
|
+ **kwargs: 其他pipeline参数
|
|
|
+
|
|
|
+ Returns:
|
|
|
+ 符合OmniDocBench格式的结果字典
|
|
|
+ """
|
|
|
+ print(f"正在处理图像: {image_path}")
|
|
|
+
|
|
|
+ # 读取图像获取尺寸信息
|
|
|
+ image = cv2.imread(image_path)
|
|
|
+ height, width = image.shape[:2]
|
|
|
+
|
|
|
+ # 运行PaddleX pipeline
|
|
|
+ start_time = time.time()
|
|
|
+
|
|
|
+ output = list(self.pipeline.predict(
|
|
|
+ input=image_path,
|
|
|
+ device="gpu" if use_gpu else "cpu",
|
|
|
+ use_doc_orientation_classify=True,
|
|
|
+ use_doc_unwarping=False,
|
|
|
+ use_seal_recognition=True,
|
|
|
+ use_chart_recognition=True,
|
|
|
+ use_table_recognition=True,
|
|
|
+ use_formula_recognition=True,
|
|
|
+ **kwargs
|
|
|
+ ))
|
|
|
+
|
|
|
+ process_time = time.time() - start_time
|
|
|
+ print(f"处理耗时: {process_time:.2f}秒")
|
|
|
+
|
|
|
+ # 转换为OmniDocBench格式
|
|
|
+ result = self._convert_to_omnidocbench_format(
|
|
|
+ output, image_path, width, height
|
|
|
+ )
|
|
|
+
|
|
|
+ return result
|
|
|
+
|
|
|
+ def _convert_to_omnidocbench_format(self,
|
|
|
+ paddlex_output: List,
|
|
|
+ image_path: str,
|
|
|
+ width: int,
|
|
|
+ height: int) -> Dict[str, Any]:
|
|
|
+ """
|
|
|
+ 将PaddleX输出转换为OmniDocBench格式
|
|
|
+
|
|
|
+ Args:
|
|
|
+ paddlex_output: PaddleX的输出结果列表
|
|
|
+ image_path: 图像路径
|
|
|
+ width: 图像宽度
|
|
|
+ height: 图像高度
|
|
|
+
|
|
|
+ Returns:
|
|
|
+ OmniDocBench格式的结果
|
|
|
+ """
|
|
|
+ layout_dets = []
|
|
|
+ anno_id_counter = 0
|
|
|
+
|
|
|
+ # 处理PaddleX的输出
|
|
|
+ for res in paddlex_output:
|
|
|
+ # 从parsing_res_list中提取布局信息
|
|
|
+ if hasattr(res, 'parsing_res_list') and res.parsing_res_list:
|
|
|
+ parsing_list = res.parsing_res_list
|
|
|
+
|
|
|
+ for item in parsing_list:
|
|
|
+ # 提取边界框和类别
|
|
|
+ bbox = item.get('block_bbox', [])
|
|
|
+ category = item.get('block_label', 'text_block')
|
|
|
+ content = item.get('block_content', '')
|
|
|
+
|
|
|
+ # 转换bbox格式 [x1, y1, x2, y2] -> [x1, y1, x2, y1, x2, y2, x1, y2]
|
|
|
+ if len(bbox) == 4:
|
|
|
+ x1, y1, x2, y2 = bbox
|
|
|
+ poly = [x1, y1, x2, y1, x2, y2, x1, y2]
|
|
|
+ else:
|
|
|
+ poly = bbox
|
|
|
+
|
|
|
+ # 映射类别
|
|
|
+ omni_category = self.category_mapping.get(category, 'text_block')
|
|
|
+
|
|
|
+ # 创建layout检测结果
|
|
|
+ layout_det = {
|
|
|
+ "category_type": omni_category,
|
|
|
+ "poly": poly,
|
|
|
+ "ignore": False,
|
|
|
+ "order": anno_id_counter,
|
|
|
+ "anno_id": anno_id_counter,
|
|
|
+ }
|
|
|
+
|
|
|
+ # 添加文本识别结果
|
|
|
+ if content and content.strip():
|
|
|
+ if omni_category == 'table':
|
|
|
+ # 表格内容作为HTML存储
|
|
|
+ layout_det["html"] = content
|
|
|
+ else:
|
|
|
+ # 其他类型作为文本存储
|
|
|
+ layout_det["text"] = content.strip()
|
|
|
+
|
|
|
+ # 添加span级别的标注（从OCR结果中提取）
|
|
|
+ layout_det["line_with_spans"] = self._extract_spans_from_ocr(
|
|
|
+ res, bbox, omni_category
|
|
|
+ )
|
|
|
+
|
|
|
+ # 添加属性标签
|
|
|
+ layout_det["attribute"] = self._extract_attributes(item, omni_category)
|
|
|
+
|
|
|
+ layout_dets.append(layout_det)
|
|
|
+ anno_id_counter += 1
|
|
|
+
|
|
|
+ # 构建完整结果
|
|
|
+ result = {
|
|
|
+ "layout_dets": layout_dets,
|
|
|
+ "page_info": {
|
|
|
+ "page_no": 0,
|
|
|
+ "height": height,
|
|
|
+ "width": width,
|
|
|
+ "image_path": Path(image_path).name,
|
|
|
+ "page_attribute": self._extract_page_attributes(paddlex_output)
|
|
|
+ },
|
|
|
+ "extra": {
|
|
|
+ "relation": [] # 关系信息，需要根据具体情况提取
|
|
|
+ }
|
|
|
+ }
|
|
|
+
|
|
|
+ return result
|
|
|
+
|
|
|
+ def _extract_spans_from_ocr(self, res, block_bbox: List, category: str) -> List[Dict]:
|
|
|
+ """从OCR结果中提取span级别的标注"""
|
|
|
+ spans = []
|
|
|
+
|
|
|
+ # 如果有OCR结果，提取相关的文本行
|
|
|
+ if hasattr(res, 'overall_ocr_res') and res.overall_ocr_res:
|
|
|
+ ocr_res = res.overall_ocr_res
|
|
|
+
|
|
|
+ if hasattr(ocr_res, 'rec_texts') and hasattr(ocr_res, 'rec_boxes'):
|
|
|
+ texts = ocr_res.rec_texts
|
|
|
+ boxes = ocr_res.rec_boxes
|
|
|
+ scores = getattr(ocr_res, 'rec_scores', [1.0] * len(texts))
|
|
|
+
|
|
|
+ # 检查哪些OCR结果在当前block内
|
|
|
+ if len(block_bbox) == 4:
|
|
|
+ x1, y1, x2, y2 = block_bbox
|
|
|
+
|
|
|
+ for i, (text, box, score) in enumerate(zip(texts, boxes, scores)):
|
|
|
+ if len(box) >= 4:
|
|
|
+ # 检查OCR框是否在block内
|
|
|
+ ocr_x1, ocr_y1, ocr_x2, ocr_y2 = box[:4]
|
|
|
+
|
|
|
+ # 简单的包含检查
|
|
|
+ if (ocr_x1 >= x1 and ocr_y1 >= y1 and
|
|
|
+ ocr_x2 <= x2 and ocr_y2 <= y2):
|
|
|
+
|
|
|
+ span = {
|
|
|
+ "category_type": "text_span",
|
|
|
+ "poly": [ocr_x1, ocr_y1, ocr_x2, ocr_y1,
|
|
|
+ ocr_x2, ocr_y2, ocr_x1, ocr_y2],
|
|
|
+ "ignore": False,
|
|
|
+ "text": text,
|
|
|
+ }
|
|
|
+
|
|
|
+ # 如果置信度太低，可能需要忽略
|
|
|
+ if score < 0.5:
|
|
|
+ span["ignore"] = True
|
|
|
+
|
|
|
+ spans.append(span)
|
|
|
+
|
|
|
+ return spans
|
|
|
+
|
|
|
+ def _extract_attributes(self, item: Dict, category: str) -> Dict:
|
|
|
+ """提取属性标签"""
|
|
|
+ attributes = {}
|
|
|
+
|
|
|
+ # 根据类别提取不同的属性
|
|
|
+ if category == 'table':
|
|
|
+ # 表格属性
|
|
|
+ attributes.update({
|
|
|
+ "table_layout": "vertical", # 需要根据实际情况判断
|
|
|
+ "with_span": False, # 需要检查是否有合并单元格
|
|
|
+ "line": "full_line", # 需要检查线框类型
|
|
|
+ "language": "table_simplified_chinese", # 需要语言检测
|
|
|
+ "include_equation": False,
|
|
|
+ "include_backgroud": False,
|
|
|
+ "table_vertical": False
|
|
|
+ })
|
|
|
+
|
|
|
+ # 检查表格内容是否有合并单元格
|
|
|
+ content = item.get('block_content', '')
|
|
|
+ if 'colspan' in content or 'rowspan' in content:
|
|
|
+ attributes["with_span"] = True
|
|
|
+
|
|
|
+ elif category in ['text_block', 'title']:
|
|
|
+ # 文本属性
|
|
|
+ attributes.update({
|
|
|
+ "text_language": "text_simplified_chinese",
|
|
|
+ "text_background": "white",
|
|
|
+ "text_rotate": "normal"
|
|
|
+ })
|
|
|
+
|
|
|
+ elif 'equation' in category:
|
|
|
+ # 公式属性
|
|
|
+ attributes.update({
|
|
|
+ "formula_type": "print"
|
|
|
+ })
|
|
|
+
|
|
|
+ return attributes
|
|
|
+
|
|
|
+ def _extract_page_attributes(self, paddlex_output) -> Dict:
|
|
|
+ """提取页面级别的属性"""
|
|
|
+ return {
|
|
|
+ "data_source": "research_report", # 需要根据实际情况判断
|
|
|
+ "language": "simplified_chinese",
|
|
|
+ "layout": "single_column",
|
|
|
+ "watermark": False,
|
|
|
+ "fuzzy_scan": False,
|
|
|
+ "colorful_backgroud": False
|
|
|
+ }
|
|
|
+
|
|
|
+ def load_existing_result(self, result_path: str) -> Dict[str, Any]:
|
|
|
+ """
|
|
|
+ 从已有的PaddleX结果文件加载数据进行转换
|
|
|
+
|
|
|
+ Args:
|
|
|
+ result_path: PaddleX结果JSON文件路径
|
|
|
+
|
|
|
+ Returns:
|
|
|
+ OmniDocBench格式的结果字典
|
|
|
+ """
|
|
|
+ with open(result_path, 'r', encoding='utf-8') as f:
|
|
|
+ data = json.load(f)
|
|
|
+
|
|
|
+ # 从结果文件中提取图像信息
|
|
|
+ input_path = data.get('input_path', '')
|
|
|
+
|
|
|
+ # 读取图像获取尺寸
|
|
|
+ if input_path and Path(input_path).exists():
|
|
|
+ image = cv2.imread(input_path)
|
|
|
+ height, width = image.shape[:2]
|
|
|
+ image_name = Path(input_path).name
|
|
|
+ else:
|
|
|
+ # 如果图像路径不存在，使用默认值
|
|
|
+ height, width = 1600, 1200
|
|
|
+ image_name = "unknown.png"
|
|
|
+
|
|
|
+ # 转换格式
|
|
|
+ result = self._convert_paddlex_result_to_omnidocbench(
|
|
|
+ data, image_name, width, height
|
|
|
+ )
|
|
|
+
|
|
|
+ return result
|
|
|
+
|
|
|
+ def _convert_paddlex_result_to_omnidocbench(self,
|
|
|
+ paddlex_result: Dict,
|
|
|
+ image_name: str,
|
|
|
+ width: int,
|
|
|
+ height: int) -> Dict[str, Any]:
|
|
|
+ """
|
|
|
+ 将已有的PaddleX结果转换为OmniDocBench格式
|
|
|
+ """
|
|
|
+ layout_dets = []
|
|
|
+ anno_id_counter = 0
|
|
|
+
|
|
|
+ # 从parsing_res_list中提取布局信息
|
|
|
+ parsing_list = paddlex_result.get('parsing_res_list', [])
|
|
|
+
|
|
|
+ for item in parsing_list:
|
|
|
+ # 提取边界框和类别
|
|
|
+ bbox = item.get('block_bbox', [])
|
|
|
+ category = item.get('block_label', 'text_block')
|
|
|
+ content = item.get('block_content', '')
|
|
|
+
|
|
|
+ # 转换bbox格式
|
|
|
+ if len(bbox) == 4:
|
|
|
+ x1, y1, x2, y2 = bbox
|
|
|
+ poly = [x1, y1, x2, y1, x2, y2, x1, y2]
|
|
|
+ else:
|
|
|
+ poly = bbox
|
|
|
+
|
|
|
+ # 映射类别
|
|
|
+ omni_category = self.category_mapping.get(category, 'text_block')
|
|
|
+
|
|
|
+ # 创建layout检测结果
|
|
|
+ layout_det = {
|
|
|
+ "category_type": omni_category,
|
|
|
+ "poly": poly,
|
|
|
+ "ignore": False,
|
|
|
+ "order": anno_id_counter,
|
|
|
+ "anno_id": anno_id_counter,
|
|
|
+ }
|
|
|
+
|
|
|
+ # 添加内容
|
|
|
+ if content and content.strip():
|
|
|
+ if omni_category == 'table':
|
|
|
+ layout_det["html"] = content
|
|
|
+ else:
|
|
|
+ layout_det["text"] = content.strip()
|
|
|
+
|
|
|
+ # 添加属性
|
|
|
+ layout_det["attribute"] = self._extract_attributes(item, omni_category)
|
|
|
+ layout_det["line_with_spans"] = [] # 简化处理
|
|
|
+
|
|
|
+ layout_dets.append(layout_det)
|
|
|
+ anno_id_counter += 1
|
|
|
+
|
|
|
+ # 构建完整结果
|
|
|
+ result = {
|
|
|
+ "layout_dets": layout_dets,
|
|
|
+ "page_info": {
|
|
|
+ "page_no": 0,
|
|
|
+ "height": height,
|
|
|
+ "width": width,
|
|
|
+ "image_path": image_name,
|
|
|
+ "page_attribute": {
|
|
|
+ "data_source": "research_report",
|
|
|
+ "language": "simplified_chinese",
|
|
|
+ "layout": "single_column",
|
|
|
+ "watermark": False,
|
|
|
+ "fuzzy_scan": False,
|
|
|
+ "colorful_backgroud": False
|
|
|
+ }
|
|
|
+ },
|
|
|
+ "extra": {
|
|
|
+ "relation": []
|
|
|
+ }
|
|
|
+ }
|
|
|
+
|
|
|
+ return result
|
|
|
+
|
|
|
+def convert_existing_results():
|
|
|
+ """转换已有的PaddleX结果"""
|
|
|
+ evaluator = OmniDocBenchEvaluator()
|
|
|
+
|
|
|
+ # 示例：转换单个结果文件
|
|
|
+ result_file = "./sample_data/single_pipeline_output/PP-StructureV3-zhch/300674-母公司现金流量表-扫描_res.json"
|
|
|
+
|
|
|
+ if Path(result_file).exists():
|
|
|
+ print(f"正在转换结果文件: {result_file}")
|
|
|
+
|
|
|
+ omnidocbench_result = evaluator.load_existing_result(result_file)
|
|
|
+
|
|
|
+ # 保存转换后的结果
|
|
|
+ output_file = "./omnidocbench_converted_result.json"
|
|
|
+ with open(output_file, 'w', encoding='utf-8') as f:
|
|
|
+ json.dump([omnidocbench_result], f, ensure_ascii=False, indent=2)
|
|
|
+
|
|
|
+ print(f"转换完成，结果保存至: {output_file}")
|
|
|
+ print(f"检测到的布局元素数量: {len(omnidocbench_result['layout_dets'])}")
|
|
|
+
|
|
|
+ # 显示检测到的元素
|
|
|
+ for i, item in enumerate(omnidocbench_result['layout_dets']):
|
|
|
+ print(f" {i+1}. {item['category_type']}: {item.get('text', item.get('html', ''))[:50]}...")
|
|
|
+
|
|
|
+ else:
|
|
|
+ print(f"结果文件不存在: {result_file}")
|
|
|
+
|
|
|
+if __name__ == "__main__":
|
|
|
+ convert_existing_results()
|
