PaddleX/zhch/omnidocbench_parallel_eval.py

# zhch/omnidocbench_parallel_eval.py
import json
import time
import os
import glob
import traceback
from pathlib import Path
from typing import List, Dict, Any, Tuple
from concurrent.futures import ThreadPoolExecutor, ProcessPoolExecutor, as_completed
from multiprocessing import Queue, Manager
import cv2
import numpy as np
from paddlex import create_pipeline
from tqdm import tqdm
import threading

class OmniDocBenchParallelEvaluator:
    """
    OmniDocBench并行评估器，支持多进程批处理
    """
    
    def __init__(self, pipeline_config_path: str = "PP-StructureV3"):
        """
        初始化评估器
        
        Args:
            pipeline_config_path: PaddleX pipeline配置文件路径
        """
        self.pipeline_config = pipeline_config_path
        self.category_mapping = self._get_category_mapping()
        
    def _get_category_mapping(self) -> Dict[str, str]:
        """获取PaddleX类别到OmniDocBench类别的映射"""
        return {
            '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 create_pipeline(self):
        """创建pipeline实例（每个进程单独创建）"""
        return create_pipeline(pipeline=self.pipeline_config)
    
    def process_single_image(self, image_path: str, use_gpu: bool = True) -> Dict[str, Any]:
        """
        处理单张图像
        
        Args:
            image_path: 图像路径
            use_gpu: 是否使用GPU
            
        Returns:
            OmniDocBench格式的结果字典
        """
        try:
            # 每个进程创建自己的pipeline
            pipeline = self.create_pipeline()
            
            # 读取图像获取尺寸信息
            image = cv2.imread(image_path)
            if image is None:
                return None
                
            height, width = image.shape[:2]
            
            # 运行PaddleX pipeline
            start_time = time.time()
            
            output = list(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,
            ))
            
            process_time = time.time() - start_time
            
            # 转换为OmniDocBench格式
            result = self._convert_to_omnidocbench_format(
                output, image_path, width, height
            )
            
            # 添加处理时间信息
            if result:
                result["processing_time"] = process_time
                result["success"] = True
            
            return result
            
        except Exception as e:
            return {
                "image_path": Path(image_path).name,
                "error": str(e),
                "success": False,
                "processing_time": 0
            }
    
    def process_batch(self, image_paths: List[str], use_gpu: bool = True) -> List[Dict[str, Any]]:
        """
        批处理图像
        
        Args:
            image_paths: 图像路径列表
            use_gpu: 是否使用GPU
            
        Returns:
            结果列表
        """
        results = []
        pipeline = self.create_pipeline()
        
        for image_path in image_paths:
            try:
                result = self._process_with_pipeline(pipeline, image_path, use_gpu)
                if result:
                    results.append(result)
            except Exception as e:
                error_result = {
                    "image_path": Path(image_path).name,
                    "error": str(e),
                    "success": False,
                    "processing_time": 0
                }
                results.append(error_result)
        
        return results
    
    def _process_with_pipeline(self, pipeline, image_path: str, use_gpu: bool) -> Dict[str, Any]:
        """使用给定的pipeline处理图像"""
        # 读取图像获取尺寸信息
        image = cv2.imread(image_path)
        if image is None:
            return None
            
        height, width = image.shape[:2]
        
        # 运行pipeline
        start_time = time.time()
        
        output = list(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,
        ))
        
        process_time = time.time() - start_time
        
        # 转换格式
        result = self._convert_to_omnidocbench_format(
            output, image_path, width, height
        )
        
        if result:
            result["processing_time"] = process_time
            result["success"] = True
        
        return result
    
    def _convert_to_omnidocbench_format(self, 
                                      paddlex_output: List, 
                                      image_path: str,
                                      width: int, 
                                      height: int) -> Dict[str, Any]:
        """将PaddleX输出转换为OmniDocBench格式"""
        layout_dets = []
        anno_id_counter = 0
        
        # 处理PaddleX的输出
        for res in paddlex_output:
            res_json = res.json.get('res', {})
            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格式
                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": Path(image_path).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 _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 parallel_process_with_threading(self, 
                                      image_paths: List[str], 
                                      batch_size: int = 4,
                                      max_workers: int = 4,
                                      use_gpu: bool = True) -> List[Dict[str, Any]]:
        """
        使用多线程并行处理（推荐用于GPU）
        
        Args:
            image_paths: 图像路径列表
            batch_size: 批处理大小
            max_workers: 最大工作线程数
            use_gpu: 是否使用GPU
            
        Returns:
            处理结果列表
        """
        # 将图像路径分批
        batches = [image_paths[i:i + batch_size] for i in range(0, len(image_paths), batch_size)]
        
        all_results = []
        completed_count = 0
        total_images = len(image_paths)
        
        # 创建进度条
        with tqdm(total=total_images, desc="处理图像", unit="张") as pbar:
            with ThreadPoolExecutor(max_workers=max_workers) as executor:
                # 提交所有批处理任务
                future_to_batch = {
                    executor.submit(self.process_batch, batch, use_gpu): batch 
                    for batch in batches
                }
                
                # 收集结果
                for future in as_completed(future_to_batch):
                    batch = future_to_batch[future]
                    try:
                        batch_results = future.result()
                        all_results.extend(batch_results)
                        completed_count += len(batch)
                        pbar.update(len(batch))
                        
                        # 更新进度条描述
                        success_count = sum(1 for r in batch_results if r.get('success', False))
                        pbar.set_postfix({
                            'batch_success': f"{success_count}/{len(batch)}",
                            'total_success': f"{sum(1 for r in all_results if r.get('success', False))}/{completed_count}"
                        })
                        
                    except Exception as e:
                        print(f"批处理失败: {e}")
                        # 为失败的批次创建错误结果
                        for img_path in batch:
                            error_result = {
                                "image_path": Path(img_path).name,
                                "error": str(e),
                                "success": False,
                                "processing_time": 0
                            }
                            all_results.append(error_result)
                        pbar.update(len(batch))
        
        return all_results
    
    def parallel_process_with_multiprocessing(self, 
                                            image_paths: List[str], 
                                            batch_size: int = 4,
                                            max_workers: int = 4,
                                            use_gpu: bool = False) -> List[Dict[str, Any]]:
        """
        使用多进程并行处理（推荐用于CPU）
        
        Args:
            image_paths: 图像路径列表
            batch_size: 批处理大小
            max_workers: 最大工作进程数
            use_gpu: 是否使用GPU
            
        Returns:
            处理结果列表
        """
        # 将图像路径分批
        batches = [image_paths[i:i + batch_size] for i in range(0, len(image_paths), batch_size)]
        
        all_results = []
        completed_count = 0
        total_images = len(image_paths)
        
        # 创建进度条
        with tqdm(total=total_images, desc="处理图像", unit="张") as pbar:
            with ProcessPoolExecutor(max_workers=max_workers) as executor:
                # 提交所有批处理任务
                future_to_batch = {
                    executor.submit(process_batch_worker, batch, self.pipeline_config, use_gpu): batch 
                    for batch in batches
                }
                
                # 收集结果
                for future in as_completed(future_to_batch):
                    batch = future_to_batch[future]
                    try:
                        batch_results = future.result()
                        all_results.extend(batch_results)
                        completed_count += len(batch)
                        pbar.update(len(batch))
                        
                        # 更新进度条描述
                        success_count = sum(1 for r in batch_results if r.get('success', False))
                        pbar.set_postfix({
                            'batch_success': f"{success_count}/{len(batch)}",
                            'total_success': f"{sum(1 for r in all_results if r.get('success', False))}/{completed_count}"
                        })
                        
                    except Exception as e:
                        print(f"批处理失败: {e}")
                        # 为失败的批次创建错误结果
                        for img_path in batch:
                            error_result = {
                                "image_path": Path(img_path).name,
                                "error": str(e),
                                "success": False,
                                "processing_time": 0
                            }
                            all_results.append(error_result)
                        pbar.update(len(batch))
        
        return all_results
    
    def save_results_incrementally(self, 
                                 results: List[Dict[str, Any]], 
                                 output_file: str,
                                 save_interval: int = 50):
        """
        增量保存结果
        
        Args:
            results: 结果列表
            output_file: 输出文件路径
            save_interval: 保存间隔
        """
        if len(results) % save_interval == 0 and len(results) > 0:
            try:
                with open(output_file, 'w', encoding='utf-8') as f:
                    json.dump(results, f, ensure_ascii=False, indent=2)
                print(f"已保存 {len(results)} 个结果到 {output_file}")
            except Exception as e:
                print(f"保存结果时出错: {e}")

def process_batch_worker(image_paths: List[str], pipeline_config: str, use_gpu: bool) -> List[Dict[str, Any]]:
    """
    多进程工作函数
    """
    try:
        # 在每个进程中创建评估器
        evaluator = OmniDocBenchParallelEvaluator(pipeline_config)
        return evaluator.process_batch(image_paths, use_gpu)
    except Exception as e:
        # 返回错误结果
        error_results = []
        for img_path in image_paths:
            error_results.append({
                "image_path": Path(img_path).name,
                "error": str(e),
                "success": False,
                "processing_time": 0
            })
        return error_results

def main():
    """主函数 - 并行处理OmniDocBench数据集"""
    
    # 配置参数
    dataset_path = "/Users/zhch158/workspace/repository.git/OmniDocBench/OpenDataLab___OmniDocBench/images"
    output_dir = "/Users/zhch158/workspace/repository.git/PaddleX/zhch/OmniDocBench_Results"
    pipeline_config = "PP-StructureV3"
    
    # 并行处理参数
    batch_size = 4          # 批处理大小
    max_workers = 4         # 最大工作进程/线程数
    use_gpu = True          # 是否使用GPU
    use_multiprocessing = False  # False=多线程(GPU推荐), True=多进程(CPU推荐)
    
    # 确保输出目录存在
    os.makedirs(output_dir, exist_ok=True)
    
    print("="*60)
    print("OmniDocBench 并行评估开始")
    print("="*60)
    print(f"数据集路径: {dataset_path}")
    print(f"输出目录: {output_dir}")
    print(f"批处理大小: {batch_size}")
    print(f"最大工作线程/进程数: {max_workers}")
    print(f"使用GPU: {use_gpu}")
    print(f"并行方式: {'多进程' if use_multiprocessing else '多线程'}")
    
    # 查找所有图像文件
    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)))
    
    print(f"找到 {len(image_files)} 个图像文件")
    
    if not image_files:
        print("未找到任何图像文件，程序终止")
        return
    
    # 创建评估器
    evaluator = OmniDocBenchParallelEvaluator(pipeline_config)
    
    # 开始处理
    start_time = time.time()
    
    if use_multiprocessing:
        # 多进程处理（推荐用于CPU）
        print("使用多进程并行处理...")
        results = evaluator.parallel_process_with_multiprocessing(
            image_files, batch_size, max_workers, use_gpu
        )
    else:
        # 多线程处理（推荐用于GPU）
        print("使用多线程并行处理...")
        results = evaluator.parallel_process_with_threading(
            image_files, batch_size, max_workers, use_gpu
        )
    
    total_time = time.time() - start_time
    
    # 保存最终结果
    output_file = os.path.join(output_dir, f"OmniDocBench_PPStructureV3_batch{batch_size}.json")
    try:
        with open(output_file, 'w', encoding='utf-8') as f:
            json.dump(results, f, ensure_ascii=False, indent=2)
        
        print("\n" + "="*60)
        print("处理完成!")
        print("="*60)
        
        # 统计信息
        success_count = sum(1 for r in results if r.get('success', False))
        error_count = len(results) - success_count
        total_processing_time = sum(r.get('processing_time', 0) for r in results if r.get('success', False))
        avg_processing_time = total_processing_time / success_count if success_count > 0 else 0
        
        print(f"总文件数: {len(image_files)}")
        print(f"成功处理: {success_count}")
        print(f"失败数量: {error_count}")
        print(f"成功率: {success_count / len(image_files) * 100:.2f}%")
        print(f"总耗时: {total_time:.2f}秒")
        print(f"平均处理时间: {avg_processing_time:.2f}秒/张")
        print(f"吞吐量: {len(image_files) / total_time:.2f}张/秒")
        print(f"结果保存至: {output_file}")
        
        # 保存统计信息
        stats = {
            "total_files": len(image_files),
            "success_count": success_count,
            "error_count": error_count,
            "success_rate": success_count / len(image_files),
            "total_time": total_time,
            "avg_processing_time": avg_processing_time,
            "throughput": len(image_files) / total_time,
            "batch_size": batch_size,
            "max_workers": max_workers,
            "use_gpu": use_gpu,
            "use_multiprocessing": use_multiprocessing
        }
        
        stats_file = os.path.join(output_dir, f"processing_stats_batch{batch_size}.json")
        with open(stats_file, 'w', encoding='utf-8') as f:
            json.dump(stats, f, ensure_ascii=False, indent=2)
        
        print(f"统计信息保存至: {stats_file}")
        
    except Exception as e:
        print(f"保存结果文件时发生错误: {str(e)}")
        traceback.print_exc()

if __name__ == "__main__":
    main()