PaddleX/zhch/omnidocbench_eval.py

# 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
import os
import glob
import traceback

class OmniDocBenchEvaluator:
    """
    OmniDocBench评估器（修正版），用于生成符合评测格式的结果
    pipeline_config_path = "paddlex/configs/pipelines/PP-StructureV3.yaml"
    """
    
    def __init__(self, pipeline_config_path: str = "PP-StructureV3"):
        """
        初始化评估器
        
        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)
        if image is None:
            print(f"无法读取图像: {image_path}")
            return None
            
        height, width = image.shape[:2]
        
        # 运行PaddleX pipeline
        start_time = time.time()
        
        try:
            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
            ))
        except Exception as e:
            print(f"处理图像 {image_path} 时发生错误: {str(e)}")
            traceback.print_exc()
            return None
        
        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:
            res_json = res.json.get('res', {})
            # 从parsing_res_list中提取布局信息
            parsing_list = res_json.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格式 [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_json, 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结果，提取相关的文本行
        ocr_res = res.get('overall_ocr_res', None)
        if ocr_res:
            texts = ocr_res.get('rec_texts', [])
            boxes = ocr_res.get('rec_boxes', [])
            scores = ocr_res.get('rec_scores', [1.0] * len(texts)) if 'rec_scores' in ocr_res else [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格式的结果字典
        """
        if not Path(result_path).exists():
            print(f"结果文件不存在: {result_path}")
            return None
            
        try:
            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)
                if image is None:
                    print(f"无法读取图像: {input_path}")
                    height, width = 1600, 1200
                    image_name = "unknown.png"
                else:
                    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
            
        except Exception as e:
            print(f"加载结果文件 {result_path} 时发生错误: {str(e)}")
            traceback.print_exc()
            return None
    
    def load_omnidocbench_dataset(self, dataset_dir: str) -> Dict[str, Any]:
        """
        加载OmniDocBench数据集
        
        Args:
            dataset_dir: 数据集目录路径
            
        Returns:
            加载的数据集字典
        """
        dataset = {}
        
        # 遍历数据集目录
        for file_path in Path(dataset_dir).rglob('*.json'):
            if "pred" in file_path.name or "result" in file_path.name:
                continue
                
            try:
                with open(file_path, 'r', encoding='utf-8') as f:
                    data = json.load(f)
                
                # 提取文件ID
                file_id = file_path.stem
                dataset[file_id] = {
                    "ground_truth": data,
                    "predictions": None,
                    "image_path": self._find_image_file(file_path.parent, file_id)
                }
                
            except Exception as e:
                print(f"加载文件 {file_path} 出错: {str(e)}")
        
        return dataset
    
    def _find_image_file(self, search_dir: Path, file_id: str) -> str:
        """查找对应的图像文件"""
        image_extensions = ['.jpg', '.jpeg', '.png', '.bmp', '.tiff']
        
        for ext in image_extensions:
            image_path = search_dir / f"{file_id}{ext}"
            if image_path.exists():
                return str(image_path)
        
        # 如果找不到图像文件，尝试在子目录中查找
        for subdir in search_dir.iterdir():
            if subdir.is_dir():
                result = self._find_image_file(subdir, file_id)
                if result:
                    return result
        
        return ""
    
    def generate_predictions(self, dataset: Dict[str, Any]) -> Dict[str, Any]:
        """
        为数据集生成预测
        
        Args:
            dataset: 数据集字典
            
        Returns:
            包含预测结果的数据集字典
        """
        for file_id, item in dataset.items():
            if item["image_path"] and Path(item["image_path"]).exists():
                try:
                    result = self.evaluate_single_image(item["image_path"])
                    dataset[file_id]["predictions"] = result
                    print(f"成功处理文件: {file_id}")
                except Exception as e:
                    print(f"处理文件 {file_id} 出错: {str(e)}")
            else:
                print(f"图像文件不存在: {item['image_path']}")
        
        return dataset
    
    def evaluate_dataset(self, dataset: Dict[str, Any]) -> Dict[str, float]:
        """
        评估数据集的预测结果
        
        Args:
            dataset: 包含预测和真实标签的数据集字典
            
        Returns:
            包含评估指标的字典
        """
        metrics = {
            "precision": 0.0,
            "recall": 0.0,
            "f1_score": 0.0,
            "iou": 0.0
        }
        
        # 实现具体的评估逻辑
        # 这里只是一个示例框架
        total_precision = 0.0
        total_recall = 0.0
        total_f1 = 0.0
        total_iou = 0.0
        count = 0
        
        for file_id, item in dataset.items():
            if item["predictions"] is None:
                continue
                
            # 计算单个样本的评估指标
            sample_metrics = self._calculate_sample_metrics(item["ground_truth"], item["predictions"])
            
            total_precision += sample_metrics["precision"]
            total_recall += sample_metrics["recall"]
            total_f1 += sample_metrics["f1_score"]
            total_iou += sample_metrics["iou"]
            count += 1
        
        if count > 0:
            metrics["precision"] = total_precision / count
            metrics["recall"] = total_recall / count
            metrics["f1_score"] = total_f1 / count
            metrics["iou"] = total_iou / count
        
        return metrics
    
    def _calculate_sample_metrics(self, ground_truth: Dict, prediction: Dict) -> Dict[str, float]:
        """
        计算单个样本的评估指标
        
        Args:
            ground_truth: 真实标签数据
            prediction: 预测结果
            
        Returns:
            包含评估指标的字典
        """
        metrics = {
            "precision": 0.0,
            "recall": 0.0,
            "f1_score": 0.0,
            "iou": 0.0
        }
        
        # 实现具体的评估逻辑
        # 这里只是一个示例框架
        gt_layouts = ground_truth.get("layout_dets", [])
        pred_layouts = prediction.get("layout_dets", [])
        
        # 简单的类别匹配计算
        gt_categories = set(item.get("category_type", "") for item in gt_layouts)
        pred_categories = set(item.get("category_type", "") for item in pred_layouts)
        
        # 计算交并集
        intersection = len(gt_categories.intersection(pred_categories))
        union = len(gt_categories.union(pred_categories))
        
        if union > 0:
            metrics["iou"] = intersection / union
        
        # 计算精确度、召回率、F1分数
        if len(pred_categories) > 0:
            metrics["precision"] = intersection / len(pred_categories)
        
        if len(gt_categories) > 0:
            metrics["recall"] = intersection / len(gt_categories)
        
        if metrics["precision"] + metrics["recall"] > 0:
            metrics["f1_score"] = 2 * (metrics["precision"] * metrics["recall"]) / (metrics["precision"] + metrics["recall"])
        
        return metrics
    
    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结果"""
    try:
        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)
            if omnidocbench_result is None:
                print(f"转换结果为空: {result_file}")
                return
                
            # 保存转换后的结果
            output_file = "./omnidocbench_converted_result.json"
            try:
                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]}...")
            except Exception as e:
                print(f"保存结果到文件 {output_file} 时发生错误: {str(e)}")
                traceback.print_exc()
        else:
            print(f"结果文件不存在: {result_file}")
            
    except Exception as e:
        print(f"转换过程中发生致命错误: {str(e)}")
        traceback.print_exc()

def process_omnidocbench_dataset():
    """处理OmniDocBench数据集"""
    try:
        # 初始化评估器
        evaluator = OmniDocBenchEvaluator()
        
        # 数据集路径
        dataset_path = "/home/ubuntu/zhch/OmniDocBench/OpenDataLab___OmniDocBench"
        result_dir = "/home/ubuntu/zhch/PaddleX/zhch/OmniDocBench_Result"
        
        # 确保结果目录存在
        os.makedirs(result_dir, exist_ok=True)
        
        # 查找所有图像文件
        image_extensions = ['*.jpg', '*.jpeg', '*.png', '*.bmp', '*.tiff']
        image_files = []
        
        for ext in image_extensions:
            image_files.extend(glob.glob(os.path.join(dataset_path, '**', ext), recursive=True))
        
        print(f"找到 {len(image_files)} 个图像文件")
        
        if not image_files:
            print("未找到任何图像文件，程序终止")
            return
        
        # 存储所有结果
        all_results = []
        
        # 处理每个图像
        for i, image_path in enumerate(image_files[:10]):  # 限制处理前10个文件用于测试
            try:
                print(f"处理进度: {i+1}/{len(image_files[:10])}")
                
                # 处理单个图像
                result = evaluator.evaluate_single_image(image_path)
                if result is not None:
                    all_results.append(result)
                
            except Exception as e:
                print(f"处理文件 {image_path} 时出错: {str(e)}")
                traceback.print_exc()
                continue
        
        # 保存结果
        output_file = os.path.join(result_dir, "OmniDocBench-PPStructureV3.json")
        try:
            with open(output_file, 'w', encoding='utf-8') as f:
                json.dump(all_results, f, ensure_ascii=False, indent=2)
            
            print(f"处理完成，结果保存至: {output_file}")
            print(f"共处理 {len(all_results)} 个文件")
        except Exception as e:
            print(f"保存结果文件时发生错误: {str(e)}")
            traceback.print_exc()
            
    except Exception as e:
        print(f"处理数据集时发生致命错误: {str(e)}")
        traceback.print_exc()


if __name__ == "__main__":
    # convert_existing_results()
    process_omnidocbench_dataset()