OmniDocBench_DotsOCR_multthreads.py 28 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743
  1. """
  2. 批量处理 OmniDocBench 图片并生成符合评测要求的预测结果
  3. 根据 OmniDocBench 评测要求:
  4. - 输入:OpenDataLab___OmniDocBench/images 下的所有 .jpg 图片,以及PDF文件
  5. - 输出:每个图片对应的 .md、.json 和带标注的 layout 图片文件
  6. - 输出目录:用于后续的 end2end 评测
  7. """
  8. import os
  9. import sys
  10. import json
  11. import tempfile
  12. import uuid
  13. import shutil
  14. import time
  15. import traceback
  16. import warnings
  17. from pathlib import Path
  18. from typing import List, Dict, Any
  19. from PIL import Image
  20. from tqdm import tqdm
  21. import argparse
  22. # 导入 dots.ocr 相关模块
  23. from dots_ocr.parser import DotsOCRParser
  24. from dots_ocr.utils import dict_promptmode_to_prompt
  25. from dots_ocr.utils.consts import MIN_PIXELS, MAX_PIXELS
  26. from dots_ocr.utils.doc_utils import load_images_from_pdf
  27. # 导入工具函数
  28. from utils import (
  29. get_image_files_from_dir,
  30. get_image_files_from_list,
  31. get_image_files_from_csv,
  32. collect_pid_files
  33. )
  34. def convert_pdf_to_images(pdf_file: str, output_dir: str | None = None, dpi: int = 200) -> List[str]:
  35. """
  36. 将PDF转换为图像文件
  37. Args:
  38. pdf_file: PDF文件路径
  39. output_dir: 输出目录
  40. dpi: 图像分辨率
  41. Returns:
  42. 生成的图像文件路径列表
  43. """
  44. pdf_path = Path(pdf_file)
  45. if not pdf_path.exists() or pdf_path.suffix.lower() != '.pdf':
  46. print(f"❌ Invalid PDF file: {pdf_path}")
  47. return []
  48. # 如果没有指定输出目录,使用PDF同名目录
  49. if output_dir is None:
  50. output_path = pdf_path.parent / f"{pdf_path.stem}"
  51. else:
  52. output_path = Path(output_dir) / f"{pdf_path.stem}"
  53. output_path = output_path.resolve()
  54. output_path.mkdir(parents=True, exist_ok=True)
  55. try:
  56. # 使用utils中的函数加载PDF图像
  57. images = load_images_from_pdf(str(pdf_path), dpi=dpi)
  58. image_paths = []
  59. for i, image in enumerate(images):
  60. # 生成图像文件名
  61. image_filename = f"{pdf_path.stem}_page_{i+1:03d}.png"
  62. image_path = output_path / image_filename
  63. # 保存图像
  64. image.save(str(image_path))
  65. image_paths.append(str(image_path))
  66. print(f"✅ Converted {len(images)} pages from {pdf_path.name} to images")
  67. return image_paths
  68. except Exception as e:
  69. print(f"❌ Error converting PDF {pdf_path}: {e}")
  70. traceback.print_exc()
  71. return []
  72. def get_input_files(args) -> List[str]:
  73. """
  74. 获取输入文件列表,统一处理PDF和图像文件
  75. Args:
  76. args: 命令行参数
  77. Returns:
  78. 处理后的图像文件路径列表
  79. """
  80. input_files = []
  81. # 获取原始输入文件
  82. if args.input_csv:
  83. raw_files = get_image_files_from_csv(args.input_csv, "fail")
  84. elif args.input_file_list:
  85. raw_files = get_image_files_from_list(args.input_file_list)
  86. elif args.input_file:
  87. raw_files = [Path(args.input_file).resolve()]
  88. else:
  89. input_dir = Path(args.input_dir).resolve()
  90. if not input_dir.exists():
  91. print(f"❌ Input directory does not exist: {input_dir}")
  92. return []
  93. # 获取所有支持的文件(图像和PDF)
  94. image_extensions = ['.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.tif']
  95. pdf_extensions = ['.pdf']
  96. raw_files = []
  97. for ext in image_extensions + pdf_extensions:
  98. raw_files.extend(list(input_dir.glob(f"*{ext}")))
  99. raw_files.extend(list(input_dir.glob(f"*{ext.upper()}")))
  100. raw_files = [str(f) for f in raw_files]
  101. # 分别处理PDF和图像文件
  102. pdf_count = 0
  103. image_count = 0
  104. for file_path in raw_files:
  105. file_path = Path(file_path)
  106. if file_path.suffix.lower() == '.pdf':
  107. # 转换PDF为图像
  108. print(f"📄 Processing PDF: {file_path.name}")
  109. pdf_images = convert_pdf_to_images(
  110. str(file_path),
  111. args.output_dir,
  112. dpi=args.dpi
  113. )
  114. input_files.extend(pdf_images)
  115. pdf_count += 1
  116. else:
  117. # 直接添加图像文件
  118. if file_path.exists():
  119. input_files.append(str(file_path))
  120. image_count += 1
  121. print(f"📊 Input summary:")
  122. print(f" PDF files processed: {pdf_count}")
  123. print(f" Image files found: {image_count}")
  124. print(f" Total image files to process: {len(input_files)}")
  125. return input_files
  126. class DotsOCRProcessor:
  127. """DotsOCR 处理器"""
  128. def __init__(self,
  129. ip: str = "127.0.0.1",
  130. port: int = 8101,
  131. model_name: str = "DotsOCR",
  132. prompt_mode: str = "prompt_layout_all_en",
  133. dpi: int = 200,
  134. min_pixels: int = MIN_PIXELS,
  135. max_pixels: int = MAX_PIXELS):
  136. """
  137. 初始化处理器
  138. Args:
  139. ip: vLLM 服务器 IP
  140. port: vLLM 服务器端口
  141. model_name: 模型名称
  142. prompt_mode: 提示模式
  143. dpi: PDF 处理 DPI
  144. min_pixels: 最小像素数
  145. max_pixels: 最大像素数
  146. """
  147. self.ip = ip
  148. self.port = port
  149. self.model_name = model_name
  150. self.prompt_mode = prompt_mode
  151. self.dpi = dpi
  152. self.min_pixels = min_pixels
  153. self.max_pixels = max_pixels
  154. # 初始化解析器
  155. self.parser = DotsOCRParser(
  156. ip=ip,
  157. port=port,
  158. dpi=dpi,
  159. min_pixels=min_pixels,
  160. max_pixels=max_pixels,
  161. model_name=model_name
  162. )
  163. print(f"DotsOCR Parser 初始化完成:")
  164. print(f" - 服务器: {ip}:{port}")
  165. print(f" - 模型: {model_name}")
  166. print(f" - 提示模式: {prompt_mode}")
  167. print(f" - 像素范围: {min_pixels} - {max_pixels}")
  168. def create_temp_session_dir(self) -> tuple:
  169. """创建临时会话目录"""
  170. session_id = uuid.uuid4().hex[:8]
  171. temp_dir = os.path.join(tempfile.gettempdir(), f"omnidocbench_batch_{session_id}")
  172. os.makedirs(temp_dir, exist_ok=True)
  173. return temp_dir, session_id
  174. def save_results_to_output_dir(self, result: Dict, image_name: str, output_dir: str) -> Dict[str, str]:
  175. """
  176. 将处理结果保存到输出目录
  177. Args:
  178. result: 解析结果
  179. image_name: 图片文件名(不含扩展名)
  180. output_dir: 输出目录
  181. Returns:
  182. dict: 保存的文件路径
  183. """
  184. saved_files = {}
  185. try:
  186. # 1. 保存 Markdown 文件(OmniDocBench 评测必需)
  187. output_md_path = os.path.join(output_dir, f"{image_name}.md")
  188. md_content = ""
  189. # 优先使用无页眉页脚的版本(符合 OmniDocBench 评测要求)
  190. if 'md_content_nohf_path' in result and os.path.exists(result['md_content_nohf_path']):
  191. with open(result['md_content_nohf_path'], 'r', encoding='utf-8') as f:
  192. md_content = f.read()
  193. elif 'md_content_path' in result and os.path.exists(result['md_content_path']):
  194. with open(result['md_content_path'], 'r', encoding='utf-8') as f:
  195. md_content = f.read()
  196. else:
  197. md_content = "# 解析失败\n\n未能提取到有效的文档内容。"
  198. with open(output_md_path, 'w', encoding='utf-8') as f:
  199. f.write(md_content)
  200. saved_files['md'] = output_md_path
  201. # 2. 保存 JSON 文件
  202. output_json_path = os.path.join(output_dir, f"{image_name}.json")
  203. json_data = {}
  204. if 'layout_info_path' in result and os.path.exists(result['layout_info_path']):
  205. with open(result['layout_info_path'], 'r', encoding='utf-8') as f:
  206. json_data = json.load(f)
  207. else:
  208. json_data = {
  209. "error": "未能提取到有效的布局信息",
  210. "cells": []
  211. }
  212. with open(output_json_path, 'w', encoding='utf-8') as f:
  213. json.dump(json_data, f, ensure_ascii=False, indent=2)
  214. saved_files['json'] = output_json_path
  215. # 3. 保存带标注的布局图片
  216. output_layout_image_path = os.path.join(output_dir, f"{image_name}_layout.jpg")
  217. if 'layout_image_path' in result and os.path.exists(result['layout_image_path']):
  218. # 直接复制布局图片
  219. shutil.copy2(result['layout_image_path'], output_layout_image_path)
  220. saved_files['layout_image'] = output_layout_image_path
  221. else:
  222. # 如果没有布局图片,使用原始图片作为占位符
  223. try:
  224. original_image = Image.open(result.get('original_image_path', ''))
  225. original_image.save(output_layout_image_path, 'JPEG', quality=95)
  226. saved_files['layout_image'] = output_layout_image_path
  227. except Exception as e:
  228. saved_files['layout_image'] = None
  229. except Exception as e:
  230. print(f"Error saving results for {image_name}: {e}")
  231. return saved_files
  232. def process_single_image(self, image_path: str, output_dir: str) -> Dict[str, Any]:
  233. """
  234. 处理单张图片
  235. Args:
  236. image_path: 图片路径
  237. output_dir: 输出目录
  238. Returns:
  239. dict: 处理结果
  240. """
  241. start_time = time.time()
  242. image_name = Path(image_path).stem
  243. result_info = {
  244. "image_path": image_path,
  245. "processing_time": 0,
  246. "success": False,
  247. "device": f"{self.ip}:{self.port}",
  248. "error": None,
  249. "output_files": {},
  250. "is_pdf_page": "_page_" in Path(image_path).name # 标记是否为PDF页面
  251. }
  252. try:
  253. # 检查输出文件是否已存在
  254. output_md_path = os.path.join(output_dir, f"{image_name}.md")
  255. output_json_path = os.path.join(output_dir, f"{image_name}.json")
  256. output_layout_path = os.path.join(output_dir, f"{image_name}_layout.jpg")
  257. if all(os.path.exists(p) for p in [output_md_path, output_json_path, output_layout_path]):
  258. result_info.update({
  259. "success": True,
  260. "processing_time": 0,
  261. "output_files": {
  262. "md": output_md_path,
  263. "json": output_json_path,
  264. "layout_image": output_layout_path
  265. },
  266. "skipped": True
  267. })
  268. return result_info
  269. # 创建临时会话目录
  270. temp_dir, session_id = self.create_temp_session_dir()
  271. try:
  272. # 读取图片
  273. image = Image.open(image_path)
  274. # 使用 DotsOCRParser 处理图片
  275. filename = f"omnidocbench_{session_id}"
  276. results = self.parser.parse_image(
  277. input_path=image,
  278. filename=filename,
  279. prompt_mode=self.prompt_mode,
  280. save_dir=temp_dir,
  281. fitz_preprocess=True # 对图片使用 fitz 预处理
  282. )
  283. # 解析结果
  284. if not results:
  285. raise Exception("未返回解析结果")
  286. result = results[0] # parse_image 返回单个结果的列表
  287. # 保存所有结果文件到输出目录
  288. saved_files = self.save_results_to_output_dir(result, image_name, output_dir)
  289. # 验证保存结果
  290. success_count = sum(1 for path in saved_files.values() if path and os.path.exists(path))
  291. if success_count >= 2: # 至少保存了 md 和 json
  292. result_info.update({
  293. "success": True,
  294. "output_files": saved_files
  295. })
  296. else:
  297. raise Exception(f"保存文件不完整 ({success_count}/3)")
  298. finally:
  299. # 清理临时目录
  300. if os.path.exists(temp_dir):
  301. shutil.rmtree(temp_dir, ignore_errors=True)
  302. except Exception as e:
  303. result_info["error"] = str(e)
  304. print(f"❌ Error processing {image_name}: {e}")
  305. finally:
  306. result_info["processing_time"] = time.time() - start_time
  307. return result_info
  308. def process_images_single_process(image_paths: List[str],
  309. processor: DotsOCRProcessor,
  310. batch_size: int = 1,
  311. output_dir: str = "./output") -> List[Dict[str, Any]]:
  312. """
  313. 单进程版本的图像处理函数
  314. Args:
  315. image_paths: 图像路径列表
  316. processor: DotsOCR处理器实例
  317. batch_size: 批处理大小
  318. output_dir: 输出目录
  319. Returns:
  320. 处理结果列表
  321. """
  322. # 创建输出目录
  323. output_path = Path(output_dir)
  324. output_path.mkdir(parents=True, exist_ok=True)
  325. all_results = []
  326. total_images = len(image_paths)
  327. print(f"Processing {total_images} images with batch size {batch_size}")
  328. # 使用tqdm显示进度,添加更多统计信息
  329. with tqdm(total=total_images, desc="Processing images", unit="img",
  330. bar_format='{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}]') as pbar:
  331. # 按批次处理图像(DotsOCR通常单张处理)
  332. for i in range(0, total_images, batch_size):
  333. batch = image_paths[i:i + batch_size]
  334. batch_start_time = time.time()
  335. batch_results = []
  336. try:
  337. # 处理批次中的每张图片
  338. for image_path in batch:
  339. try:
  340. result = processor.process_single_image(image_path, output_dir)
  341. batch_results.append(result)
  342. except Exception as e:
  343. print(f"Error processing {image_path}: {e}", file=sys.stderr)
  344. traceback.print_exc()
  345. batch_results.append({
  346. "image_path": image_path,
  347. "processing_time": 0,
  348. "success": False,
  349. "device": f"{processor.ip}:{processor.port}",
  350. "error": str(e)
  351. })
  352. batch_processing_time = time.time() - batch_start_time
  353. all_results.extend(batch_results)
  354. # 更新进度条
  355. success_count = sum(1 for r in batch_results if r.get('success', False))
  356. skipped_count = sum(1 for r in batch_results if r.get('skipped', False))
  357. total_success = sum(1 for r in all_results if r.get('success', False))
  358. total_skipped = sum(1 for r in all_results if r.get('skipped', False))
  359. avg_time = batch_processing_time / len(batch)
  360. pbar.update(len(batch))
  361. pbar.set_postfix({
  362. 'batch_time': f"{batch_processing_time:.2f}s",
  363. 'avg_time': f"{avg_time:.2f}s/img",
  364. 'success': f"{total_success}/{len(all_results)}",
  365. 'skipped': f"{total_skipped}",
  366. 'rate': f"{total_success/len(all_results)*100:.1f}%"
  367. })
  368. except Exception as e:
  369. print(f"Error processing batch {[Path(p).name for p in batch]}: {e}", file=sys.stderr)
  370. traceback.print_exc()
  371. # 为批次中的所有图像添加错误结果
  372. error_results = []
  373. for img_path in batch:
  374. error_results.append({
  375. "image_path": str(img_path),
  376. "processing_time": 0,
  377. "success": False,
  378. "device": f"{processor.ip}:{processor.port}",
  379. "error": str(e)
  380. })
  381. all_results.extend(error_results)
  382. pbar.update(len(batch))
  383. return all_results
  384. def process_images_concurrent(image_paths: List[str],
  385. processor: DotsOCRProcessor,
  386. batch_size: int = 1,
  387. output_dir: str = "./output",
  388. max_workers: int = 3) -> List[Dict[str, Any]]:
  389. """并发版本的图像处理函数"""
  390. from concurrent.futures import ThreadPoolExecutor, as_completed
  391. Path(output_dir).mkdir(parents=True, exist_ok=True)
  392. def process_batch(batch_images):
  393. """处理一批图像"""
  394. batch_results = []
  395. for image_path in batch_images:
  396. try:
  397. result = processor.process_single_image(image_path, output_dir)
  398. batch_results.append(result)
  399. except Exception as e:
  400. batch_results.append({
  401. "image_path": image_path,
  402. "processing_time": 0,
  403. "success": False,
  404. "device": f"{processor.ip}:{processor.port}",
  405. "error": str(e)
  406. })
  407. return batch_results
  408. # 将图像分批
  409. batches = [image_paths[i:i + batch_size] for i in range(0, len(image_paths), batch_size)]
  410. all_results = []
  411. with ThreadPoolExecutor(max_workers=max_workers) as executor:
  412. # 提交所有批次
  413. future_to_batch = {executor.submit(process_batch, batch): batch for batch in batches}
  414. # 使用 tqdm 显示进度
  415. with tqdm(total=len(image_paths), desc="Processing images") as pbar:
  416. for future in as_completed(future_to_batch):
  417. try:
  418. batch_results = future.result()
  419. all_results.extend(batch_results)
  420. # 更新进度
  421. success_count = sum(1 for r in batch_results if r.get('success', False))
  422. pbar.update(len(batch_results))
  423. pbar.set_postfix({'batch_success': f"{success_count}/{len(batch_results)}"})
  424. except Exception as e:
  425. batch = future_to_batch[future]
  426. # 为批次中的所有图像添加错误结果
  427. error_results = [
  428. {
  429. "image_path": img_path,
  430. "processing_time": 0,
  431. "success": False,
  432. "device": f"{processor.ip}:{processor.port}",
  433. "error": str(e)
  434. }
  435. for img_path in batch
  436. ]
  437. all_results.extend(error_results)
  438. pbar.update(len(batch))
  439. return all_results
  440. def main():
  441. """主函数"""
  442. parser = argparse.ArgumentParser(description="DotsOCR OmniDocBench Processing with PDF Support")
  443. # 输入参数组
  444. input_group = parser.add_mutually_exclusive_group(required=True)
  445. input_group.add_argument("--input_file", type=str, help="Input file (supports both PDF and image file)")
  446. input_group.add_argument("--input_dir", type=str, help="Input directory (supports both PDF and image files)")
  447. input_group.add_argument("--input_file_list", type=str, help="Input file list (one file per line)")
  448. input_group.add_argument("--input_csv", type=str, help="Input CSV file with image_path and status columns")
  449. # 输出参数
  450. parser.add_argument("--output_dir", type=str, help="Output directory")
  451. # DotsOCR 参数
  452. parser.add_argument("--ip", type=str, default="127.0.0.1", help="vLLM server IP")
  453. parser.add_argument("--port", type=int, default=8101, help="vLLM server port")
  454. parser.add_argument("--model_name", type=str, default="DotsOCR", help="Model name")
  455. parser.add_argument("--prompt_mode", type=str, default="prompt_layout_all_en",
  456. choices=list(dict_promptmode_to_prompt.keys()), help="Prompt mode")
  457. parser.add_argument("--min_pixels", type=int, default=MIN_PIXELS, help="Minimum pixels")
  458. parser.add_argument("--max_pixels", type=int, default=MAX_PIXELS, help="Maximum pixels")
  459. parser.add_argument("--dpi", type=int, default=200, help="PDF processing DPI")
  460. # 处理参数
  461. parser.add_argument("--batch_size", type=int, default=1, help="Batch size")
  462. parser.add_argument("--input_pattern", type=str, default="*", help="Input file pattern")
  463. parser.add_argument("--test_mode", action="store_true", help="Test mode (process only 10 images)")
  464. parser.add_argument("--collect_results", type=str, help="收集处理结果到指定CSV文件")
  465. # 并发参数
  466. parser.add_argument("--max_workers", type=int, default=3,
  467. help="Maximum number of concurrent workers (should match vLLM data-parallel-size)")
  468. parser.add_argument("--use_threading", action="store_true",
  469. help="Use multi-threading")
  470. args = parser.parse_args()
  471. try:
  472. # 获取并预处理输入文件
  473. print("🔄 Preprocessing input files...")
  474. image_files = get_input_files(args)
  475. if not image_files:
  476. print("❌ No input files found or processed")
  477. return 1
  478. output_dir = Path(args.output_dir).resolve()
  479. print(f"📁 Output dir: {output_dir}")
  480. print(f"📊 Found {len(image_files)} image files to process")
  481. if args.test_mode:
  482. image_files = image_files[:10]
  483. print(f"🧪 Test mode: processing only {len(image_files)} images")
  484. print(f"🌐 Using server: {args.ip}:{args.port}")
  485. print(f"📦 Batch size: {args.batch_size}")
  486. print(f"🎯 Prompt mode: {args.prompt_mode}")
  487. # 创建处理器
  488. processor = DotsOCRProcessor(
  489. ip=args.ip,
  490. port=args.port,
  491. model_name=args.model_name,
  492. prompt_mode=args.prompt_mode,
  493. dpi=args.dpi,
  494. min_pixels=args.min_pixels,
  495. max_pixels=args.max_pixels
  496. )
  497. # 开始处理
  498. start_time = time.time()
  499. # 选择处理方式
  500. if args.use_threading:
  501. results = process_images_concurrent(
  502. image_files,
  503. processor,
  504. args.batch_size,
  505. str(output_dir),
  506. args.max_workers
  507. )
  508. else:
  509. results = process_images_single_process(
  510. image_files,
  511. processor,
  512. args.batch_size,
  513. str(output_dir)
  514. )
  515. total_time = time.time() - start_time
  516. # 统计结果
  517. success_count = sum(1 for r in results if r.get('success', False))
  518. skipped_count = sum(1 for r in results if r.get('skipped', False))
  519. error_count = len(results) - success_count
  520. pdf_page_count = sum(1 for r in results if r.get('is_pdf_page', False))
  521. print(f"\n" + "="*60)
  522. print(f"✅ Processing completed!")
  523. print(f"📊 Statistics:")
  524. print(f" Total files processed: {len(image_files)}")
  525. print(f" PDF pages processed: {pdf_page_count}")
  526. print(f" Regular images processed: {len(image_files) - pdf_page_count}")
  527. print(f" Successful: {success_count}")
  528. print(f" Skipped: {skipped_count}")
  529. print(f" Failed: {error_count}")
  530. if len(image_files) > 0:
  531. print(f" Success rate: {success_count / len(image_files) * 100:.2f}%")
  532. print(f"⏱️ Performance:")
  533. print(f" Total time: {total_time:.2f} seconds")
  534. if total_time > 0:
  535. print(f" Throughput: {len(image_files) / total_time:.2f} images/second")
  536. print(f" Avg time per image: {total_time / len(image_files):.2f} seconds")
  537. # 保存结果统计
  538. stats = {
  539. "total_files": len(image_files),
  540. "pdf_pages": pdf_page_count,
  541. "regular_images": len(image_files) - pdf_page_count,
  542. "success_count": success_count,
  543. "skipped_count": skipped_count,
  544. "error_count": error_count,
  545. "success_rate": success_count / len(image_files) if len(image_files) > 0 else 0,
  546. "total_time": total_time,
  547. "throughput": len(image_files) / total_time if total_time > 0 else 0,
  548. "avg_time_per_image": total_time / len(image_files) if len(image_files) > 0 else 0,
  549. "batch_size": args.batch_size,
  550. "server": f"{args.ip}:{args.port}",
  551. "model": args.model_name,
  552. "prompt_mode": args.prompt_mode,
  553. "pdf_dpi": args.dpi,
  554. "timestamp": time.strftime("%Y-%m-%d %H:%M:%S")
  555. }
  556. # 保存最终结果
  557. output_file_name = Path(output_dir).name
  558. output_file = os.path.join(output_dir, f"{output_file_name}_results.json")
  559. final_results = {
  560. "stats": stats,
  561. "results": results
  562. }
  563. with open(output_file, 'w', encoding='utf-8') as f:
  564. json.dump(final_results, f, ensure_ascii=False, indent=2)
  565. print(f"💾 Results saved to: {output_file}")
  566. # 收集处理结果
  567. if not args.collect_results:
  568. output_file_processed = Path(args.output_dir) / f"processed_files_{time.strftime('%Y%m%d_%H%M%S')}.csv"
  569. else:
  570. output_file_processed = Path(args.collect_results).resolve()
  571. processed_files = collect_pid_files(output_file)
  572. with open(output_file_processed, 'w', encoding='utf-8') as f:
  573. f.write("image_path,status\n")
  574. for file_path, status in processed_files:
  575. f.write(f"{file_path},{status}\n")
  576. print(f"💾 Processed files saved to: {output_file_processed}")
  577. return 0
  578. except Exception as e:
  579. print(f"❌ Processing failed: {e}", file=sys.stderr)
  580. traceback.print_exc()
  581. return 1
  582. if __name__ == "__main__":
  583. print(f"🚀 启动DotsOCR统一PDF/图像处理程序...")
  584. print(f"🔧 CUDA_VISIBLE_DEVICES: {os.environ.get('CUDA_VISIBLE_DEVICES', 'Not set')}")
  585. if len(sys.argv) == 1:
  586. # 如果没有命令行参数,使用默认配置运行
  587. print("ℹ️ No command line arguments provided. Running with default configuration...")
  588. # 默认配置
  589. default_config = {
  590. "input_dir": "../../OmniDocBench/OpenDataLab___OmniDocBench/images",
  591. "output_dir": "./OmniDocBench_DotsOCR_Results",
  592. # "ip": "10.192.72.11",
  593. "ip": "127.0.0.1",
  594. "port": "8101",
  595. "model_name": "DotsOCR",
  596. "prompt_mode": "prompt_layout_all_en",
  597. "batch_size": "1",
  598. "max_workers": "3",
  599. "dpi": "200",
  600. "collect_results": "./OmniDocBench_DotsOCR_Results/processed_files.csv",
  601. }
  602. # 如果需要处理失败的文件,可以使用这个配置
  603. # default_config = {
  604. # "input_csv": "./OmniDocBench_DotsOCR_Results/processed_files.csv",
  605. # "output_dir": "./OmniDocBench_DotsOCR_Results",
  606. # "ip": "127.0.0.1",
  607. # "port": "8101",
  608. # "collect_results": f"./OmniDocBench_DotsOCR_Results/processed_files_{time.strftime('%Y%m%d_%H%M%S')}.csv",
  609. # }
  610. # 构造参数
  611. sys.argv = [sys.argv[0]]
  612. for key, value in default_config.items():
  613. sys.argv.extend([f"--{key}", str(value)])
  614. # 测试模式
  615. sys.argv.append("--use_threading")
  616. # sys.argv.append("--test_mode")
  617. sys.exit(main())