|
|
@@ -5,11 +5,13 @@ PDF处理工具模块
|
|
|
- PDF加载与分类
|
|
|
- PDF文本提取
|
|
|
- 跨页表格合并
|
|
|
+- 页面范围解析与过滤
|
|
|
"""
|
|
|
-from typing import Dict, List, Any, Optional, Tuple
|
|
|
+from typing import Dict, List, Any, Optional, Tuple, Set
|
|
|
from pathlib import Path
|
|
|
from PIL import Image
|
|
|
from loguru import logger
|
|
|
+import re
|
|
|
|
|
|
# 导入 MinerU 组件
|
|
|
try:
|
|
|
@@ -30,26 +32,84 @@ class PDFUtils:
|
|
|
"""PDF处理工具类"""
|
|
|
|
|
|
@staticmethod
|
|
|
+ def parse_page_range(page_range: Optional[str], total_pages: int) -> Set[int]:
|
|
|
+ """
|
|
|
+ 解析页面范围字符串
|
|
|
+
|
|
|
+ 支持格式:
|
|
|
+ - "1-5" → {0, 1, 2, 3, 4}(页码从1开始,内部转为0-based索引)
|
|
|
+ - "3" → {2}
|
|
|
+ - "1-5,7,9-12" → {0, 1, 2, 3, 4, 6, 8, 9, 10, 11}
|
|
|
+ - "1-" → 从第1页到最后
|
|
|
+ - "-5" → 从第1页到第5页
|
|
|
+
|
|
|
+ Args:
|
|
|
+ page_range: 页面范围字符串(页码从1开始)
|
|
|
+ total_pages: 总页数
|
|
|
+
|
|
|
+ Returns:
|
|
|
+ 页面索引集合(0-based)
|
|
|
+ """
|
|
|
+ if not page_range or not page_range.strip():
|
|
|
+ return set(range(total_pages))
|
|
|
+
|
|
|
+ pages = set()
|
|
|
+ parts = page_range.replace(' ', '').split(',')
|
|
|
+
|
|
|
+ for part in parts:
|
|
|
+ part = part.strip()
|
|
|
+ if not part:
|
|
|
+ continue
|
|
|
+
|
|
|
+ if '-' in part:
|
|
|
+ # 范围格式
|
|
|
+ match = re.match(r'^(\d*)-(\d*)$', part)
|
|
|
+ if match:
|
|
|
+ start_str, end_str = match.groups()
|
|
|
+ start = int(start_str) if start_str else 1
|
|
|
+ end = int(end_str) if end_str else total_pages
|
|
|
+
|
|
|
+ # 转换为 0-based 索引
|
|
|
+ start = max(0, start - 1)
|
|
|
+ end = min(total_pages, end)
|
|
|
+
|
|
|
+ pages.update(range(start, end))
|
|
|
+ else:
|
|
|
+ # 单页
|
|
|
+ try:
|
|
|
+ page_num = int(part)
|
|
|
+ if 1 <= page_num <= total_pages:
|
|
|
+ pages.add(page_num - 1) # 转换为 0-based 索引
|
|
|
+ except ValueError:
|
|
|
+ logger.warning(f"Invalid page number: {part}")
|
|
|
+
|
|
|
+ return pages
|
|
|
+
|
|
|
+ @staticmethod
|
|
|
def load_and_classify_document(
|
|
|
document_path: Path,
|
|
|
- dpi: int = 200
|
|
|
+ dpi: int = 200,
|
|
|
+ page_range: Optional[str] = None
|
|
|
) -> Tuple[List[Dict], str, Optional[Any]]:
|
|
|
"""
|
|
|
- 加载文档并分类
|
|
|
+ 加载文档并分类,支持页面范围过滤
|
|
|
|
|
|
Args:
|
|
|
document_path: 文档路径
|
|
|
dpi: PDF渲染DPI
|
|
|
+ page_range: 页面范围字符串,如 "1-5,7,9-12"
|
|
|
+ - PDF:按页码(从1开始)
|
|
|
+ - 图片目录:按文件名排序后的位置(从1开始)
|
|
|
|
|
|
Returns:
|
|
|
(images_list, pdf_type, pdf_doc)
|
|
|
- - images_list: 图像列表,每个元素包含 {'img_pil': PIL.Image, 'scale': float}
|
|
|
+ - images_list: 图像列表,每个元素包含 {'img_pil': PIL.Image, 'scale': float, 'page_idx': int}
|
|
|
- pdf_type: 'ocr' 或 'txt'
|
|
|
- pdf_doc: PDF文档对象(如果是PDF)
|
|
|
"""
|
|
|
pdf_doc = None
|
|
|
pdf_type = 'ocr' # 默认使用OCR模式
|
|
|
- images = []
|
|
|
+ all_images = []
|
|
|
|
|
|
if document_path.is_dir():
|
|
|
# 处理目录:遍历所有图片
|
|
|
@@ -59,14 +119,26 @@ class PDFUtils:
|
|
|
if f.suffix.lower() in image_extensions
|
|
|
])
|
|
|
|
|
|
- for img_file in image_files:
|
|
|
+ # 解析页面范围
|
|
|
+ total_pages = len(image_files)
|
|
|
+ selected_pages = PDFUtils.parse_page_range(page_range, total_pages)
|
|
|
+
|
|
|
+ if page_range:
|
|
|
+ logger.info(f"📋 图片目录共 {total_pages} 张,选择处理 {len(selected_pages)} 张")
|
|
|
+
|
|
|
+ for idx, img_file in enumerate(image_files):
|
|
|
+ if idx not in selected_pages:
|
|
|
+ continue
|
|
|
+
|
|
|
img = Image.open(img_file)
|
|
|
if img.mode != 'RGB':
|
|
|
img = img.convert('RGB')
|
|
|
- images.append({
|
|
|
+ all_images.append({
|
|
|
'img_pil': img,
|
|
|
'scale': 1.0,
|
|
|
- 'source_path': str(img_file)
|
|
|
+ 'source_path': str(img_file),
|
|
|
+ 'page_idx': idx, # 原始索引
|
|
|
+ 'page_name': img_file.stem # 文件名(不含扩展名)
|
|
|
})
|
|
|
|
|
|
pdf_type = 'ocr' # 图片目录始终使用OCR模式
|
|
|
@@ -90,11 +162,22 @@ class PDFUtils:
|
|
|
image_type=ImageType.PIL
|
|
|
)
|
|
|
|
|
|
- for img_dict in images_list:
|
|
|
- images.append({
|
|
|
+ # 解析页面范围
|
|
|
+ total_pages = len(images_list)
|
|
|
+ selected_pages = PDFUtils.parse_page_range(page_range, total_pages)
|
|
|
+
|
|
|
+ if page_range:
|
|
|
+ logger.info(f"📋 PDF 共 {total_pages} 页,选择处理 {len(selected_pages)} 页")
|
|
|
+
|
|
|
+ for idx, img_dict in enumerate(images_list):
|
|
|
+ if idx not in selected_pages:
|
|
|
+ continue
|
|
|
+
|
|
|
+ all_images.append({
|
|
|
'img_pil': img_dict['img_pil'],
|
|
|
'scale': img_dict.get('scale', dpi / 72),
|
|
|
- 'source_path': str(document_path)
|
|
|
+ 'source_path': str(document_path),
|
|
|
+ 'page_idx': idx # 原始页码索引
|
|
|
})
|
|
|
|
|
|
elif document_path.suffix.lower() in ['.png', '.jpg', '.jpeg', '.bmp', '.tiff', '.tif']:
|
|
|
@@ -102,17 +185,19 @@ class PDFUtils:
|
|
|
img = Image.open(document_path)
|
|
|
if img.mode != 'RGB':
|
|
|
img = img.convert('RGB')
|
|
|
- images.append({
|
|
|
+ all_images.append({
|
|
|
'img_pil': img,
|
|
|
'scale': 1.0,
|
|
|
- 'source_path': str(document_path)
|
|
|
+ 'source_path': str(document_path),
|
|
|
+ 'page_idx': 0,
|
|
|
+ 'page_name': document_path.stem
|
|
|
})
|
|
|
pdf_type = 'ocr'
|
|
|
|
|
|
else:
|
|
|
raise ValueError(f"Unsupported file format: {document_path.suffix}")
|
|
|
|
|
|
- return images, pdf_type, pdf_doc
|
|
|
+ return all_images, pdf_type, pdf_doc
|
|
|
|
|
|
@staticmethod
|
|
|
def extract_text_from_pdf(
|