Browse Source

feat: add page range argument to process_single_input for selective document processing, enhancing user control over which pages to process

zhch158_admin 5 days ago
parent
commit
4fb6cf4eb9
1 changed files with 27 additions and 5 deletions
  1. 27 5
      zhch/universal_doc_parser/main_v2.py

+ 27 - 5
zhch/universal_doc_parser/main_v2.py

@@ -72,7 +72,8 @@ def process_single_input(
     config_path: Path,
     output_dir: Path,
     debug: bool = False,
-    scene: str = None
+    scene: str = None,
+    page_range: str = None
 ) -> dict:
     """
     处理单个输入(文件或目录)
@@ -83,6 +84,7 @@ def process_single_input(
         output_dir: 输出目录
         debug: 是否开启debug模式
         scene: 场景类型覆盖
+        page_range: 页面范围(如 "1-5,7,9-12")
         
     Returns:
         处理结果和输出路径
@@ -99,10 +101,12 @@ def process_single_input(
             logger.info(f"🚀 开始处理: {input_path}")
             logger.info(f"📋 场景配置: {pipeline.scene_name}")
             logger.info(f"📁 输出目录: {output_dir}")
+            if page_range:
+                logger.info(f"📄 页面范围: {page_range}")
             
             # 处理文档
             start_time = datetime.now()
-            results = pipeline.process_document(str(input_path))
+            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}秒")
@@ -200,6 +204,12 @@ def main():
   
   # 指定输出目录
   python main_v2.py -i doc.pdf -c config.yaml -o ./my_output/
+  
+  # 指定页面范围(PDF按页码,图片目录按排序位置)
+  python main_v2.py -i doc.pdf -c config.yaml -p 1-5      # 处理第1-5页
+  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 5-       # 从第5页到最后
         """
     )
     
@@ -243,6 +253,10 @@ def main():
         action="store_true",
         help="仅验证配置,不执行处理"
     )
+    parser.add_argument(
+        "--pages", "-p",
+        help="页面范围(PDF按页码,图片目录按排序位置),如: 1-5,7,9-12"
+    )
     
     args = parser.parse_args()
     
@@ -272,7 +286,8 @@ def main():
         config_path=config_path,
         output_dir=Path(args.output_dir),
         debug=args.debug,
-        scene=args.scene
+        scene=args.scene,
+        page_range=args.pages
     )
     
     return 0 if result.get('success') else 1
@@ -302,8 +317,11 @@ if __name__ == "__main__":
             # "input": "/Users/zhch158/workspace/data/流水分析/2023年度报告母公司/mineru_vllm_results/2023年度报告母公司/2023年度报告母公司_page_003.png",
             # "output_dir": "./output/2023年度报告母公司_bank_statement_v2",
             
-            "input": "/Users/zhch158/workspace/data/流水分析/B用户_扫描流水.pdf",
-            "output_dir": "/Users/zhch158/workspace/data/流水分析/B用户_扫描流水/bank_statement_yusys_v2",
+            # "input": "/Users/zhch158/workspace/data/流水分析/B用户_扫描流水.pdf",
+            # "output_dir": "/Users/zhch158/workspace/data/流水分析/B用户_扫描流水/bank_statement_yusys_v2",
+
+            "input": "/Users/zhch158/workspace/data/流水分析/2023年度报告母公司.pdf",
+            "output_dir": "/Users/zhch158/workspace/data/流水分析/2023年度报告母公司/bank_statement_yusys_v2",
 
             # 配置文件
             "config": "./config/bank_statement_yusys_v2.yaml",
@@ -313,6 +331,10 @@ if __name__ == "__main__":
             # 场景
             "scene": "bank_statement",
             
+            # 页面范围(可选)
+            # "pages": "3",  # 只处理前5页
+            # "pages": "1-3,5,7-10",  # 处理指定页面
+            
             # Debug模式
             "debug": True,