340 lines
12 KiB
Python
340 lines
12 KiB
Python
|
#
|
|||
|
|
# Copyright 2025 The InfiniFlow Authors. All Rights Reserved.
|
|||
|
|
#
|
|||
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|||
|
|
# you may not use this file except in compliance with the License.
|
|||
|
|
# You may obtain a copy of the License at
|
|||
|
|
#
|
|||
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|||
|
|
#
|
|||
|
|
# Unless required by applicable law or agreed to in writing, software
|
|||
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|||
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|||
|
|
# See the License for the specific language governing permissions and
|
|||
|
|
# limitations under the License.
|
|||
|
|
#
|
|||
|
|
|
|||
|
|
"""
|
|||
|
|
简化的PDF解析器，只使用OCR处理PDF文档
|
|||
|
|
|
|||
|
|
从 RAGFlow 的 RAGFlowPdfParser 中提取OCR相关功能，移除了：
|
|||
|
|
- 布局识别（Layout Recognition）
|
|||
|
|
- 表格结构识别（Table Structure Recognition）
|
|||
|
|
- 文本合并和语义分析
|
|||
|
|
- RAG相关功能
|
|||
|
|
|
|||
|
|
只保留：
|
|||
|
|
- PDF转图片
|
|||
|
|
- OCR文本检测和识别
|
|||
|
|
- 基本的文本和位置信息返回
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
import logging
|
|||
|
|
import sys
|
|||
|
|
import threading
|
|||
|
|
from io import BytesIO
|
|||
|
|
from pathlib import Path
|
|||
|
|
from timeit import default_timer as timer
|
|||
|
|
|
|||
|
|
import numpy as np
|
|||
|
|
import pdfplumber
|
|||
|
|
import trio
|
|||
|
|
|
|||
|
|
# 处理导入问题：支持直接运行和模块导入
|
|||
|
|
try:
|
|||
|
|
_package = __package__
|
|||
|
|
except NameError:
|
|||
|
|
_package = None
|
|||
|
|
|
|||
|
|
if _package is None:
|
|||
|
|
# 直接运行时，添加父目录到路径并使用绝对导入
|
|||
|
|
parent_dir = Path(__file__).parent.parent
|
|||
|
|
if str(parent_dir) not in sys.path:
|
|||
|
|
sys.path.insert(0, str(parent_dir))
|
|||
|
|
from ocr.config import PARALLEL_DEVICES
|
|||
|
|
from ocr.ocr import OCR
|
|||
|
|
else:
|
|||
|
|
# 作为模块导入时使用相对导入
|
|||
|
|
from config import PARALLEL_DEVICES
|
|||
|
|
from ocr import OCR
|
|||
|
|
|
|||
|
|
LOCK_KEY_pdfplumber = "global_shared_lock_pdfplumber"
|
|||
|
|
if LOCK_KEY_pdfplumber not in sys.modules:
|
|||
|
|
sys.modules[LOCK_KEY_pdfplumber] = threading.Lock()
|
|||
|
|
|
|||
|
|
|
|||
|
|
class SimplePdfParser:
|
|||
|
|
"""
|
|||
|
|
简化的PDF解析器，只使用OCR处理PDF
|
|||
|
|
|
|||
|
|
使用方法:
|
|||
|
|
parser = SimplePdfParser()
|
|||
|
|
result = parser.parse_pdf("file.pdf") # 或传入二进制数据
|
|||
|
|
# result 格式:
|
|||
|
|
# {
|
|||
|
|
# "pages": [
|
|||
|
|
# {
|
|||
|
|
# "page_number": 1,
|
|||
|
|
# "boxes": [
|
|||
|
|
# {
|
|||
|
|
# "text": "识别到的文本",
|
|||
|
|
# "bbox": [[x0, y0], [x1, y0], [x1, y1], [x0, y1]],
|
|||
|
|
# "confidence": 0.95
|
|||
|
|
# },
|
|||
|
|
# ...
|
|||
|
|
# ]
|
|||
|
|
# },
|
|||
|
|
# ...
|
|||
|
|
# ]
|
|||
|
|
# }
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
def __init__(self, model_dir=None):
|
|||
|
|
"""
|
|||
|
|
初始化PDF解析器
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
model_dir: OCR模型目录，如果为None则使用默认路径
|
|||
|
|
"""
|
|||
|
|
self.ocr = OCR(model_dir=model_dir)
|
|||
|
|
self.parallel_limiter = None
|
|||
|
|
if PARALLEL_DEVICES > 1:
|
|||
|
|
self.parallel_limiter = [trio.CapacityLimiter(1) for _ in range(PARALLEL_DEVICES)]
|
|||
|
|
|
|||
|
|
def __ocr_page(self, page_num, img, zoomin=3, device_id=None):
|
|||
|
|
"""
|
|||
|
|
对单页进行OCR处理
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
page_num: 页码
|
|||
|
|
img: PIL Image对象
|
|||
|
|
zoomin: 放大倍数（用于坐标缩放）
|
|||
|
|
device_id: GPU设备ID
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
list: OCR结果列表，每个元素为 {"text": str, "bbox": list, "confidence": float}
|
|||
|
|
"""
|
|||
|
|
start = timer()
|
|||
|
|
img_np = np.array(img)
|
|||
|
|
|
|||
|
|
# 文本检测
|
|||
|
|
# detect方法返回: zip对象，格式为 (box_coords, (text, score))
|
|||
|
|
# 但检测阶段text和score都是默认值，需要后续识别
|
|||
|
|
detection_result = self.ocr.detect(img_np, device_id)
|
|||
|
|
|
|||
|
|
if detection_result is None:
|
|||
|
|
return []
|
|||
|
|
|
|||
|
|
# 转换为列表并提取box坐标
|
|||
|
|
# detect返回的格式是zip，每个元素是 (box_coords, (text, score))
|
|||
|
|
# 在检测阶段，text是空字符串，score是0
|
|||
|
|
bxs = list(detection_result)
|
|||
|
|
|
|||
|
|
logging.info(f"Page {page_num}: OCR detection found {len(bxs)} boxes in {timer() - start:.2f}s")
|
|||
|
|
|
|||
|
|
if not bxs:
|
|||
|
|
return []
|
|||
|
|
|
|||
|
|
# 解析检测结果并准备识别
|
|||
|
|
boxes_to_reg = []
|
|||
|
|
|
|||
|
|
start = timer()
|
|||
|
|
for box_coords, _, _ in bxs:
|
|||
|
|
# box_coords 是四边形坐标: [[x0, y0], [x1, y0], [x1, y1], [x0, y1]]
|
|||
|
|
# 转换为原始坐标（考虑zoomin）
|
|||
|
|
box_coords_np = np.array(box_coords, dtype=np.float32)
|
|||
|
|
original_coords = box_coords_np / zoomin # 缩放回原始坐标
|
|||
|
|
|
|||
|
|
# 裁剪图像用于识别
|
|||
|
|
# 使用放大后的坐标裁剪（因为img_np是放大后的图像）
|
|||
|
|
crop_box = box_coords_np
|
|||
|
|
crop_img = self.ocr.get_rotate_crop_image(img_np, crop_box)
|
|||
|
|
boxes_to_reg.append({
|
|||
|
|
"bbox": original_coords.tolist(),
|
|||
|
|
"crop_img": crop_img
|
|||
|
|
})
|
|||
|
|
|
|||
|
|
# 批量识别文本
|
|||
|
|
ocr_results = []
|
|||
|
|
if boxes_to_reg:
|
|||
|
|
crop_imgs = [b["crop_img"] for b in boxes_to_reg]
|
|||
|
|
texts = self.ocr.recognize_batch(crop_imgs, device_id)
|
|||
|
|
|
|||
|
|
# 组装结果
|
|||
|
|
for i, b in enumerate(boxes_to_reg):
|
|||
|
|
if i < len(texts) and texts[i]: # 过滤空文本
|
|||
|
|
ocr_results.append({
|
|||
|
|
"text": texts[i],
|
|||
|
|
"bbox": b["bbox"],
|
|||
|
|
"confidence": 0.9 # 简化版本，不计算具体置信度
|
|||
|
|
})
|
|||
|
|
|
|||
|
|
logging.info(f"Page {page_num}: OCR recognition {len(ocr_results)} boxes cost {timer() - start:.2f}s")
|
|||
|
|
return ocr_results
|
|||
|
|
|
|||
|
|
async def __ocr_page_async(self, page_num, img, zoomin, device_id, limiter, callback):
|
|||
|
|
"""
|
|||
|
|
异步OCR处理单页
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
page_num: 页码
|
|||
|
|
img: PIL Image对象
|
|||
|
|
zoomin: 放大倍数
|
|||
|
|
device_id: GPU设备ID
|
|||
|
|
limiter: 并发限制器
|
|||
|
|
callback: 进度回调函数
|
|||
|
|
"""
|
|||
|
|
if limiter:
|
|||
|
|
async with limiter:
|
|||
|
|
result = await trio.to_thread.run_sync(
|
|||
|
|
lambda: self.__ocr_page(page_num, img, zoomin, device_id)
|
|||
|
|
)
|
|||
|
|
else:
|
|||
|
|
result = await trio.to_thread.run_sync(
|
|||
|
|
lambda: self.__ocr_page(page_num, img, zoomin, device_id)
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
if callback and page_num % 5 == 0:
|
|||
|
|
callback(prog=page_num * 0.9 / 100, msg=f"Processing page {page_num}...")
|
|||
|
|
|
|||
|
|
return result
|
|||
|
|
|
|||
|
|
def __convert_pdf_to_images(self, pdf_source, zoomin=3, page_from=0, page_to=299):
|
|||
|
|
"""
|
|||
|
|
将PDF转换为图片
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
pdf_source: PDF文件路径（str）或二进制数据（bytes）
|
|||
|
|
zoomin: 放大倍数，默认3（72*3=216 DPI）
|
|||
|
|
page_from: 起始页码（从0开始）
|
|||
|
|
page_to: 结束页码
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
list: PIL Image对象列表
|
|||
|
|
"""
|
|||
|
|
start = timer()
|
|||
|
|
page_images = []
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
with sys.modules[LOCK_KEY_pdfplumber]:
|
|||
|
|
pdf = pdfplumber.open(pdf_source) if isinstance(pdf_source, str) else pdfplumber.open(BytesIO(pdf_source))
|
|||
|
|
try:
|
|||
|
|
# 转换为图片，resolution = 72 * zoomin
|
|||
|
|
page_images = [
|
|||
|
|
p.to_image(resolution=72 * zoomin, antialias=True).annotated
|
|||
|
|
for i, p in enumerate(pdf.pages[page_from:page_to])
|
|||
|
|
]
|
|||
|
|
pdf.close()
|
|||
|
|
except Exception as e:
|
|||
|
|
logging.warning(f"Failed to convert PDF pages {page_from}-{page_to}: {str(e)}")
|
|||
|
|
if hasattr(pdf, 'close'):
|
|||
|
|
pdf.close()
|
|||
|
|
except Exception as e:
|
|||
|
|
logging.exception(f"Error converting PDF to images: {str(e)}")
|
|||
|
|
|
|||
|
|
logging.info(f"Converted {len(page_images)} pages to images in {timer() - start:.2f}s")
|
|||
|
|
return page_images
|
|||
|
|
|
|||
|
|
def parse_pdf(self, pdf_source, zoomin=3, page_from=0, page_to=299, callback=None):
|
|||
|
|
"""
|
|||
|
|
解析PDF文档，使用OCR识别文本
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
pdf_source: PDF文件路径（str）或二进制数据（bytes）
|
|||
|
|
zoomin: 放大倍数，默认3
|
|||
|
|
page_from: 起始页码（从0开始）
|
|||
|
|
page_to: 结束页码
|
|||
|
|
callback: 进度回调函数，格式: callback(prog: float, msg: str)
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
dict: 解析结果
|
|||
|
|
{
|
|||
|
|
"pages": [
|
|||
|
|
{
|
|||
|
|
"page_number": int,
|
|||
|
|
"boxes": [
|
|||
|
|
{
|
|||
|
|
"text": str,
|
|||
|
|
"bbox": [[x0, y0], [x1, y0], [x1, y1], [x0, y1]],
|
|||
|
|
"confidence": float
|
|||
|
|
},
|
|||
|
|
...
|
|||
|
|
]
|
|||
|
|
},
|
|||
|
|
...
|
|||
|
|
]
|
|||
|
|
}
|
|||
|
|
"""
|
|||
|
|
if callback:
|
|||
|
|
callback(0.0, "Starting PDF parsing...")
|
|||
|
|
|
|||
|
|
# 1. 转换为图片
|
|||
|
|
if callback:
|
|||
|
|
callback(0.1, "Converting PDF to images...")
|
|||
|
|
page_images = self.__convert_pdf_to_images(pdf_source, zoomin, page_from, page_to)
|
|||
|
|
|
|||
|
|
if not page_images:
|
|||
|
|
logging.warning("No pages converted from PDF")
|
|||
|
|
return {"pages": []}
|
|||
|
|
|
|||
|
|
# 2. OCR处理
|
|||
|
|
async def process_all_pages():
|
|||
|
|
pages_result = []
|
|||
|
|
|
|||
|
|
if self.parallel_limiter:
|
|||
|
|
# 并行处理（多GPU）
|
|||
|
|
async with trio.open_nursery() as nursery:
|
|||
|
|
tasks = []
|
|||
|
|
for i, img in enumerate(page_images):
|
|||
|
|
page_num = page_from + i + 1
|
|||
|
|
device_id = i % PARALLEL_DEVICES
|
|||
|
|
task = nursery.start_soon(
|
|||
|
|
self.__ocr_page_async,
|
|||
|
|
page_num, img, zoomin, device_id,
|
|||
|
|
self.parallel_limiter[device_id], callback
|
|||
|
|
)
|
|||
|
|
tasks.append(task)
|
|||
|
|
|
|||
|
|
# 等待所有任务完成并收集结果
|
|||
|
|
for i, task in enumerate(tasks):
|
|||
|
|
result = await task
|
|||
|
|
pages_result.append({
|
|||
|
|
"page_number": page_from + i + 1,
|
|||
|
|
"boxes": result
|
|||
|
|
})
|
|||
|
|
else:
|
|||
|
|
# 串行处理（单GPU或CPU）
|
|||
|
|
for i, img in enumerate(page_images):
|
|||
|
|
page_num = page_from + i + 1
|
|||
|
|
result = await trio.to_thread.run_sync(
|
|||
|
|
lambda img=img, pn=page_num: self.__ocr_page(pn, img, zoomin, 0)
|
|||
|
|
)
|
|||
|
|
pages_result.append({
|
|||
|
|
"page_number": page_num,
|
|||
|
|
"boxes": result
|
|||
|
|
})
|
|||
|
|
if callback:
|
|||
|
|
callback(0.1 + (i + 1) * 0.9 / len(page_images), f"Processing page {page_num}...")
|
|||
|
|
|
|||
|
|
return pages_result
|
|||
|
|
|
|||
|
|
# 运行异步处理
|
|||
|
|
if callback:
|
|||
|
|
callback(0.2, "Starting OCR processing...")
|
|||
|
|
|
|||
|
|
start = timer()
|
|||
|
|
pages_result = trio.run(process_all_pages)
|
|||
|
|
logging.info(f"OCR processing completed in {timer() - start:.2f}s")
|
|||
|
|
|
|||
|
|
if callback:
|
|||
|
|
callback(1.0, "OCR processing completed")
|
|||
|
|
|
|||
|
|
return {
|
|||
|
|
"pages": pages_result
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
|
|||
|
|
# 向后兼容的别名
|
|||
|
|
PdfParser = SimplePdfParser
|
|||
|
|
|