# 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 PPStructureV3ParallelPredictor: """ PP-StructureV3并行预测器，支持多进程批处理 """ def __init__(self, pipeline_config_path: str = "PP-StructureV3", output_path: str = "output", use_gpu: bool = True): """ 初始化预测器 Args: pipeline_config_path: PaddleX pipeline配置文件路径 """ self.pipeline_config = pipeline_config_path self.pipeline = create_pipeline(pipeline=self.pipeline_config) self.output_path = output_path self.use_gpu = use_gpu def create_pipeline(self): """创建pipeline实例（每个进程单独创建）""" if self.pipeline is not None: return self.pipeline return create_pipeline(pipeline=self.pipeline_config) def process_single_image(self, image_path: str) -> Dict[str, Any]: """ 处理单张图像 Args: image_path: 图像路径 output_path: 输出路径 use_gpu: 是否使用GPU Returns: 处理结果{"image_path": str, "success": bool, "processing_time": float, "error": str} """ try: # 读取图像获取尺寸信息 image = cv2.imread(image_path) if image is None: return { "image_path": Path(image_path).name, "error": "无法读取图像", "success": False, "processing_time": 0 } height, width = image.shape[:2] # 运行PaddleX pipeline start_time = time.time() output = self.pipeline.predict( input=image_path, device="gpu" if self.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, ) # 可视化结果并保存 json 结果 for res in output: res.save_to_json(save_path=self.output_path) # 保存所有结果到指定路径 res.save_to_markdown(save_path=self.output_path) # 保存所有结果到指定路径 process_time = time.time() - start_time # 添加处理时间信息 result = {"image_path": Path(image_path).name} if output: 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]) -> List[Dict[str, Any]]: """ 批处理图像 Args: image_paths: 图像路径列表 use_gpu: 是否使用GPU Returns: 结果列表 """ results = [] for image_path in image_paths: result = self.process_single_image(image_path=image_path) results.append(result) return results def parallel_process_with_threading(self, image_paths: List[str], batch_size: int = 4, max_workers: int = 4 ) -> List[Dict[str, Any]]: """ 使用多线程并行处理（推荐用于GPU） Args: image_paths: 图像路径列表 batch_size: 批处理大小 max_workers: 最大工作线程数 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): 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, output_path: str, use_gpu: bool) -> List[Dict[str, Any]]: """ 多进程工作函数 """ try: # 在每个进程中创建pipeline实例 predictor = PPStructureV3ParallelPredictor(pipeline_config, output_path=output_path, use_gpu=use_gpu) return predictor.process_batch(image_paths) 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 parallel_process_with_multiprocessing(image_paths: List[str], batch_size: int = 4, max_workers: int = 4, pipeline_config: str = "PP-StructureV3", output_path: str = "./output", use_gpu: bool = True ) -> 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, pipeline_config, output_path, 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 main(): """主函数 - 并行处理OmniDocBench数据集""" # 配置参数 dataset_path = "../../OmniDocBench/OpenDataLab___OmniDocBench/images" output_dir = "./OmniDocBench_Results" pipeline_config = "PP-StructureV3" # 并行处理参数 batch_size = 4 # 批处理大小 max_workers = 4 # 最大工作进程/线程数 use_gpu = True # 是否使用GPU use_multiprocessing = True # False=多线程(GPU推荐), True=多进程(CPU推荐) # 确保输出目录存在 print(f"输出目录: {Path(output_dir).absolute()}") os.makedirs(output_dir, exist_ok=True) dataset_path = Path(dataset_path).resolve() output_dir = Path(output_dir).resolve() 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 # 开始处理 start_time = time.time() if use_multiprocessing: # 多进程处理（推荐用于CPU） print("使用多进程并行处理...") results = parallel_process_with_multiprocessing( image_files, batch_size, max_workers, pipeline_config, output_dir, use_gpu ) else: # 多线程处理（推荐用于GPU） print("使用多线程并行处理...") predictor = PPStructureV3ParallelPredictor(pipeline_config, output_path=output_dir, use_gpu=use_gpu) results = predictor.parallel_process_with_threading( image_files, batch_size, max_workers ) total_time = time.time() - start_time # 保存最终结果 output_file = os.path.join(output_dir, f"OmniDocBench_PPStructureV3_batch{batch_size}.json") try: # 统计信息 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 } results['stats'] = stats 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) except Exception as e: print(f"保存结果文件时发生错误: {str(e)}") traceback.print_exc() if __name__ == "__main__": main()