|
|
@@ -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",
|
|
|
}
|
|
|
|