254 lines
9.1 KiB
Python
254 lines
9.1 KiB
Python
"""
|
|
RapidOCR 处理器 - 纯OCR文字识别
|
|
|
|
使用 RapidOCR (PP-OCRv4) 进行文字识别
|
|
"""
|
|
|
|
import os
|
|
import tempfile
|
|
import time
|
|
from pathlib import Path
|
|
|
|
import fitz
|
|
import numpy as np
|
|
from PIL import Image
|
|
from rapidocr_onnxruntime import RapidOCR
|
|
|
|
from src.plugins.document_processor_base import BaseDocumentProcessor, OCRException
|
|
from src.utils import logger
|
|
|
|
|
|
class RapidOCRProcessor(BaseDocumentProcessor):
|
|
"""RapidOCR 处理器 - 使用 ONNX 模型进行文字识别"""
|
|
|
|
def __init__(self, det_box_thresh: float = 0.3):
|
|
self.ocr = None
|
|
self.det_box_thresh = det_box_thresh
|
|
self.model_dir_root = (
|
|
os.getenv("MODEL_DIR") if not os.getenv("RUNNING_IN_DOCKER") else os.getenv("MODEL_DIR_IN_DOCKER")
|
|
)
|
|
|
|
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_paths(self) -> tuple[str, str]:
|
|
"""获取模型文件路径"""
|
|
model_dir = os.path.join(self.model_dir_root, "SWHL/RapidOCR")
|
|
det_model_path = os.path.join(model_dir, "PP-OCRv4/ch_PP-OCRv4_det_infer.onnx")
|
|
rec_model_path = os.path.join(model_dir, "PP-OCRv4/ch_PP-OCRv4_rec_infer.onnx")
|
|
return det_model_path, rec_model_path
|
|
|
|
def check_health(self) -> dict:
|
|
"""检查 RapidOCR 模型是否可用"""
|
|
try:
|
|
det_model_path, rec_model_path = self._get_model_paths()
|
|
model_dir = os.path.dirname(os.path.dirname(det_model_path))
|
|
|
|
if not os.path.exists(model_dir):
|
|
return {
|
|
"status": "unavailable",
|
|
"message": f"模型目录不存在: {model_dir}",
|
|
"details": {"model_dir": model_dir},
|
|
}
|
|
|
|
if not os.path.exists(det_model_path) or not os.path.exists(rec_model_path):
|
|
return {
|
|
"status": "unavailable",
|
|
"message": "模型文件缺失",
|
|
"details": {"det_model": det_model_path, "rec_model": rec_model_path},
|
|
}
|
|
|
|
# 尝试加载模型
|
|
try:
|
|
test_ocr = RapidOCR(
|
|
det_box_thresh=self.det_box_thresh, det_model_path=det_model_path, rec_model_path=rec_model_path
|
|
)
|
|
del test_ocr # 释放资源
|
|
return {
|
|
"status": "healthy",
|
|
"message": "RapidOCR模型可用",
|
|
"details": {"model_path": self._get_model_paths()},
|
|
}
|
|
except Exception as e:
|
|
return {"status": "error", "message": f"模型加载失败: {str(e)}", "details": {"error": str(e)}}
|
|
|
|
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 模型...")
|
|
|
|
# 先检查健康状态
|
|
health = self.check_health()
|
|
if health["status"] != "healthy":
|
|
raise OCRException(health["message"], self.get_service_name(), health["status"])
|
|
|
|
try:
|
|
det_model_path, rec_model_path = self._get_model_paths()
|
|
self.ocr = RapidOCR(
|
|
det_box_thresh=self.det_box_thresh, det_model_path=det_model_path, rec_model_path=rec_model_path
|
|
)
|
|
logger.info(f"RapidOCR 模型加载成功 (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:
|
|
text = "\n".join([line[1] for line in result])
|
|
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)
|