tcsenpai/pensieve
1
0
Fork 0
mirror of https://github.com/tcsenpai/pensieve.git synced 2025-06-10 13:07:15 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(ocr): update import flow

Browse Source
This commit is contained in:
arkohut 2024-09-05 18:15:27 +08:00
parent c1226117d6
commit e062636c08
1 changed files with 18 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

19
memos/plugins/ocr/server.py
View File

@ -1,4 +1,3 @@
from rapidocr_onnxruntime import RapidOCR
from PIL import Image from PIL import Image
import numpy as np import numpy as np
import logging import logging
@ -32,8 +31,22 @@ def init_thread_pool(max_workers):
def init_ocr(use_gpu): def init_ocr(use_gpu):
global ocr global ocr
config_path = "ppocr-gpu.yaml" if use_gpu else "ppocr.yaml" if use_gpu:
ocr = RapidOCR(config_path=config_path) try:
from rapidocr_paddle import RapidOCR as RapidOCRPaddle
ocr = RapidOCRPaddle(det_use_cuda=True, cls_use_cuda=True, rec_use_cuda=True)
logger.info("Initialized OCR with RapidOCR Paddle (GPU)")
except ImportError:
logger.error("Failed to import rapidocr_paddle. Make sure it's installed for GPU usage.")
raise
else:
try:
from rapidocr_onnxruntime import RapidOCR
ocr = RapidOCR(config_path="ppocr.yaml")
logger.info("Initialized OCR with RapidOCR ONNX Runtime (CPU)")
except ImportError:
logger.error("Failed to import rapidocr_onnxruntime. Make sure it's installed for CPU usage.")
raise
def convert_ocr_results(results): def convert_ocr_results(results):