Yuxi-Know
Yuxi-Know copied to clipboard
Feat: OCR 在 PyTorch 上的并行加速
#️⃣ 您的功能请求是否与某个问题相关?请描述。
目前的 OCR 基于 rapidocr_onnxruntime 推理引擎,只能在 CPU 上使用,这样固然方便项目的部署与使用,但在大量数据处理的时候会出现问题。
#️⃣ 描述您期望的解决方案
对于 OCR 部分考虑使用目前成熟的端到端 PDF2TXT 的解决方案,比如 MinerU
#️⃣ 附加背景信息
from rapidocr_onnxruntime import RapidOCR
class OCRPlugin:
"""OCR 插件"""
def __init__(self, **kwargs):
self.ocr = None
self.det_box_thresh = kwargs.get('det_box_thresh', 0.3)
def load_model(self):
"""加载 OCR 模型"""
logger.info(f"加载 OCR 模型,仅在第一次调用时加载")
model_dir = os.path.join(os.getenv("MODEL_DIR", ""), "SWHL/RapidOCR")
det_model_dir = os.path.join(model_dir, "PP-OCRv4/ch_PP-OCRv4_det_infer.onnx")
rec_model_dir = os.path.join(model_dir, "PP-OCRv4/ch_PP-OCRv4_rec_infer.onnx")
assert os.path.exists(model_dir), (
f"模型文件不存在,请下载 SWHL/RapidOCR 到 {model_dir},"
"并确认是否在 docker-compose.dev.yml 中添加 MODEL_DIR 环境变量"
)
self.ocr = RapidOCR(det_box_thresh=0.3, det_model_path=det_model_dir, rec_model_path=rec_model_dir)
logger.info(f"OCR Plugin for det_box_thresh = {self.det_box_thresh} loaded.")