|
@@ -25,6 +25,7 @@ import json
|
|
|
import sys
|
|
import sys
|
|
|
import os
|
|
import os
|
|
|
from pathlib import Path
|
|
from pathlib import Path
|
|
|
|
|
+from typing import Optional
|
|
|
from loguru import logger
|
|
from loguru import logger
|
|
|
from datetime import datetime
|
|
from datetime import datetime
|
|
|
|
|
|
|
@@ -43,10 +44,11 @@ from dotenv import load_dotenv
|
|
|
load_dotenv(override=True)
|
|
load_dotenv(override=True)
|
|
|
|
|
|
|
|
from core.pipeline_manager_v2 import EnhancedDocPipeline
|
|
from core.pipeline_manager_v2 import EnhancedDocPipeline
|
|
|
|
|
+from core.pipeline_manager_v2_streaming import StreamingDocPipeline
|
|
|
from universal_doc_parser.utils import OutputFormatterV2
|
|
from universal_doc_parser.utils import OutputFormatterV2
|
|
|
|
|
|
|
|
|
|
|
|
|
-def setup_logging(log_level: str = "INFO", log_file: str = None):
|
|
|
|
|
|
|
+def setup_logging(log_level: str = "INFO", log_file: Optional[str] = None):
|
|
|
"""设置日志"""
|
|
"""设置日志"""
|
|
|
logger.remove()
|
|
logger.remove()
|
|
|
|
|
|
|
@@ -72,8 +74,9 @@ def process_single_input(
|
|
|
config_path: Path,
|
|
config_path: Path,
|
|
|
output_dir: Path,
|
|
output_dir: Path,
|
|
|
debug: bool = False,
|
|
debug: bool = False,
|
|
|
- scene: str = None,
|
|
|
|
|
- page_range: str = None
|
|
|
|
|
|
|
+ scene: Optional[str] = None,
|
|
|
|
|
+ page_range: Optional[str] = None,
|
|
|
|
|
+ streaming: bool = False
|
|
|
) -> dict:
|
|
) -> dict:
|
|
|
"""
|
|
"""
|
|
|
处理单个输入(文件或目录)
|
|
处理单个输入(文件或目录)
|
|
@@ -85,15 +88,32 @@ def process_single_input(
|
|
|
debug: 是否开启debug模式
|
|
debug: 是否开启debug模式
|
|
|
scene: 场景类型覆盖
|
|
scene: 场景类型覆盖
|
|
|
page_range: 页面范围(如 "1-5,7,9-12")
|
|
page_range: 页面范围(如 "1-5,7,9-12")
|
|
|
|
|
+ streaming: 是否使用流式处理模式(按页处理,立即保存,节省内存)
|
|
|
|
|
|
|
|
Returns:
|
|
Returns:
|
|
|
处理结果和输出路径
|
|
处理结果和输出路径
|
|
|
"""
|
|
"""
|
|
|
try:
|
|
try:
|
|
|
- # 初始化处理流水线
|
|
|
|
|
- with EnhancedDocPipeline(str(config_path)) as pipeline:
|
|
|
|
|
|
|
+ # 选择处理模式
|
|
|
|
|
+ if streaming:
|
|
|
|
|
+ logger.info("🔄 Using streaming processing mode (memory-efficient)")
|
|
|
|
|
+ pipeline_streaming = StreamingDocPipeline(str(config_path), str(output_dir))
|
|
|
|
|
+ use_context = False # StreamingDocPipeline 不使用 context manager
|
|
|
|
|
+ else:
|
|
|
|
|
+ logger.info("🔄 Using batch processing mode (all pages in memory)")
|
|
|
|
|
+ pipeline_batch = EnhancedDocPipeline(str(config_path))
|
|
|
|
|
+ use_context = hasattr(pipeline_batch, '__enter__')
|
|
|
|
|
+ if use_context:
|
|
|
|
|
+ pipeline_batch = pipeline_batch.__enter__()
|
|
|
|
|
+
|
|
|
|
|
+ try:
|
|
|
|
|
|
|
|
# 覆盖场景设置
|
|
# 覆盖场景设置
|
|
|
|
|
+ if streaming:
|
|
|
|
|
+ pipeline = pipeline_streaming
|
|
|
|
|
+ else:
|
|
|
|
|
+ pipeline = pipeline_batch
|
|
|
|
|
+
|
|
|
if scene:
|
|
if scene:
|
|
|
pipeline.scene_name = scene
|
|
pipeline.scene_name = scene
|
|
|
logger.info(f"🔄 Scene overridden to: {scene}")
|
|
logger.info(f"🔄 Scene overridden to: {scene}")
|
|
@@ -104,38 +124,70 @@ def process_single_input(
|
|
|
if page_range:
|
|
if page_range:
|
|
|
logger.info(f"📄 页面范围: {page_range}")
|
|
logger.info(f"📄 页面范围: {page_range}")
|
|
|
|
|
|
|
|
- # 处理文档
|
|
|
|
|
- start_time = datetime.now()
|
|
|
|
|
- results = pipeline.process_document(str(input_path), page_range=page_range)
|
|
|
|
|
- process_time = (datetime.now() - start_time).total_seconds()
|
|
|
|
|
-
|
|
|
|
|
- logger.info(f"⏱️ 处理耗时: {process_time:.2f}秒")
|
|
|
|
|
-
|
|
|
|
|
# 构建输出配置
|
|
# 构建输出配置
|
|
|
output_config = {
|
|
output_config = {
|
|
|
'save_json': True,
|
|
'save_json': True,
|
|
|
'save_markdown': True,
|
|
'save_markdown': True,
|
|
|
'save_html': True,
|
|
'save_html': True,
|
|
|
|
|
+ 'save_page_json': True,
|
|
|
|
|
+ 'save_images': True,
|
|
|
'save_layout_image': debug,
|
|
'save_layout_image': debug,
|
|
|
'save_ocr_image': debug,
|
|
'save_ocr_image': debug,
|
|
|
- 'draw_type_label': True,
|
|
|
|
|
- 'draw_bbox_number': True,
|
|
|
|
|
|
|
+ 'normalize_numbers': True,
|
|
|
|
|
+ 'merge_cross_page_tables': True,
|
|
|
|
|
+ 'cleanup_temp_files': True,
|
|
|
}
|
|
}
|
|
|
|
|
|
|
|
- # 格式化输出
|
|
|
|
|
- logger.info("💾 保存结果...")
|
|
|
|
|
- formatter = OutputFormatterV2(str(output_dir))
|
|
|
|
|
- output_paths = formatter.save_results(results, output_config)
|
|
|
|
|
-
|
|
|
|
|
- # 打印摘要
|
|
|
|
|
- _print_summary(results, output_paths, process_time)
|
|
|
|
|
|
|
+ # 处理文档
|
|
|
|
|
+ start_time = datetime.now()
|
|
|
|
|
|
|
|
- return {
|
|
|
|
|
- 'success': True,
|
|
|
|
|
- 'results': results,
|
|
|
|
|
- 'output_paths': output_paths,
|
|
|
|
|
- 'process_time': process_time
|
|
|
|
|
- }
|
|
|
|
|
|
|
+ if streaming:
|
|
|
|
|
+ # 流式处理模式
|
|
|
|
|
+ results = pipeline.process_document_streaming( # type: ignore
|
|
|
|
|
+ str(input_path),
|
|
|
|
|
+ page_range=page_range,
|
|
|
|
|
+ output_config=output_config
|
|
|
|
|
+ )
|
|
|
|
|
+ process_time = (datetime.now() - start_time).total_seconds()
|
|
|
|
|
+
|
|
|
|
|
+ # 流式模式已经保存了所有结果,只需要返回摘要
|
|
|
|
|
+ output_paths = results.get('output_paths', {})
|
|
|
|
|
+
|
|
|
|
|
+ # 打印摘要
|
|
|
|
|
+ _print_summary_streaming(results, process_time)
|
|
|
|
|
+
|
|
|
|
|
+ return {
|
|
|
|
|
+ 'success': True,
|
|
|
|
|
+ 'results': results,
|
|
|
|
|
+ 'output_paths': output_paths,
|
|
|
|
|
+ 'process_time': process_time
|
|
|
|
|
+ }
|
|
|
|
|
+ else:
|
|
|
|
|
+ # 批量处理模式(原有逻辑)
|
|
|
|
|
+ results = pipeline.process_document(str(input_path), page_range=page_range)
|
|
|
|
|
+ process_time = (datetime.now() - start_time).total_seconds()
|
|
|
|
|
+
|
|
|
|
|
+ logger.info(f"⏱️ 处理耗时: {process_time:.2f}秒")
|
|
|
|
|
+
|
|
|
|
|
+ # 格式化输出
|
|
|
|
|
+ logger.info("💾 保存结果...")
|
|
|
|
|
+ formatter = OutputFormatterV2(str(output_dir))
|
|
|
|
|
+ output_paths = formatter.save_results(results, output_config)
|
|
|
|
|
+
|
|
|
|
|
+ # 打印摘要
|
|
|
|
|
+ _print_summary(results, output_paths, process_time)
|
|
|
|
|
+
|
|
|
|
|
+ return {
|
|
|
|
|
+ 'success': True,
|
|
|
|
|
+ 'results': results,
|
|
|
|
|
+ 'output_paths': output_paths,
|
|
|
|
|
+ 'process_time': process_time
|
|
|
|
|
+ }
|
|
|
|
|
+
|
|
|
|
|
+ finally:
|
|
|
|
|
+ # 关闭context manager
|
|
|
|
|
+ if not streaming and use_context:
|
|
|
|
|
+ pipeline_batch.__exit__(None, None, None)
|
|
|
|
|
|
|
|
except Exception as e:
|
|
except Exception as e:
|
|
|
logger.error(f"❌ 处理失败: {e}")
|
|
logger.error(f"❌ 处理失败: {e}")
|
|
@@ -187,6 +239,28 @@ def _print_summary(results: dict, output_paths: dict, process_time: float):
|
|
|
print(f"{'='*60}\n")
|
|
print(f"{'='*60}\n")
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
+def _print_summary_streaming(results_summary: dict, process_time: float):
|
|
|
|
|
+ """打印流式处理结果摘要"""
|
|
|
|
|
+ print(f"\n{'='*60}")
|
|
|
|
|
+ print(f"📊 处理摘要(流式模式)")
|
|
|
|
|
+ print(f"{'='*60}")
|
|
|
|
|
+ print(f" 📄 文档: {results_summary.get('document_path', 'N/A')}")
|
|
|
|
|
+ print(f" 🎯 场景: {results_summary.get('scene', 'N/A')}")
|
|
|
|
|
+ print(f" 📋 PDF类型: {results_summary.get('metadata', {}).get('pdf_type', 'N/A')}")
|
|
|
|
|
+ print(f" 📖 页面数: {results_summary.get('total_pages', 0)}")
|
|
|
|
|
+ print(f" ⏱️ 耗时: {process_time:.2f}秒")
|
|
|
|
|
+ print(f"{'='*60}")
|
|
|
|
|
+ print(f"📁 输出文件:")
|
|
|
|
|
+ output_paths = results_summary.get('output_paths', {})
|
|
|
|
|
+ if output_paths.get('middle_json'):
|
|
|
|
|
+ print(f" - {output_paths['middle_json']}")
|
|
|
|
|
+ if output_paths.get('json_pages'):
|
|
|
|
|
+ print(f" - {len(output_paths['json_pages'])} 个页面JSON文件")
|
|
|
|
|
+ if output_paths.get('images'):
|
|
|
|
|
+ print(f" - {len(output_paths['images'])} 个图片文件")
|
|
|
|
|
+ print(f"{'='*60}\n")
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
def main():
|
|
def main():
|
|
|
parser = argparse.ArgumentParser(
|
|
parser = argparse.ArgumentParser(
|
|
|
description="金融文档处理工具 v2",
|
|
description="金融文档处理工具 v2",
|
|
@@ -210,6 +284,9 @@ def main():
|
|
|
python main_v2.py -i doc.pdf -c config.yaml -p 3,7,10 # 处理第3、7、10页
|
|
python main_v2.py -i doc.pdf -c config.yaml -p 3,7,10 # 处理第3、7、10页
|
|
|
python main_v2.py -i doc.pdf -c config.yaml -p 1-5,8-10 # 处理第1-5、8-10页
|
|
python main_v2.py -i doc.pdf -c config.yaml -p 1-5,8-10 # 处理第1-5、8-10页
|
|
|
python main_v2.py -i doc.pdf -c config.yaml -p 5- # 从第5页到最后
|
|
python main_v2.py -i doc.pdf -c config.yaml -p 5- # 从第5页到最后
|
|
|
|
|
+
|
|
|
|
|
+ # 使用流式处理模式(节省内存,适合大文档)
|
|
|
|
|
+ python main_v2.py -i large_doc.pdf -c config.yaml --streaming
|
|
|
"""
|
|
"""
|
|
|
)
|
|
)
|
|
|
|
|
|
|
@@ -257,6 +334,11 @@ def main():
|
|
|
"--pages", "-p",
|
|
"--pages", "-p",
|
|
|
help="页面范围(PDF按页码,图片目录按排序位置),如: 1-5,7,9-12"
|
|
help="页面范围(PDF按页码,图片目录按排序位置),如: 1-5,7,9-12"
|
|
|
)
|
|
)
|
|
|
|
|
+ parser.add_argument(
|
|
|
|
|
+ "--streaming",
|
|
|
|
|
+ action="store_true",
|
|
|
|
|
+ help="使用流式处理模式(按页处理,立即保存,节省内存,适合大文档)"
|
|
|
|
|
+ )
|
|
|
|
|
|
|
|
args = parser.parse_args()
|
|
args = parser.parse_args()
|
|
|
|
|
|
|
@@ -287,7 +369,8 @@ def main():
|
|
|
output_dir=Path(args.output_dir),
|
|
output_dir=Path(args.output_dir),
|
|
|
debug=args.debug,
|
|
debug=args.debug,
|
|
|
scene=args.scene,
|
|
scene=args.scene,
|
|
|
- page_range=args.pages
|
|
|
|
|
|
|
+ page_range=args.pages,
|
|
|
|
|
+ streaming=args.streaming
|
|
|
)
|
|
)
|
|
|
|
|
|
|
|
return 0 if result.get('success') else 1
|
|
return 0 if result.get('success') else 1
|
|
@@ -332,9 +415,11 @@ if __name__ == "__main__":
|
|
|
"scene": "bank_statement",
|
|
"scene": "bank_statement",
|
|
|
|
|
|
|
|
# 页面范围(可选)
|
|
# 页面范围(可选)
|
|
|
- # "pages": "3", # 只处理前5页
|
|
|
|
|
|
|
+ "pages": "4", # 只处理前4页
|
|
|
# "pages": "1-3,5,7-10", # 处理指定页面
|
|
# "pages": "1-3,5,7-10", # 处理指定页面
|
|
|
-
|
|
|
|
|
|
|
+
|
|
|
|
|
+ "streaming": True,
|
|
|
|
|
+
|
|
|
# Debug模式
|
|
# Debug模式
|
|
|
"debug": True,
|
|
"debug": True,
|
|
|
|
|
|