|
|
@@ -0,0 +1,292 @@
|
|
|
+import tempfile
|
|
|
+from pathlib import Path
|
|
|
+from typing import List, Tuple
|
|
|
+import json
|
|
|
+from .doc_utils import load_images_from_pdf
|
|
|
+import traceback
|
|
|
+
|
|
|
+def split_files(file_list: List[str], num_splits: int) -> List[List[str]]:
|
|
|
+ """
|
|
|
+ 将文件列表分割成指定数量的子列表
|
|
|
+
|
|
|
+ Args:
|
|
|
+ file_list: 文件路径列表
|
|
|
+ num_splits: 分割数量
|
|
|
+
|
|
|
+ Returns:
|
|
|
+ 分割后的文件列表
|
|
|
+ """
|
|
|
+ if num_splits <= 0:
|
|
|
+ return [file_list]
|
|
|
+
|
|
|
+ chunk_size = len(file_list) // num_splits
|
|
|
+ remainder = len(file_list) % num_splits
|
|
|
+
|
|
|
+ chunks = []
|
|
|
+ start = 0
|
|
|
+
|
|
|
+ for i in range(num_splits):
|
|
|
+ # 前remainder个chunk多分配一个文件
|
|
|
+ current_chunk_size = chunk_size + (1 if i < remainder else 0)
|
|
|
+ if current_chunk_size > 0:
|
|
|
+ chunks.append(file_list[start:start + current_chunk_size])
|
|
|
+ start += current_chunk_size
|
|
|
+
|
|
|
+ return [chunk for chunk in chunks if chunk] # 过滤空列表
|
|
|
+
|
|
|
+def create_temp_file_list(file_chunk: List[str]) -> str:
|
|
|
+ """
|
|
|
+ 创建临时文件列表文件
|
|
|
+
|
|
|
+ Args:
|
|
|
+ file_chunk: 文件路径列表
|
|
|
+
|
|
|
+ Returns:
|
|
|
+ 临时文件路径
|
|
|
+ """
|
|
|
+ with tempfile.NamedTemporaryFile(mode='w', suffix='.txt', delete=False) as f:
|
|
|
+ for file_path in file_chunk:
|
|
|
+ f.write(f"{file_path}\n")
|
|
|
+ return f.name
|
|
|
+
|
|
|
+def get_image_files_from_dir(input_dir: Path, pattern: str = "*", max_files: int = None) -> List[str]:
|
|
|
+ """
|
|
|
+ 从目录获取图像文件列表
|
|
|
+
|
|
|
+ Args:
|
|
|
+ input_dir: 输入目录
|
|
|
+ max_files: 最大文件数量限制
|
|
|
+
|
|
|
+ Returns:
|
|
|
+ 图像文件路径列表
|
|
|
+ """
|
|
|
+ image_extensions = ['.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.tif']
|
|
|
+ image_files = []
|
|
|
+
|
|
|
+ for ext in image_extensions:
|
|
|
+ image_files.extend(list(input_dir.glob(f"{pattern}{ext}")))
|
|
|
+ image_files.extend(list(input_dir.glob(f"{pattern}{ext.upper()}")))
|
|
|
+
|
|
|
+ # 去重并排序
|
|
|
+ image_files = sorted(list(set(str(f) for f in image_files)))
|
|
|
+
|
|
|
+ # 限制文件数量
|
|
|
+ if max_files:
|
|
|
+ image_files = image_files[:max_files]
|
|
|
+
|
|
|
+ return image_files
|
|
|
+
|
|
|
+def get_image_files_from_list(file_list_path: str) -> List[str]:
|
|
|
+ """
|
|
|
+ 从文件列表获取图像文件列表
|
|
|
+
|
|
|
+ Args:
|
|
|
+ file_list_path: 文件列表路径
|
|
|
+
|
|
|
+ Returns:
|
|
|
+ 图像文件路径列表
|
|
|
+ """
|
|
|
+ print(f"📄 Reading file list from: {file_list_path}")
|
|
|
+
|
|
|
+ with open(file_list_path, 'r', encoding='utf-8') as f:
|
|
|
+ image_files = [line.strip() for line in f if line.strip()]
|
|
|
+
|
|
|
+ # 验证文件存在性
|
|
|
+ valid_files = []
|
|
|
+ missing_files = []
|
|
|
+
|
|
|
+ for file_path in image_files:
|
|
|
+ if Path(file_path).exists():
|
|
|
+ valid_files.append(file_path)
|
|
|
+ else:
|
|
|
+ missing_files.append(file_path)
|
|
|
+
|
|
|
+ if missing_files:
|
|
|
+ print(f"⚠️ Warning: {len(missing_files)} files not found:")
|
|
|
+ for missing_file in missing_files[:5]: # 只显示前5个
|
|
|
+ print(f" - {missing_file}")
|
|
|
+ if len(missing_files) > 5:
|
|
|
+ print(f" ... and {len(missing_files) - 5} more")
|
|
|
+
|
|
|
+ print(f"✅ Found {len(valid_files)} valid files out of {len(image_files)} in list")
|
|
|
+ return valid_files
|
|
|
+
|
|
|
+def get_image_files_from_csv(csv_file: str, status_filter: str = "fail") -> List[str]:
|
|
|
+ """
|
|
|
+ 从CSV文件获取图像文件列表
|
|
|
+
|
|
|
+ Args:
|
|
|
+ csv_file: CSV文件路径
|
|
|
+ status_filter: 状态过滤器
|
|
|
+
|
|
|
+ Returns:
|
|
|
+ 图像文件路径列表
|
|
|
+ """
|
|
|
+ print(f"📄 Reading image files from CSV: {csv_file}")
|
|
|
+
|
|
|
+ # 读取CSV文件, 表头:image_path,status
|
|
|
+ image_files = []
|
|
|
+ with open(csv_file, 'r', encoding='utf-8') as f:
|
|
|
+ for line in f:
|
|
|
+ # 需要去掉表头, 按“,”分割,读取文件名,状态
|
|
|
+ image_file, status = line.strip().split(",")
|
|
|
+ if status.lower() == status_filter.lower():
|
|
|
+ image_files.append(image_file)
|
|
|
+
|
|
|
+ return image_files
|
|
|
+
|
|
|
+
|
|
|
+def collect_pid_files(pid_output_file: str) -> List[Tuple[str, str]]:
|
|
|
+ """
|
|
|
+ 从进程输出文件中收集文件
|
|
|
+
|
|
|
+ Args:
|
|
|
+ pid_output_file: 进程输出文件路径
|
|
|
+
|
|
|
+ Returns:
|
|
|
+ 文件列表(文件路径,处理结果)
|
|
|
+ """
|
|
|
+
|
|
|
+ """
|
|
|
+ 单进程结果统计文件格式
|
|
|
+ "results": [
|
|
|
+ {
|
|
|
+ "image_path": "docstructbench_dianzishu_zhongwenzaixian-o.O-61520612.pdf_140.jpg",
|
|
|
+ "processing_time": 2.0265579223632812e-06,
|
|
|
+ "success": true,
|
|
|
+ "device": "gpu:3",
|
|
|
+ "output_json": "/home/ubuntu/zhch/PaddleX/zhch/OmniDocBench_Results_Scheduler/process_3/docstructbench_dianzishu_zhongwenzaixian-o.O-61520612.pdf_140.json",
|
|
|
+ "output_md": "/home/ubuntu/zhch/PaddleX/zhch/OmniDocBench_Results_Scheduler/process_3/docstructbench_dianzishu_zhongwenzaixian-o.O-61520612.pdf_140.md"
|
|
|
+ },
|
|
|
+ ...
|
|
|
+ """
|
|
|
+ if not Path(pid_output_file).exists():
|
|
|
+ print(f"⚠️ Warning: PID output file not found: {pid_output_file}")
|
|
|
+ return []
|
|
|
+
|
|
|
+ with open(pid_output_file, 'r', encoding='utf-8') as f:
|
|
|
+ data = json.load(f)
|
|
|
+
|
|
|
+ if not isinstance(data, dict) or "results" not in data:
|
|
|
+ print(f"⚠️ Warning: Invalid PID output file format: {pid_output_file}")
|
|
|
+ return []
|
|
|
+ # 返回文件路径和处理状态, 如果“success”: True, 则状态为“success”, 否则为“fail”
|
|
|
+ file_list = []
|
|
|
+ for file_result in data.get("results", []):
|
|
|
+ image_path = file_result.get("image_path", "")
|
|
|
+ status = "success" if file_result.get("success", False) else "fail"
|
|
|
+ file_list.append((image_path, status))
|
|
|
+ return file_list
|
|
|
+
|
|
|
+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:
|
|
|
+ # 使用doc_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.pdf_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 sorted(list(set(str(f) for f in input_files)))
|