tcsenpai/pensieve
1
0
Fork 0
mirror of https://github.com/tcsenpai/pensieve.git synced 2025-06-06 19:25:24 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(ocr): use openvino if possible

Browse Source
This commit is contained in:
arkohut 2024-10-22 00:00:50 +08:00
parent 780a571bed
commit 2174e9ed2d
2 changed files with 11 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
memos/plugins/ocr/main.py
View File

@ -7,12 +7,12 @@ import json
import base64 import base64
from PIL import Image from PIL import Image
import numpy as np import numpy as np
from rapidocr_onnxruntime import RapidOCR
from concurrent.futures import ThreadPoolExecutor from concurrent.futures import ThreadPoolExecutor
from functools import partial from functools import partial
import yaml import yaml
import io import io
import platform import platform
import cpuinfo
MAX_THUMBNAIL_SIZE = (1920, 1920) MAX_THUMBNAIL_SIZE = (1920, 1920)
@ -224,7 +224,12 @@ def init_plugin(config):
with open(temp_config_path, 'w') as f: with open(temp_config_path, 'w') as f:
yaml.safe_dump(ocr_config, f) yaml.safe_dump(ocr_config, f)
ocr = RapidOCR(config_path=temp_config_path) if platform.system() == 'Windows' and 'Intel' in cpuinfo.get_cpu_info()['brand_raw']:
from rapidocr_openvino import RapidOCR
ocr = RapidOCR(config_path=temp_config_path)
else:
from rapidocr_onnxruntime import RapidOCR
ocr = RapidOCR(config_path=temp_config_path)
thread_pool = ThreadPoolExecutor(max_workers=concurrency) thread_pool = ThreadPoolExecutor(max_workers=concurrency)
logger.info("OCR plugin initialized") logger.info("OCR plugin initialized")
@ -232,4 +237,5 @@ def init_plugin(config):
logger.info(f"Token: {token}") logger.info(f"Token: {token}")
logger.info(f"Concurrency: {concurrency}") logger.info(f"Concurrency: {concurrency}")
logger.info(f"Use local: {use_local}") logger.info(f"Use local: {use_local}")
if use_local:
logger.info(f"OCR library: {'rapidocr_openvino' if platform.system() == 'Windows' and 'Intel' in cpuinfo.get_cpu_info()['brand_raw'] else 'rapidocr_onnxruntime'}")

2
pyproject.toml
View File

@ -30,6 +30,8 @@ dependencies = [
"piexif", "piexif",
"imagehash", "imagehash",
"rapidocr_onnxruntime", "rapidocr_onnxruntime",
"rapidocr_openvino; sys_platform == 'win32'",
"py-cpuinfo",
"screeninfo", "screeninfo",
"pywin32; sys_platform == 'win32'", "pywin32; sys_platform == 'win32'",
"psutil; sys_platform == 'win32'", "psutil; sys_platform == 'win32'",