|
|
@@ -53,6 +53,76 @@ except ImportError:
|
|
|
from utils import OutputFormatterV2
|
|
|
|
|
|
|
|
|
+# ==================== Helper Functions ====================
|
|
|
+
|
|
|
+def _print_environment_info():
|
|
|
+ """打印环境变量信息"""
|
|
|
+ env_vars = [
|
|
|
+ 'CUDA_VISIBLE_DEVICES', 'HF_HOME', 'HF_ENDPOINT', 'HF_HUB_OFFLINE',
|
|
|
+ 'TORCH_HOME', 'MODELSCOPE_CACHE', 'USE_MODELSCOPE_HUB', 'MINERU_MODEL_SOURCE'
|
|
|
+ ]
|
|
|
+ for var in env_vars:
|
|
|
+ print(f"🔧 {var}: {os.environ.get(var, 'Not set')}")
|
|
|
+
|
|
|
+
|
|
|
+def _validate_arguments(args: argparse.Namespace) -> bool:
|
|
|
+ """验证命令行参数"""
|
|
|
+ input_path = Path(args.input)
|
|
|
+ if not input_path.exists():
|
|
|
+ logger.error(f"❌ 输入路径不存在: {input_path}")
|
|
|
+ return False
|
|
|
+
|
|
|
+ config_path = Path(args.config)
|
|
|
+ if not config_path.exists():
|
|
|
+ logger.error(f"❌ 配置文件不存在: {config_path}")
|
|
|
+ return False
|
|
|
+
|
|
|
+ return True
|
|
|
+
|
|
|
+
|
|
|
+def _handle_dry_run(args: argparse.Namespace) -> bool:
|
|
|
+ """处理dry run模式"""
|
|
|
+ if args.dry_run:
|
|
|
+ if _validate_arguments(args):
|
|
|
+ logger.info("✅ 配置验证通过(dry run)")
|
|
|
+ return True
|
|
|
+ return False
|
|
|
+ return False
|
|
|
+
|
|
|
+
|
|
|
+def _create_pipeline(streaming: bool, config_path: str, output_dir: str):
|
|
|
+ """创建并初始化处理流水线"""
|
|
|
+ if streaming:
|
|
|
+ logger.info("🔄 Using streaming processing mode (memory-efficient)")
|
|
|
+ pipeline = StreamingDocPipeline(config_path, output_dir)
|
|
|
+ else:
|
|
|
+ logger.info("🔄 Using batch processing mode (all pages in memory)")
|
|
|
+ pipeline = EnhancedDocPipeline(config_path)
|
|
|
+
|
|
|
+ return pipeline
|
|
|
+
|
|
|
+
|
|
|
+def _get_default_output_config(debug: bool) -> dict:
|
|
|
+ """获取默认输出配置"""
|
|
|
+ return {
|
|
|
+ 'create_subdir': True,
|
|
|
+ 'save_pdf_images': False,
|
|
|
+ 'save_json': True,
|
|
|
+ 'save_markdown': True,
|
|
|
+ 'save_html': True,
|
|
|
+ 'save_page_json': True,
|
|
|
+ 'save_images': True,
|
|
|
+ 'save_layout_image': debug,
|
|
|
+ 'save_ocr_image': debug,
|
|
|
+ 'draw_type_label': True,
|
|
|
+ 'draw_bbox_number': True,
|
|
|
+ 'save_enhanced_json': True,
|
|
|
+ 'coordinate_precision': 2,
|
|
|
+ 'normalize_numbers': True,
|
|
|
+ 'merge_cross_page_tables': True,
|
|
|
+ }
|
|
|
+
|
|
|
+
|
|
|
def setup_logging(log_level: str = "INFO", log_file: Optional[str] = None):
|
|
|
"""设置日志"""
|
|
|
logger.remove()
|
|
|
@@ -99,26 +169,15 @@ def process_single_input(
|
|
|
处理结果和输出路径
|
|
|
"""
|
|
|
try:
|
|
|
- # 选择处理模式
|
|
|
- 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__()
|
|
|
+ # 创建流水线
|
|
|
+ pipeline = _create_pipeline(streaming, str(config_path), str(output_dir))
|
|
|
+ output_config = pipeline.config.get('output', {}) or _get_default_output_config(debug)
|
|
|
+
|
|
|
+ use_context = not streaming and hasattr(pipeline, '__enter__')
|
|
|
+ if use_context:
|
|
|
+ pipeline = pipeline.__enter__()
|
|
|
|
|
|
try:
|
|
|
-
|
|
|
- # 覆盖场景设置
|
|
|
- if streaming:
|
|
|
- pipeline = pipeline_streaming
|
|
|
- else:
|
|
|
- pipeline = pipeline_batch
|
|
|
-
|
|
|
if scene:
|
|
|
pipeline.scene_name = scene
|
|
|
logger.info(f"🔄 Scene overridden to: {scene}")
|
|
|
@@ -129,62 +188,37 @@ def process_single_input(
|
|
|
if page_range:
|
|
|
logger.info(f"📄 页面范围: {page_range}")
|
|
|
|
|
|
- # 构建输出配置
|
|
|
- output_config = {
|
|
|
- 'save_json': True,
|
|
|
- 'save_markdown': True,
|
|
|
- 'save_html': True,
|
|
|
- 'save_page_json': True,
|
|
|
- 'save_images': True,
|
|
|
- 'save_layout_image': debug,
|
|
|
- 'save_ocr_image': debug,
|
|
|
- 'normalize_numbers': True,
|
|
|
- 'merge_cross_page_tables': True,
|
|
|
- 'cleanup_temp_files': True,
|
|
|
- }
|
|
|
-
|
|
|
- # 处理文档
|
|
|
start_time = datetime.now()
|
|
|
|
|
|
if streaming:
|
|
|
# 流式处理模式
|
|
|
- results = pipeline.process_document_streaming( # type: ignore
|
|
|
+ results = pipeline.process_document_streaming(
|
|
|
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,
|
|
|
+ 'output_paths': results.get('output_paths', {}),
|
|
|
'process_time': process_time
|
|
|
}
|
|
|
else:
|
|
|
- # 批量处理模式(原有逻辑)
|
|
|
- # 批量处理模式(原有逻辑)
|
|
|
+ # 批量处理模式
|
|
|
results = pipeline.process_document(
|
|
|
str(input_path),
|
|
|
page_range=page_range,
|
|
|
output_dir=str(output_dir)
|
|
|
)
|
|
|
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 {
|
|
|
@@ -195,9 +229,8 @@ def process_single_input(
|
|
|
}
|
|
|
|
|
|
finally:
|
|
|
- # 关闭context manager
|
|
|
- if not streaming and use_context:
|
|
|
- pipeline_batch.__exit__(None, None, None)
|
|
|
+ if use_context:
|
|
|
+ pipeline.__exit__(None, None, None)
|
|
|
|
|
|
except Exception as e:
|
|
|
logger.error(f"❌ 处理失败: {e}")
|
|
|
@@ -352,30 +385,17 @@ def main():
|
|
|
|
|
|
args = parser.parse_args()
|
|
|
|
|
|
- # 设置日志
|
|
|
setup_logging(args.log_level, args.log_file)
|
|
|
|
|
|
- # 验证输入
|
|
|
- input_path = Path(args.input)
|
|
|
- if not input_path.exists():
|
|
|
- logger.error(f"❌ 输入路径不存在: {input_path}")
|
|
|
- return 1
|
|
|
+ if _handle_dry_run(args):
|
|
|
+ return 0
|
|
|
|
|
|
- # 验证配置文件
|
|
|
- config_path = Path(args.config)
|
|
|
- if not config_path.exists():
|
|
|
- logger.error(f"❌ 配置文件不存在: {config_path}")
|
|
|
+ if not _validate_arguments(args):
|
|
|
return 1
|
|
|
|
|
|
- # 仅验证模式
|
|
|
- if args.dry_run:
|
|
|
- logger.info("✅ 配置验证通过(dry run)")
|
|
|
- return 0
|
|
|
-
|
|
|
- # 处理文档
|
|
|
result = process_single_input(
|
|
|
- input_path=input_path,
|
|
|
- config_path=config_path,
|
|
|
+ input_path=Path(args.input),
|
|
|
+ config_path=Path(args.config),
|
|
|
output_dir=Path(args.output_dir),
|
|
|
debug=args.debug,
|
|
|
scene=args.scene,
|
|
|
@@ -387,21 +407,12 @@ def main():
|
|
|
|
|
|
|
|
|
if __name__ == "__main__":
|
|
|
- # 打印环境变量
|
|
|
- print(f"🔧 CUDA_VISIBLE_DEVICES: {os.environ.get('CUDA_VISIBLE_DEVICES', 'Not set')}")
|
|
|
- print(f"🔧 HF_HOME: {os.environ.get('HF_HOME', 'Not set')}")
|
|
|
- print(f"🔧 HF_ENDPOINT: {os.environ.get('HF_ENDPOINT', 'Not set')}")
|
|
|
- print(f"🔧 HF_HUB_OFFLINE: {os.environ.get('HF_HUB_OFFLINE', 'Not set')}")
|
|
|
- print(f"🔧 TORCH_HOME: {os.environ.get('TORCH_HOME', 'Not set')}")
|
|
|
- print(f"🔧 MODELSCOPE_CACHE: {os.environ.get('MODELSCOPE_CACHE', 'Not set')}")
|
|
|
- print(f"🔧 USE_MODELSCOPE_HUB: {os.environ.get('USE_MODELSCOPE_HUB', 'Not set')}")
|
|
|
- print(f"🔧 MINERU_MODEL_SOURCE: {os.environ.get('MINERU_MODEL_SOURCE', 'Not set')}")
|
|
|
+ _print_environment_info()
|
|
|
|
|
|
if len(sys.argv) == 1:
|
|
|
- # 没有命令行参数时,使用默认配置运行
|
|
|
print("ℹ️ 未提供命令行参数,使用默认配置运行...")
|
|
|
|
|
|
- # 默认配置
|
|
|
+ # 默认配置(用于开发测试)
|
|
|
default_config = {
|
|
|
# 测试输入
|
|
|
# "input": "/Users/zhch158/workspace/data/流水分析/康强_北京农村商业银行.pdf",
|
|
|
@@ -414,15 +425,15 @@ if __name__ == "__main__":
|
|
|
# "output_dir": "/Users/zhch158/workspace/data/流水分析/B用户_扫描流水/bank_statement_yusys_v2",
|
|
|
|
|
|
# "input": "/Users/zhch158/workspace/repository.git/ocr_platform/ocr_tools/universal_doc_parser/tests/2023年度报告母公司_page_005.png",
|
|
|
- # "input": "/Users/zhch158/workspace/repository.git/ocr_platform/ocr_tools/universal_doc_parser/tests/2023年度报告母公司_page_003_270.png",
|
|
|
+ # "input": "/Users/zhch158/workspace/repository.git/ocr_platform/ocr_tools/universal_doc_parser/tests/2023年度报告母公司_page_003.png",
|
|
|
# "input": "/Users/zhch158/workspace/repository.git/ocr_platform/ocr_tools/universal_doc_parser/tests/2023年度报告母公司_page_003_270_skew(-0.4).png",
|
|
|
- "input": "/Users/zhch158/workspace/data/流水分析/2023年度报告母公司.pdf",
|
|
|
- "output_dir": "./output/2023年度报告母公司/bank_statement_wired_unet",
|
|
|
+ # "input": "/Users/zhch158/workspace/data/流水分析/2023年度报告母公司.pdf",
|
|
|
+ # "output_dir": "./output/2023年度报告母公司/bank_statement_wired_unet",
|
|
|
|
|
|
# "input": "/Users/zhch158/workspace/data/流水分析/2023年度报告母公司.pdf",
|
|
|
# "output_dir": "/Users/zhch158/workspace/data/流水分析/2023年度报告母公司/bank_statement_yusys_v2",
|
|
|
|
|
|
- # "input": "/Users/zhch158/workspace/data/流水分析/A用户_单元格扫描流水.pdf",
|
|
|
+ # # "input": "/Users/zhch158/workspace/data/流水分析/A用户_单元格扫描流水.pdf",
|
|
|
# "output_dir": "/Users/zhch158/workspace/data/流水分析/A用户_单元格扫描流水/bank_statement_wired_unet",
|
|
|
|
|
|
# "input": "/Users/zhch158/workspace/repository.git/ocr_platform/ocr_tools/universal_doc_parser/tests/600916_中国黄金_2022年报_page_096.png",
|
|
|
@@ -430,6 +441,10 @@ if __name__ == "__main__":
|
|
|
# "input": "/Users/zhch158/workspace/data/流水分析/600916_中国黄金_2022年报.pdf",
|
|
|
# "output_dir": "./output/600916_中国黄金_2022年报/bank_statement_wired_unet",
|
|
|
|
|
|
+
|
|
|
+ # "input": "/Users/zhch158/workspace/repository.git/ocr_platform/ocr_tools/universal_doc_parser/tests/提取自赤峰黄金2023年报.pdf",
|
|
|
+ # "output_dir": "./output/提取自赤峰黄金2023年报/bank_statement_wired_unet",
|
|
|
+
|
|
|
# "input": "/Users/zhch158/workspace/data/流水分析/施博深.pdf",
|
|
|
# "output_dir": "/Users/zhch158/workspace/data/流水分析/施博深/bank_statement_yusys_v2",
|
|
|
|
|
|
@@ -439,6 +454,9 @@ if __name__ == "__main__":
|
|
|
# "input": "/Users/zhch158/workspace/data/流水分析/施博深.wiredtable",
|
|
|
# "output_dir": "/Users/zhch158/workspace/data/流水分析/施博深/bank_statement_wired_unet",
|
|
|
|
|
|
+ "input": "/Users/zhch158/workspace/data/流水分析/山西云集科技有限公司.pdf",
|
|
|
+ "output_dir": "/Users/zhch158/workspace/data/流水分析/山西云集科技有限公司/bank_statement_wired_unet",
|
|
|
+
|
|
|
# 配置文件
|
|
|
"config": "./config/bank_statement_wired_unet.yaml",
|
|
|
# "config": "./config/bank_statement_yusys_v2.yaml",
|
|
|
@@ -448,7 +466,7 @@ if __name__ == "__main__":
|
|
|
"scene": "bank_statement",
|
|
|
|
|
|
# 页面范围(可选)
|
|
|
- # "pages": "6", # 只处理前1页
|
|
|
+ # "pages": "12,26,27", # 只处理前1页
|
|
|
# "pages": "1-3,5,7-10", # 处理指定页面
|
|
|
# "pages": "83-109", # 处理指定页面
|
|
|
|
|
|
@@ -459,6 +477,9 @@ if __name__ == "__main__":
|
|
|
|
|
|
# 日志级别
|
|
|
"log_level": "DEBUG",
|
|
|
+
|
|
|
+ # 日志文件
|
|
|
+ "log_file": "./logs/bank_statement_wired_unet/process.log",
|
|
|
}
|
|
|
|
|
|
# 构造参数
|