batch_process_pdf.py 32 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890
  1. #!/usr/bin/env python3
  2. """
  3. PDF 批量处理脚本
  4. 支持多种处理器,配置文件驱动
  5. 支持自动切换虚拟环境
  6. 支持执行器输出日志重定向
  7. """
  8. import os
  9. import sys
  10. import argparse
  11. import subprocess
  12. import json
  13. import yaml
  14. from pathlib import Path
  15. from datetime import datetime
  16. from typing import List, Dict, Optional, Any
  17. from dataclasses import dataclass, field
  18. import logging
  19. from tqdm import tqdm
  20. import time
  21. # ============================================================================
  22. # 数据类定义
  23. # ============================================================================
  24. @dataclass
  25. class ProcessorConfig:
  26. """处理器配置"""
  27. name: str
  28. script: str
  29. input_arg: str = "--input_file"
  30. output_arg: str = "--output_dir"
  31. extra_args: List[str] = field(default_factory=list)
  32. output_subdir: str = "results"
  33. log_subdir: str = "logs" # 🎯 新增:日志子目录
  34. venv: Optional[str] = None
  35. description: str = ""
  36. @dataclass
  37. class ProcessResult:
  38. """处理结果"""
  39. pdf_file: str
  40. success: bool
  41. duration: float
  42. error_message: str = ""
  43. log_file: str = "" # 🎯 新增:日志文件路径
  44. # ============================================================================
  45. # 配置管理
  46. # ============================================================================
  47. class ConfigManager:
  48. """配置管理器"""
  49. DEFAULT_CONFIG = {
  50. 'processors': {
  51. 'paddleocr_vl_single_process': {
  52. 'script': '/Users/zhch158/workspace/repository.git/PaddleX/zhch/paddleocr_vl_single_process.py',
  53. 'input_arg': '--input_file',
  54. 'output_arg': '--output_dir',
  55. 'extra_args': [
  56. '--pipeline=/Users/zhch158/workspace/repository.git/PaddleX/zhch/my_config/PaddleOCR-VL-Client-RT-DETR-H_layout_17cls.yaml',
  57. '--no-adapter'
  58. ],
  59. 'output_subdir': 'paddleocr_vl_results',
  60. 'venv': 'source /Users/zhch158/workspace/repository.git/PaddleX/paddle_env/bin/activate',
  61. 'description': 'PaddleOCR-VL 处理器',
  62. 'log_subdir': 'logs/paddleocr_vl_single_process' # 🎯 新增
  63. },
  64. 'ppstructurev3_single_process': {
  65. 'script': '/Users/zhch158/workspace/repository.git/PaddleX/zhch/ppstructurev3_single_process.py',
  66. 'input_arg': '--input_file',
  67. 'output_arg': '--output_dir',
  68. 'extra_args': [
  69. '--pipeline=/Users/zhch158/workspace/repository.git/PaddleX/zhch/my_config/PP-StructureV3.yaml'
  70. ],
  71. 'output_subdir': 'ppstructurev3_results',
  72. 'venv': 'conda activate paddle',
  73. 'description': 'PP-StructureV3 处理器',
  74. 'log_subdir': 'logs/ppstructurev3_single_process' # 🎯 新增
  75. },
  76. 'ppstructurev3_single_client': {
  77. 'script': '/Users/zhch158/workspace/repository.git/PaddleX/zhch/ppstructurev3_single_client.py',
  78. 'input_arg': '--input_file',
  79. 'output_arg': '--output_dir',
  80. 'extra_args': [
  81. '--api_url=http://10.192.72.11:8111/layout-parsing',
  82. '--timeout=300'
  83. ],
  84. 'output_subdir': 'ppstructurev3_client_results',
  85. 'venv': 'source /Users/zhch158/workspace/repository.git/PaddleX/paddle_env/bin/activate',
  86. 'description': 'PP-StructureV3 HTTP API 客户端',
  87. 'log_subdir': 'logs/ppstructurev3_single_client' # 🎯 新增
  88. },
  89. 'mineru_vllm': {
  90. 'script': '/Users/zhch158/workspace/repository.git/MinerU/zhch/mineru2_vllm_multthreads.py',
  91. 'input_arg': '--input_file',
  92. 'output_arg': '--output_dir',
  93. 'extra_args': [
  94. '--server_url=http://10.192.72.11:8121',
  95. '--timeout=300',
  96. '--batch_size=1'
  97. ],
  98. 'output_subdir': 'mineru_vllm_results',
  99. 'venv': 'conda activate mineru2',
  100. 'description': 'MinerU vLLM 处理器',
  101. 'log_subdir': 'logs/mineru_vllm' # 🎯 新增
  102. },
  103. 'dotsocr_vllm': {
  104. 'script': '/Users/zhch158/workspace/repository.git/dots.ocr/zhch/dotsocr_vllm_multthreads.py',
  105. 'input_arg': '--input_file',
  106. 'output_arg': '--output_dir',
  107. 'extra_args': [
  108. '--ip=10.192.72.11',
  109. '--port=8101',
  110. '--model_name=DotsOCR',
  111. '--prompt_mode=prompt_layout_all_en',
  112. '--batch_size=1',
  113. '--max_workers=1',
  114. '--dpi=200'
  115. ],
  116. 'output_subdir': 'dotsocr_vllm_results',
  117. 'venv': 'conda activate py312',
  118. 'description': 'DotsOCR vLLM 处理器 - 支持PDF和图片',
  119. 'log_subdir': 'logs/dotsocr_vllm' # 🎯 新增
  120. }
  121. },
  122. 'global': {
  123. 'base_dir': '/Users/zhch158/workspace/data/流水分析',
  124. 'output_subdir': 'results',
  125. 'log_dir': 'logs',
  126. 'log_retention_days': 30,
  127. 'log_level': 'INFO'
  128. }
  129. }
  130. def __init__(self, config_file: Optional[str] = None):
  131. self.config_file = config_file
  132. self.config = self._load_config()
  133. def _load_config(self) -> Dict:
  134. """加载配置文件"""
  135. if self.config_file and Path(self.config_file).exists():
  136. with open(self.config_file, 'r', encoding='utf-8') as f:
  137. if self.config_file.endswith('.yaml') or self.config_file.endswith('.yml'):
  138. return yaml.safe_load(f)
  139. else:
  140. return json.load(f)
  141. return self.DEFAULT_CONFIG.copy()
  142. def get_processor_config(self, processor_name: str) -> ProcessorConfig:
  143. """获取处理器配置"""
  144. if processor_name not in self.config['processors']:
  145. raise ValueError(f"处理器 '{processor_name}' 不存在")
  146. proc_config = self.config['processors'][processor_name]
  147. return ProcessorConfig(
  148. name=processor_name,
  149. script=proc_config['script'],
  150. input_arg=proc_config.get('input_arg', '--input_file'),
  151. output_arg=proc_config.get('output_arg', '--output_dir'),
  152. extra_args=proc_config.get('extra_args', []),
  153. output_subdir=proc_config.get('output_subdir', processor_name + '_results'),
  154. log_subdir=proc_config.get('log_subdir', f'logs/{processor_name}'), # 🎯 新增
  155. venv=proc_config.get('venv'),
  156. description=proc_config.get('description', '')
  157. )
  158. def get_global_config(self, key: str, default=None):
  159. """获取全局配置"""
  160. return self.config.get('global', {}).get(key, default)
  161. def list_processors(self) -> List[str]:
  162. """列出所有可用的处理器"""
  163. return list(self.config['processors'].keys())
  164. # ============================================================================
  165. # PDF 文件查找器
  166. # ============================================================================
  167. class PDFFileFinder:
  168. """PDF 文件查找器"""
  169. def __init__(self, base_dir: str):
  170. self.base_dir = Path(base_dir)
  171. def from_file_list(self, list_file: str) -> List[Path]:
  172. """从文件列表读取"""
  173. pdf_files = []
  174. with open(list_file, 'r', encoding='utf-8') as f:
  175. for line in f:
  176. # 跳过空行和注释
  177. line = line.strip()
  178. if not line or line.startswith('#'):
  179. continue
  180. # 构建完整路径
  181. pdf_path = self._resolve_path(line)
  182. if pdf_path:
  183. pdf_files.append(pdf_path)
  184. return pdf_files
  185. def from_list(self, pdf_list: List[str]) -> List[Path]:
  186. """从列表读取"""
  187. pdf_files = []
  188. for pdf in pdf_list:
  189. pdf_path = self._resolve_path(pdf.strip())
  190. if pdf_path:
  191. pdf_files.append(pdf_path)
  192. return pdf_files
  193. def find_all(self) -> List[Path]:
  194. """查找基础目录下所有 PDF"""
  195. return sorted(self.base_dir.rglob('*.pdf'))
  196. def _resolve_path(self, path_str: str) -> Optional[Path]:
  197. """解析路径"""
  198. path = Path(path_str)
  199. # 绝对路径
  200. if path.is_absolute():
  201. return path if path.exists() else path # 返回路径,即使不存在
  202. # 相对路径
  203. # 1. 尝试完整相对路径
  204. candidate1 = self.base_dir / path
  205. if candidate1.exists():
  206. return candidate1
  207. # 2. 尝试在同名子目录下查找
  208. if '/' not in path_str:
  209. pdf_name = path.stem
  210. candidate2 = self.base_dir / pdf_name / path.name
  211. if candidate2.exists():
  212. return candidate2
  213. # 3. 使用 glob 搜索
  214. matches = list(self.base_dir.rglob(path.name))
  215. if matches:
  216. return matches[0]
  217. # 返回候选路径(即使不存在)
  218. return candidate1
  219. # ============================================================================
  220. # PDF 批处理器
  221. # ============================================================================
  222. class PDFBatchProcessor:
  223. """PDF 批处理器"""
  224. def __init__(
  225. self,
  226. processor_config: ProcessorConfig,
  227. output_subdir: Optional[str] = None,
  228. log_base_dir: Optional[str] = None, # 🎯 新增:日志基础目录
  229. dry_run: bool = False
  230. ):
  231. self.processor_config = processor_config
  232. # 如果指定了output_subdir,使用指定的;否则使用处理器配置中的
  233. self.output_subdir = output_subdir or processor_config.output_subdir
  234. self.log_base_dir = Path(log_base_dir) if log_base_dir else Path('logs') # 🎯 新增
  235. self.dry_run = dry_run
  236. # 设置日志
  237. self.logger = self._setup_logger()
  238. # 统计信息
  239. self.results: List[ProcessResult] = []
  240. def _setup_logger(self) -> logging.Logger:
  241. """设置日志"""
  242. logger = logging.getLogger('PDFBatchProcessor')
  243. logger.setLevel(logging.INFO)
  244. # 避免重复添加handler
  245. if not logger.handlers:
  246. # 控制台输出
  247. console_handler = logging.StreamHandler()
  248. console_handler.setLevel(logging.INFO)
  249. console_format = logging.Formatter(
  250. '%(asctime)s - %(levelname)s - %(message)s',
  251. datefmt='%Y-%m-%d %H:%M:%S'
  252. )
  253. console_handler.setFormatter(console_format)
  254. logger.addHandler(console_handler)
  255. return logger
  256. def _get_log_file_path(self, pdf_file: Path) -> Path:
  257. """
  258. 🎯 获取日志文件路径
  259. 日志结构:
  260. base_dir/
  261. └── PDF名称/
  262. └── logs/
  263. └── processor_name/
  264. └── PDF名称_YYYYMMDD_HHMMSS.log
  265. """
  266. # PDF 目录
  267. pdf_dir = pdf_file.parent / pdf_file.stem
  268. # 日志目录: pdf_dir / logs / processor_name
  269. log_dir = pdf_dir / self.processor_config.log_subdir
  270. log_dir.mkdir(parents=True, exist_ok=True)
  271. # 日志文件名: PDF名称_时间戳.log
  272. timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
  273. log_file = log_dir / f"{pdf_file.stem}_{timestamp}.log"
  274. return log_file
  275. def process_files(self, pdf_files: List[Path]) -> Dict[str, Any]:
  276. """批量处理文件"""
  277. self.logger.info(f"开始处理 {len(pdf_files)} 个文件")
  278. self.logger.info(f"处理器: {self.processor_config.description}")
  279. self.logger.info(f"脚本: {self.processor_config.script}")
  280. self.logger.info(f"输出目录: {self.output_subdir}")
  281. self.logger.info(f"日志目录: {self.processor_config.log_subdir}")
  282. if self.processor_config.venv:
  283. self.logger.info(f"虚拟环境: {self.processor_config.venv}")
  284. start_time = time.time()
  285. # 使用进度条
  286. with tqdm(total=len(pdf_files), desc="处理进度", unit="file") as pbar:
  287. for pdf_file in pdf_files:
  288. result = self._process_single_file(pdf_file)
  289. self.results.append(result)
  290. pbar.update(1)
  291. # 更新进度条描述
  292. success_count = sum(1 for r in self.results if r.success)
  293. pbar.set_postfix({
  294. 'success': success_count,
  295. 'failed': len(self.results) - success_count
  296. })
  297. total_duration = time.time() - start_time
  298. # 生成统计信息
  299. stats = self._generate_stats(total_duration)
  300. self._save_summary_log(stats)
  301. return stats
  302. def _process_single_file(self, pdf_file: Path) -> ProcessResult:
  303. """🎯 处理单个文件(支持日志重定向)"""
  304. self.logger.info(f"处理: {pdf_file}")
  305. # 检查文件是否存在
  306. if not pdf_file.exists():
  307. self.logger.warning(f"跳过: 文件不存在 - {pdf_file}")
  308. return ProcessResult(
  309. pdf_file=str(pdf_file),
  310. success=False,
  311. duration=0,
  312. error_message="文件不存在"
  313. )
  314. # 确定输出目录
  315. output_dir = pdf_file.parent / pdf_file.stem / self.output_subdir
  316. # 🎯 获取日志文件路径
  317. log_file = self._get_log_file_path(pdf_file)
  318. # 构建命令
  319. cmd = self._build_command(pdf_file, output_dir)
  320. self.logger.debug(f"执行命令: {cmd if isinstance(cmd, str) else ' '.join(cmd)}")
  321. self.logger.info(f"日志输出: {log_file}")
  322. if self.dry_run:
  323. self.logger.info(f"[DRY RUN] 将执行: {cmd if isinstance(cmd, str) else ' '.join(cmd)}")
  324. return ProcessResult(
  325. pdf_file=str(pdf_file),
  326. success=True,
  327. duration=0,
  328. error_message="",
  329. log_file=str(log_file)
  330. )
  331. # 🎯 执行命令并重定向输出到日志文件
  332. start_time = time.time()
  333. try:
  334. with open(log_file, 'w', encoding='utf-8') as log_f:
  335. # 写入日志头
  336. log_f.write(f"{'='*80}\n")
  337. log_f.write(f"处理器: {self.processor_config.description}\n")
  338. log_f.write(f"PDF 文件: {pdf_file}\n")
  339. log_f.write(f"输出目录: {output_dir}\n")
  340. log_f.write(f"开始时间: {datetime.now()}\n")
  341. log_f.write(f"{'='*80}\n\n")
  342. log_f.flush()
  343. # 执行命令
  344. if isinstance(cmd, str):
  345. result = subprocess.run(
  346. cmd,
  347. shell=True,
  348. executable='/bin/bash',
  349. stdout=log_f, # 🎯 重定向 stdout
  350. stderr=subprocess.STDOUT, # 🎯 合并 stderr 到 stdout
  351. text=True,
  352. check=True
  353. )
  354. else:
  355. result = subprocess.run(
  356. cmd,
  357. stdout=log_f, # 🎯 重定向 stdout
  358. stderr=subprocess.STDOUT, # 🎯 合并 stderr
  359. text=True,
  360. check=True
  361. )
  362. # 写入日志尾
  363. log_f.write(f"\n{'='*80}\n")
  364. log_f.write(f"结束时间: {datetime.now()}\n")
  365. log_f.write(f"状态: 成功\n")
  366. log_f.write(f"{'='*80}\n")
  367. duration = time.time() - start_time
  368. self.logger.info(f"✓ 成功 (耗时: {duration:.2f}秒)")
  369. return ProcessResult(
  370. pdf_file=str(pdf_file),
  371. success=True,
  372. duration=duration,
  373. error_message="",
  374. log_file=str(log_file)
  375. )
  376. except subprocess.CalledProcessError as e:
  377. duration = time.time() - start_time
  378. error_msg = f"命令执行失败 (退出码: {e.returncode})"
  379. # 🎯 在日志文件中追加错误信息
  380. with open(log_file, 'a', encoding='utf-8') as log_f:
  381. log_f.write(f"\n{'='*80}\n")
  382. log_f.write(f"结束时间: {datetime.now()}\n")
  383. log_f.write(f"状态: 失败\n")
  384. log_f.write(f"错误: {error_msg}\n")
  385. log_f.write(f"{'='*80}\n")
  386. self.logger.error(f"✗ 失败 (耗时: {duration:.2f}秒)")
  387. self.logger.error(f"错误信息: {error_msg}")
  388. self.logger.error(f"详细日志: {log_file}")
  389. return ProcessResult(
  390. pdf_file=str(pdf_file),
  391. success=False,
  392. duration=duration,
  393. error_message=error_msg,
  394. log_file=str(log_file)
  395. )
  396. except Exception as e:
  397. duration = time.time() - start_time
  398. error_msg = str(e)
  399. with open(log_file, 'a', encoding='utf-8') as log_f:
  400. log_f.write(f"\n{'='*80}\n")
  401. log_f.write(f"结束时间: {datetime.now()}\n")
  402. log_f.write(f"状态: 异常\n")
  403. log_f.write(f"错误: {error_msg}\n")
  404. log_f.write(f"{'='*80}\n")
  405. self.logger.error(f"✗ 异常 (耗时: {duration:.2f}秒)")
  406. self.logger.error(f"错误信息: {error_msg}")
  407. return ProcessResult(
  408. pdf_file=str(pdf_file),
  409. success=False,
  410. duration=duration,
  411. error_message=error_msg,
  412. log_file=str(log_file)
  413. )
  414. def _build_command(self, pdf_file: Path, output_dir: Path):
  415. """构建执行命令
  416. Returns:
  417. 如果配置了 venv,返回 shell 命令字符串
  418. 否则返回命令列表
  419. """
  420. # 构建基础 Python 命令
  421. base_cmd = [
  422. 'python', # 使用虚拟环境中的 python
  423. self.processor_config.script,
  424. self.processor_config.input_arg, str(pdf_file),
  425. self.processor_config.output_arg, str(output_dir)
  426. ]
  427. # 添加额外参数
  428. base_cmd.extend(self.processor_config.extra_args)
  429. # 如果配置了虚拟环境,构建 shell 命令
  430. if self.processor_config.venv:
  431. # 转义参数中的特殊字符
  432. escaped_cmd = []
  433. for arg in base_cmd:
  434. if ' ' in arg or '"' in arg or "'" in arg:
  435. # 使用单引号包裹,内部单引号转义
  436. arg = arg.replace("'", "'\\''")
  437. escaped_cmd.append(f"'{arg}'")
  438. else:
  439. escaped_cmd.append(arg)
  440. python_cmd = ' '.join(escaped_cmd)
  441. # 检查是否使用 conda
  442. if 'conda activate' in self.processor_config.venv:
  443. # 获取 conda 基础路径
  444. # 对于 conda,需要先 source conda.sh,然后 conda activate
  445. conda_init = """
  446. eval "$(conda shell.bash hook)"
  447. """.strip()
  448. shell_cmd = f"{conda_init} && {self.processor_config.venv} && {python_cmd}"
  449. else:
  450. # 对于 source 激活的虚拟环境
  451. shell_cmd = f"{self.processor_config.venv} && {python_cmd}"
  452. return shell_cmd
  453. else:
  454. # 没有虚拟环境,使用当前 Python 解释器
  455. base_cmd[0] = sys.executable
  456. return base_cmd
  457. def _generate_stats(self, total_duration: float) -> Dict[str, Any]:
  458. """生成统计信息"""
  459. success_count = sum(1 for r in self.results if r.success)
  460. failed_count = len(self.results) - success_count
  461. failed_files = [
  462. {
  463. 'file': r.pdf_file,
  464. 'error': r.error_message,
  465. 'log': r.log_file
  466. }
  467. for r in self.results if not r.success
  468. ]
  469. stats = {
  470. 'total': len(self.results),
  471. 'success': success_count,
  472. 'failed': failed_count,
  473. 'total_duration': total_duration,
  474. 'failed_files': failed_files,
  475. 'results': [
  476. {
  477. 'file': r.pdf_file,
  478. 'success': r.success,
  479. 'duration': r.duration,
  480. 'error': r.error_message,
  481. 'log': r.log_file
  482. }
  483. for r in self.results
  484. ]
  485. }
  486. return stats
  487. def _save_summary_log(self, stats: Dict[str, Any]):
  488. """🎯 保存汇总日志"""
  489. timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
  490. summary_log_file = self.log_base_dir / f"batch_summary_{self.processor_config.name}_{timestamp}.log"
  491. # 确保目录存在
  492. summary_log_file.parent.mkdir(parents=True, exist_ok=True)
  493. with open(summary_log_file, 'w', encoding='utf-8') as f:
  494. f.write("PDF 批量处理汇总日志\n")
  495. f.write("=" * 80 + "\n\n")
  496. f.write(f"处理器: {self.processor_config.description}\n")
  497. f.write(f"处理器名称: {self.processor_config.name}\n")
  498. f.write(f"脚本: {self.processor_config.script}\n")
  499. f.write(f"输出目录: {self.output_subdir}\n")
  500. f.write(f"日志目录: {self.processor_config.log_subdir}\n")
  501. if self.processor_config.venv:
  502. f.write(f"虚拟环境: {self.processor_config.venv}\n")
  503. f.write(f"开始时间: {datetime.now()}\n")
  504. f.write(f"总耗时: {stats['total_duration']:.2f} 秒\n\n")
  505. f.write("统计信息:\n")
  506. f.write(f" 总文件数: {stats['total']}\n")
  507. f.write(f" 成功: {stats['success']}\n")
  508. f.write(f" 失败: {stats['failed']}\n\n")
  509. if stats['failed_files']:
  510. f.write("失败的文件:\n")
  511. for item in stats['failed_files']:
  512. f.write(f" ✗ {item['file']}\n")
  513. f.write(f" 错误: {item['error']}\n")
  514. f.write(f" 日志: {item['log']}\n\n")
  515. f.write("详细结果:\n")
  516. for result in stats['results']:
  517. status = "✓" if result['success'] else "✗"
  518. f.write(f"{status} {result['file']} ({result['duration']:.2f}s)\n")
  519. f.write(f" 日志: {result['log']}\n")
  520. if result['error']:
  521. f.write(f" 错误: {result['error']}\n")
  522. self.logger.info(f"汇总日志已保存: {summary_log_file}")
  523. # ============================================================================
  524. # 命令行接口
  525. # ============================================================================
  526. def create_parser() -> argparse.ArgumentParser:
  527. """创建命令行参数解析器"""
  528. parser = argparse.ArgumentParser(
  529. description='PDF 批量处理工具 (支持虚拟环境自动切换)',
  530. formatter_class=argparse.RawDescriptionHelpFormatter,
  531. epilog="""
  532. 示例用法:
  533. 1. 使用配置文件中的处理器 (自动切换虚拟环境):
  534. python batch_process_pdf.py -p paddleocr_vl_single_process -f pdf_list.txt
  535. 2. 使用 DotsOCR 处理器 (自动切换到 py312 环境):
  536. python batch_process_pdf.py -p dotsocr_vllm -f pdf_list.txt
  537. 3. 使用 MinerU 处理器 (自动切换到 mineru2 环境):
  538. python batch_process_pdf.py -p mineru_vllm -f pdf_list.txt
  539. 4. 处理指定目录下所有 PDF:
  540. python batch_process_pdf.py -p ppstructurev3_single_client -d /path/to/pdfs
  541. 5. 列出所有可用的处理器:
  542. python batch_process_pdf.py --list-processors
  543. 6. 手动指定虚拟环境:
  544. python batch_process_pdf.py -p paddleocr_vl_single_process -f pdf_list.txt --venv "conda activate paddle"
  545. """
  546. )
  547. # 处理器选择
  548. parser.add_argument(
  549. '-p', '--processor',
  550. help='处理器名称'
  551. )
  552. # 配置文件
  553. parser.add_argument(
  554. '-c', '--config',
  555. default='processor_configs.yaml',
  556. help='配置文件路径 (默认: processor_configs.yaml)'
  557. )
  558. # 手动指定脚本
  559. parser.add_argument(
  560. '-s', '--script',
  561. help='Python 脚本路径 (覆盖配置文件)'
  562. )
  563. # 目录和文件
  564. parser.add_argument(
  565. '-d', '--base-dir',
  566. help='PDF 文件基础目录'
  567. )
  568. parser.add_argument(
  569. '-o', '--output-subdir',
  570. help='输出子目录名称 (覆盖处理器默认配置)'
  571. )
  572. parser.add_argument(
  573. '-f', '--file-list',
  574. help='PDF 文件列表文件路径'
  575. )
  576. parser.add_argument(
  577. '-l', '--pdf-list',
  578. nargs='+',
  579. help='PDF 文件列表 (空格分隔)'
  580. )
  581. # 额外参数
  582. parser.add_argument(
  583. '-e', '--extra-args',
  584. help='额外参数 (覆盖配置文件)'
  585. )
  586. # 虚拟环境
  587. parser.add_argument(
  588. '--venv',
  589. help='虚拟环境激活命令 (覆盖配置文件)'
  590. )
  591. # 工具选项
  592. parser.add_argument(
  593. '--list-processors',
  594. action='store_true',
  595. help='列出所有可用的处理器'
  596. )
  597. parser.add_argument(
  598. '--show-config',
  599. action='store_true',
  600. help='显示配置文件内容'
  601. )
  602. parser.add_argument(
  603. '--dry-run',
  604. action='store_true',
  605. help='模拟运行,不实际执行'
  606. )
  607. parser.add_argument(
  608. '-v', '--verbose',
  609. action='store_true',
  610. help='详细输出'
  611. )
  612. return parser
  613. def main():
  614. """主函数"""
  615. parser = create_parser()
  616. args = parser.parse_args()
  617. # 设置日志级别
  618. if args.verbose:
  619. logging.getLogger().setLevel(logging.DEBUG)
  620. # 加载配置
  621. config_manager = ConfigManager(args.config if Path(args.config).exists() else None)
  622. # 列出处理器
  623. if args.list_processors:
  624. print("可用的处理器:")
  625. for name in config_manager.list_processors():
  626. proc_config = config_manager.get_processor_config(name)
  627. print(f" • {name}")
  628. print(f" 描述: {proc_config.description}")
  629. print(f" 脚本: {proc_config.script}")
  630. print(f" 输出目录: {proc_config.output_subdir}")
  631. if proc_config.venv:
  632. print(f" 虚拟环境: {proc_config.venv}")
  633. print()
  634. return 0
  635. # 显示配置
  636. if args.show_config:
  637. print(yaml.dump(config_manager.config, allow_unicode=True))
  638. return 0
  639. # 获取处理器配置
  640. if args.processor:
  641. processor_config = config_manager.get_processor_config(args.processor)
  642. elif args.script:
  643. # 手动指定脚本
  644. processor_config = ProcessorConfig(
  645. name='manual',
  646. script=args.script,
  647. extra_args=args.extra_args.split() if args.extra_args else [],
  648. output_subdir=args.output_subdir or 'manual_results',
  649. venv=args.venv
  650. )
  651. else:
  652. parser.error("必须指定 -p 或 -s 参数")
  653. # 覆盖额外参数
  654. if args.extra_args and args.processor:
  655. processor_config.extra_args = args.extra_args.split()
  656. # 覆盖虚拟环境
  657. if args.venv:
  658. processor_config.venv = args.venv
  659. # 获取基础目录
  660. base_dir = args.base_dir or config_manager.get_global_config('base_dir')
  661. if not base_dir:
  662. parser.error("必须指定 -d 参数或在配置文件中设置 base_dir")
  663. log_base_dir = base_dir + '/' + config_manager.get_global_config('log_dir', 'logs')
  664. # 查找 PDF 文件
  665. finder = PDFFileFinder(base_dir)
  666. if args.file_list:
  667. pdf_files = finder.from_file_list(args.file_list)
  668. elif args.pdf_list:
  669. pdf_files = finder.from_list(args.pdf_list)
  670. else:
  671. pdf_files = finder.find_all()
  672. if not pdf_files:
  673. print("❌ 未找到任何 PDF 文件")
  674. return 1
  675. # 显示找到的文件
  676. valid_file_paths = [f.as_posix() for f in pdf_files if f.exists()]
  677. if valid_file_paths:
  678. print("\n".join(valid_file_paths))
  679. # 验证文件
  680. valid_files = [f for f in pdf_files if f.exists()]
  681. invalid_files = [f for f in pdf_files if not f.exists()]
  682. if invalid_files:
  683. print(f"\n⚠️ 警告: {len(invalid_files)} 个文件不存在:")
  684. for f in invalid_files[:5]:
  685. print(f" - {f}")
  686. if len(invalid_files) > 5:
  687. print(f" ... 还有 {len(invalid_files) - 5} 个")
  688. # 确认执行
  689. if not args.dry_run and valid_files:
  690. venv_info = f" (虚拟环境: {processor_config.venv})" if processor_config.venv else ""
  691. confirm = input(f"\n是否继续处理 {len(valid_files)} 个文件{venv_info}? [Y/n]: ")
  692. if confirm.lower() not in ['', 'y', 'yes']:
  693. print("已取消")
  694. return 0
  695. # 批量处理
  696. processor = PDFBatchProcessor(
  697. processor_config=processor_config,
  698. output_subdir=args.output_subdir,
  699. log_base_dir=log_base_dir, # 🎯 传递日志目录
  700. dry_run=args.dry_run
  701. )
  702. stats = processor.process_files(valid_files)
  703. # 显示统计信息
  704. print("\n" + "=" * 80)
  705. print("处理完成")
  706. print("=" * 80)
  707. print(f"\n📊 统计信息:")
  708. print(f" 处理器: {processor_config.description}")
  709. print(f" 输出目录: {processor.output_subdir}")
  710. print(f" 日志目录: {processor.processor_config.log_subdir}")
  711. print(f" 总文件数: {stats['total']}")
  712. print(f" ✓ 成功: {stats['success']}")
  713. print(f" ✗ 失败: {stats['failed']}")
  714. print(f" ⏱️ 总耗时: {stats['total_duration']:.2f} 秒")
  715. if stats['failed_files']:
  716. print(f"\n失败的文件:")
  717. for item in stats['failed_files']:
  718. print(f" ✗ {item['file']}")
  719. print(f" 错误: {item['error']}")
  720. print(f" 日志: {item['log']}")
  721. return 0 if stats['failed'] == 0 else 1
  722. if __name__ == '__main__':
  723. print("🚀 启动批量OCR程序...")
  724. import sys
  725. if len(sys.argv) == 1:
  726. # 如果没有命令行参数,使用默认配置运行
  727. print("ℹ️ 未提供命令行参数,使用默认配置运行...")
  728. # 默认配置
  729. default_config = {
  730. "processor": "mineru_vllm",
  731. "file-list": "pdf_list.txt",
  732. }
  733. print("⚙️ 默认参数:")
  734. for key, value in default_config.items():
  735. print(f" --{key}: {value}")
  736. # 构造参数
  737. sys.argv = [sys.argv[0]]
  738. for key, value in default_config.items():
  739. sys.argv.extend([f"--{key}", str(value)])
  740. sys.argv.append("--dry-run")
  741. sys.argv.append("--verbose") # 添加详细输出参数
  742. sys.exit(main())