Explorar el Código

feat: 添加对PDF文件的支持,更新输入文件处理逻辑

zhch158_admin hace 2 meses
padre
commit
338d06ddc1
Se han modificado 1 ficheros con 147 adiciones y 41 borrados
  1. 147 41
      zhch/OmniDocBench_DotsOCR_multthreads.py

+ 147 - 41
zhch/OmniDocBench_DotsOCR_multthreads.py

@@ -2,7 +2,7 @@
 批量处理 OmniDocBench 图片并生成符合评测要求的预测结果
 
 根据 OmniDocBench 评测要求:
-- 输入:OpenDataLab___OmniDocBench/images 下的所有 .jpg 图片
+- 输入:OpenDataLab___OmniDocBench/images 下的所有 .jpg 图片,以及PDF文件
 - 输出:每个图片对应的 .md、.json 和带标注的 layout 图片文件
 - 输出目录:用于后续的 end2end 评测
 """
@@ -26,6 +26,7 @@ import argparse
 from dots_ocr.parser import DotsOCRParser
 from dots_ocr.utils import dict_promptmode_to_prompt
 from dots_ocr.utils.consts import MIN_PIXELS, MAX_PIXELS
+from dots_ocr.utils.doc_utils import load_images_from_pdf
 
 # 导入工具函数
 from utils import (
@@ -35,6 +36,119 @@ from utils import (
     collect_pid_files
 )
 
+def convert_pdf_to_images(pdf_file: str, output_dir: str | None = None, dpi: int = 200) -> List[str]:
+    """
+    将PDF转换为图像文件
+    
+    Args:
+        pdf_file: PDF文件路径
+        output_dir: 输出目录
+        dpi: 图像分辨率
+        
+    Returns:
+        生成的图像文件路径列表
+    """
+    pdf_path = Path(pdf_file)
+    if not pdf_path.exists() or pdf_path.suffix.lower() != '.pdf':
+        print(f"❌ Invalid PDF file: {pdf_path}")
+        return []
+
+    # 如果没有指定输出目录,使用PDF同名目录
+    if output_dir is None:
+        output_path = pdf_path.parent / f"{pdf_path.stem}"
+    else:
+        output_path = Path(output_dir) / f"{pdf_path.stem}"
+    output_path = output_path.resolve()
+    output_path.mkdir(parents=True, exist_ok=True)
+
+    try:
+        # 使用utils中的函数加载PDF图像
+        images = load_images_from_pdf(str(pdf_path), dpi=dpi)
+        
+        image_paths = []
+        for i, image in enumerate(images):
+            # 生成图像文件名
+            image_filename = f"{pdf_path.stem}_page_{i+1:03d}.png"
+            image_path = output_path / image_filename
+
+            # 保存图像
+            image.save(str(image_path))
+            image_paths.append(str(image_path))
+            
+        print(f"✅ Converted {len(images)} pages from {pdf_path.name} to images")
+        return image_paths
+        
+    except Exception as e:
+        print(f"❌ Error converting PDF {pdf_path}: {e}")
+        traceback.print_exc()
+        return []
+
+def get_input_files(args) -> List[str]:
+    """
+    获取输入文件列表,统一处理PDF和图像文件
+    
+    Args:
+        args: 命令行参数
+        
+    Returns:
+        处理后的图像文件路径列表
+    """
+    input_files = []
+    
+    # 获取原始输入文件
+    if args.input_csv:
+        raw_files = get_image_files_from_csv(args.input_csv, "fail")
+    elif args.input_file_list:
+        raw_files = get_image_files_from_list(args.input_file_list)
+    elif args.input_file:
+        raw_files = [Path(args.input_file).resolve()]
+    else:
+        input_dir = Path(args.input_dir).resolve()
+        if not input_dir.exists():
+            print(f"❌ Input directory does not exist: {input_dir}")
+            return []
+        
+        # 获取所有支持的文件(图像和PDF)
+        image_extensions = ['.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.tif']
+        pdf_extensions = ['.pdf']
+        
+        raw_files = []
+        for ext in image_extensions + pdf_extensions:
+            raw_files.extend(list(input_dir.glob(f"*{ext}")))
+            raw_files.extend(list(input_dir.glob(f"*{ext.upper()}")))
+        
+        raw_files = [str(f) for f in raw_files]
+    
+    # 分别处理PDF和图像文件
+    pdf_count = 0
+    image_count = 0
+    
+    for file_path in raw_files:
+        file_path = Path(file_path)
+        
+        if file_path.suffix.lower() == '.pdf':
+            # 转换PDF为图像
+            print(f"📄 Processing PDF: {file_path.name}")
+            pdf_images = convert_pdf_to_images(
+                str(file_path), 
+                args.output_dir,
+                dpi=args.dpi
+            )
+            input_files.extend(pdf_images)
+            pdf_count += 1
+        else:
+            # 直接添加图像文件
+            if file_path.exists():
+                input_files.append(str(file_path))
+                image_count += 1
+    
+    print(f"📊 Input summary:")
+    print(f"  PDF files processed: {pdf_count}")
+    print(f"  Image files found: {image_count}")
+    print(f"  Total image files to process: {len(input_files)}")
+    
+    return input_files
+
 class DotsOCRProcessor:
     """DotsOCR 处理器"""
     
@@ -155,12 +269,6 @@ class DotsOCRProcessor:
                 except Exception as e:
                     saved_files['layout_image'] = None
             
-            # # 4. 可选:保存原始图片副本
-            # output_original_image_path = os.path.join(output_dir, f"{image_name}_original.jpg")
-            # if 'original_image_path' in result and os.path.exists(result['original_image_path']):
-            #     shutil.copy2(result['original_image_path'], output_original_image_path)
-            #     saved_files['original_image'] = output_original_image_path
-            
         except Exception as e:
             print(f"Error saving results for {image_name}: {e}")
             
@@ -186,7 +294,8 @@ class DotsOCRProcessor:
             "success": False,
             "device": f"{self.ip}:{self.port}",
             "error": None,
-            "output_files": {}
+            "output_files": {},
+            "is_pdf_page": "_page_" in Path(image_path).name  # 标记是否为PDF页面
         }
         
         try:
@@ -231,9 +340,6 @@ class DotsOCRProcessor:
                 
                 result = results[0]  # parse_image 返回单个结果的列表
                 
-                # 添加原始图片路径到结果中
-                # result['original_image_path'] = image_path
-                
                 # 保存所有结果文件到输出目录
                 saved_files = self.save_results_to_output_dir(result, image_name, output_dir)
                 
@@ -426,11 +532,12 @@ def process_images_concurrent(image_paths: List[str],
 
 def main():
     """主函数"""
-    parser = argparse.ArgumentParser(description="DotsOCR OmniDocBench Processing")
+    parser = argparse.ArgumentParser(description="DotsOCR OmniDocBench Processing with PDF Support")
     
     # 输入参数组
     input_group = parser.add_mutually_exclusive_group(required=True)
-    input_group.add_argument("--input_dir", type=str, help="Input directory")
+    input_group.add_argument("--input_file", type=str, help="Input file (supports both PDF and image file)")
+    input_group.add_argument("--input_dir", type=str, help="Input directory (supports both PDF and image files)")
     input_group.add_argument("--input_file_list", type=str, help="Input file list (one file per line)")
     input_group.add_argument("--input_csv", type=str, help="Input CSV file with image_path and status columns")
 
@@ -462,28 +569,17 @@ def main():
     args = parser.parse_args()
     
     try:
-        # 获取图像文件列表
-        if args.input_csv:
-            # 从CSV文件读取
-            image_files = get_image_files_from_csv(args.input_csv, "fail")
-            print(f"📊 Loaded {len(image_files)} files from CSV with status filter: fail")
-        elif args.input_file_list:
-            # 从文件列表读取
-            image_files = get_image_files_from_list(args.input_file_list)
-        else:
-            # 从目录读取
-            input_dir = Path(args.input_dir).resolve()
-            print(f"📁 Input dir: {input_dir}")
-            
-            if not input_dir.exists():
-                print(f"❌ Input directory does not exist: {input_dir}")
-                return 1
-
-            image_files = get_image_files_from_dir(input_dir, args.input_pattern)
+        # 获取并预处理输入文件
+        print("🔄 Preprocessing input files...")
+        image_files = get_input_files(args)
+        
+        if not image_files:
+            print("❌ No input files found or processed")
+            return 1
 
         output_dir = Path(args.output_dir).resolve()
         print(f"📁 Output dir: {output_dir}")
-        print(f"📊 Found {len(image_files)} image files")
+        print(f"📊 Found {len(image_files)} image files to process")
         
         if args.test_mode:
             image_files = image_files[:10]
@@ -530,11 +626,14 @@ def main():
         success_count = sum(1 for r in results if r.get('success', False))
         skipped_count = sum(1 for r in results if r.get('skipped', False))
         error_count = len(results) - success_count
+        pdf_page_count = sum(1 for r in results if r.get('is_pdf_page', False))
         
         print(f"\n" + "="*60)
         print(f"✅ Processing completed!")
         print(f"📊 Statistics:")
-        print(f"  Total files: {len(image_files)}")
+        print(f"  Total files processed: {len(image_files)}")
+        print(f"  PDF pages processed: {pdf_page_count}")
+        print(f"  Regular images processed: {len(image_files) - pdf_page_count}")
         print(f"  Successful: {success_count}")
         print(f"  Skipped: {skipped_count}")
         print(f"  Failed: {error_count}")
@@ -549,6 +648,8 @@ def main():
         # 保存结果统计
         stats = {
             "total_files": len(image_files),
+            "pdf_pages": pdf_page_count,
+            "regular_images": len(image_files) - pdf_page_count,
             "success_count": success_count,
             "skipped_count": skipped_count,
             "error_count": error_count,
@@ -560,6 +661,7 @@ def main():
             "server": f"{args.ip}:{args.port}",
             "model": args.model_name,
             "prompt_mode": args.prompt_mode,
+            "pdf_dpi": args.dpi,
             "timestamp": time.strftime("%Y-%m-%d %H:%M:%S")
         }
         
@@ -577,14 +679,17 @@ def main():
         print(f"💾 Results saved to: {output_file}")
 
         # 收集处理结果
-        if args.collect_results:
-            processed_files = collect_pid_files(output_file)
+        if not args.collect_results:
+            output_file_processed = Path(args.output_dir) / f"processed_files_{time.strftime('%Y%m%d_%H%M%S')}.csv"
+        else:
             output_file_processed = Path(args.collect_results).resolve()
-            with open(output_file_processed, 'w', encoding='utf-8') as f:
-                f.write("image_path,status\n")
-                for file_path, status in processed_files:
-                    f.write(f"{file_path},{status}\n")
-            print(f"💾 Processed files saved to: {output_file_processed}")
+            
+        processed_files = collect_pid_files(output_file)
+        with open(output_file_processed, 'w', encoding='utf-8') as f:
+            f.write("image_path,status\n")
+            for file_path, status in processed_files:
+                f.write(f"{file_path},{status}\n")
+        print(f"💾 Processed files saved to: {output_file_processed}")
 
         return 0
         
@@ -595,7 +700,7 @@ def main():
 
 
 if __name__ == "__main__":
-    print(f"🚀 启动DotsOCR单进程程序...")
+    print(f"🚀 启动DotsOCR统一PDF/图像处理程序...")
     print(f"🔧 CUDA_VISIBLE_DEVICES: {os.environ.get('CUDA_VISIBLE_DEVICES', 'Not set')}")
     
     if len(sys.argv) == 1:
@@ -613,6 +718,7 @@ if __name__ == "__main__":
             "prompt_mode": "prompt_layout_all_en",
             "batch_size": "1",
             "max_workers": "3",
+            "dpi": "200",
             "collect_results": "./OmniDocBench_DotsOCR_Results/processed_files.csv",
         }