wehub-resource-sync
226 lines
7.8 KiB
Python
226 lines
7.8 KiB
Python
"""
|
|
RapidOCR 解析器 - 纯OCR文字识别
|
|
|
|
使用 RapidOCR (PP-OCRv5) 进行文字识别
|
|
"""
|
|
|
|
import os
|
|
import tempfile
|
|
import time
|
|
from pathlib import Path
|
|
|
|
import fitz
|
|
import numpy as np
|
|
from PIL import Image
|
|
from rapidocr import EngineType, LangDet, LangRec, ModelType, OCRVersion, RapidOCR
|
|
|
|
from yuxi.knowledge.parser.base import BaseDocumentProcessor, OCRException
|
|
from yuxi.utils import logger
|
|
|
|
|
|
class RapidOCRParser(BaseDocumentProcessor):
|
|
"""RapidOCR 解析器 - 使用 ONNX 模型进行文字识别"""
|
|
|
|
def __init__(self, det_box_thresh: float = 0.3):
|
|
self.ocr = None
|
|
self.det_box_thresh = det_box_thresh
|
|
|
|
def get_service_name(self) -> str:
|
|
return "rapid_ocr"
|
|
|
|
def get_supported_extensions(self) -> list[str]:
|
|
return [".pdf", ".jpg", ".jpeg", ".png", ".bmp", ".tiff", ".tif"]
|
|
|
|
def _get_model_params(self) -> dict[str, object]:
|
|
return {
|
|
"Det.engine_type": EngineType.ONNXRUNTIME,
|
|
"Det.lang_type": LangDet.CH,
|
|
"Det.model_type": ModelType.MOBILE,
|
|
"Det.ocr_version": OCRVersion.PPOCRV5,
|
|
"Det.box_thresh": self.det_box_thresh,
|
|
"Cls.engine_type": EngineType.ONNXRUNTIME,
|
|
"Rec.engine_type": EngineType.ONNXRUNTIME,
|
|
"Rec.lang_type": LangRec.CH,
|
|
"Rec.model_type": ModelType.MOBILE,
|
|
"Rec.ocr_version": OCRVersion.PPOCRV5,
|
|
}
|
|
|
|
def check_health(self) -> dict:
|
|
"""检查 RapidOCR 模型是否可用"""
|
|
try:
|
|
test_ocr = RapidOCR(params=self._get_model_params())
|
|
del test_ocr
|
|
return {
|
|
"status": "healthy",
|
|
"message": "RapidOCR PP-OCRv5 模型可用",
|
|
"details": {"ocr_version": "PP-OCRv5", "engine": "onnxruntime"},
|
|
}
|
|
except Exception as e:
|
|
return {"status": "error", "message": f"模型加载失败: {str(e)}", "details": {"error": str(e)}}
|
|
|
|
def _load_model(self):
|
|
"""延迟加载 OCR 模型"""
|
|
if self.ocr is not None:
|
|
return
|
|
|
|
logger.info("加载 RapidOCR 模型...")
|
|
|
|
try:
|
|
self.ocr = RapidOCR(params=self._get_model_params())
|
|
logger.info(f"RapidOCR PP-OCRv5 模型加载成功 (det_box_thresh={self.det_box_thresh})")
|
|
except Exception as e:
|
|
raise OCRException(f"RapidOCR模型加载失败: {str(e)}", self.get_service_name(), "load_failed")
|
|
|
|
def process_image(self, image, params: dict | None = None) -> str:
|
|
"""
|
|
处理单张图像并提取文本
|
|
|
|
Args:
|
|
image: 图像数据,支持:
|
|
- str: 图像文件路径
|
|
- PIL.Image: PIL图像对象
|
|
- numpy.ndarray: numpy图像数组
|
|
params: 处理参数 (当前未使用)
|
|
|
|
Returns:
|
|
str: 提取的文本内容
|
|
"""
|
|
self._load_model()
|
|
|
|
try:
|
|
# 处理不同类型的输入
|
|
if isinstance(image, str):
|
|
image_path = image
|
|
cleanup_needed = False
|
|
else:
|
|
# 创建临时文件
|
|
image_path = self._create_temp_image_file(image)
|
|
cleanup_needed = True
|
|
|
|
try:
|
|
# 执行 OCR
|
|
start_time = time.time()
|
|
result = self.ocr(image_path)
|
|
processing_time = time.time() - start_time
|
|
|
|
# 提取文本
|
|
if result.txts:
|
|
text = "\n".join(result.txts)
|
|
logger.info(
|
|
f"RapidOCR 成功: {os.path.basename(image_path) if isinstance(image, str) else 'temp_image'}"
|
|
f" ({processing_time:.2f}s)"
|
|
)
|
|
return text
|
|
else:
|
|
logger.warning(f"RapidOCR 未识别到文本: {image_path}")
|
|
return ""
|
|
|
|
finally:
|
|
# 清理临时文件
|
|
if cleanup_needed and os.path.exists(image_path):
|
|
try:
|
|
os.remove(image_path)
|
|
except Exception as e:
|
|
logger.warning(f"临时文件清理失败: {image_path} - {e}")
|
|
|
|
except Exception as e:
|
|
error_msg = f"图像OCR处理失败: {str(e)}"
|
|
logger.error(error_msg)
|
|
raise OCRException(error_msg, self.get_service_name(), "processing_failed")
|
|
|
|
def _create_temp_image_file(self, image) -> str:
|
|
"""将图像数据保存为临时文件"""
|
|
try:
|
|
# 使用系统临时目录
|
|
with tempfile.NamedTemporaryFile(mode="wb", suffix=".png", delete=False) as tmp_file:
|
|
temp_path = tmp_file.name
|
|
|
|
if isinstance(image, Image.Image):
|
|
image.save(temp_path)
|
|
elif isinstance(image, np.ndarray):
|
|
Image.fromarray(image).save(temp_path)
|
|
else:
|
|
raise ValueError("不支持的图像类型,必须是 PIL.Image 或 numpy.ndarray")
|
|
|
|
return temp_path
|
|
|
|
except Exception as e:
|
|
raise OCRException(f"临时图像文件创建失败: {str(e)}", self.get_service_name(), "temp_file_error")
|
|
|
|
def process_pdf(self, pdf_path: str, params: dict | None = None) -> str:
|
|
"""
|
|
处理 PDF 文件并提取文本 (流式处理,避免内存占用)
|
|
|
|
Args:
|
|
pdf_path: PDF 文件路径
|
|
params: 处理参数
|
|
- zoom_x: 横向缩放 (默认 2)
|
|
- zoom_y: 纵向缩放 (默认 2)
|
|
|
|
Returns:
|
|
str: 提取的文本
|
|
"""
|
|
if not os.path.exists(pdf_path):
|
|
raise OCRException(f"PDF 文件不存在: {pdf_path}", self.get_service_name(), "file_not_found")
|
|
|
|
params = params or {}
|
|
zoom_x = params.get("zoom_x", 2)
|
|
zoom_y = params.get("zoom_y", 2)
|
|
|
|
try:
|
|
all_text = []
|
|
pdf_doc = fitz.open(pdf_path)
|
|
total_pages = pdf_doc.page_count
|
|
|
|
logger.info(f"开始处理 PDF: {os.path.basename(pdf_path)} ({total_pages} 页)")
|
|
|
|
# 流式处理每一页,避免一次性加载所有图片到内存
|
|
for page_num in range(total_pages):
|
|
page = pdf_doc[page_num]
|
|
|
|
# 转换为图像
|
|
mat = fitz.Matrix(zoom_x, zoom_y)
|
|
pix = page.get_pixmap(matrix=mat, alpha=False)
|
|
img_pil = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
|
|
|
|
# 立即处理,不保存到列表
|
|
text = self.process_image(img_pil)
|
|
all_text.append(text)
|
|
|
|
if (page_num + 1) % 10 == 0:
|
|
logger.info(f"已处理 {page_num + 1}/{total_pages} 页")
|
|
|
|
pdf_doc.close()
|
|
|
|
result_text = "\n\n".join(all_text)
|
|
logger.info(f"PDF OCR 完成: {os.path.basename(pdf_path)} - {len(result_text)} 字符")
|
|
return result_text
|
|
|
|
except OCRException:
|
|
raise
|
|
except Exception as e:
|
|
error_msg = f"PDF OCR 处理失败: {str(e)}"
|
|
logger.error(error_msg)
|
|
raise OCRException(error_msg, self.get_service_name(), "pdf_processing_failed")
|
|
|
|
def process_file(self, file_path: str, params: dict | None = None) -> str:
|
|
"""
|
|
处理文件 (PDF 或图像)
|
|
|
|
Args:
|
|
file_path: 文件路径
|
|
params: 处理参数
|
|
|
|
Returns:
|
|
str: 提取的文本
|
|
"""
|
|
file_ext = Path(file_path).suffix.lower()
|
|
|
|
if not self.supports_file_type(file_ext):
|
|
raise OCRException(f"不支持的文件类型: {file_ext}", self.get_service_name(), "unsupported_file_type")
|
|
|
|
if file_ext == ".pdf":
|
|
return self.process_pdf(file_path, params)
|
|
else:
|
|
return self.process_image(file_path, params)
|