WindyMadman/AzurLaneAutoScript
1
0
Fork 0
mirror of https://github.com/LmeSzinc/AzurLaneAutoScript.git synced 2025-01-09 06:07:40 +08:00
Code Issues Packages Projects Releases Wiki Activity

Chore: Poorly search mxnet-cu* instead of using pkg_resources

Browse Source
This commit is contained in:
LmeSzinc 2024-04-15 02:54:44 +08:00
parent 502558dbbd
commit 16ba4aaad6
2 changed files with 8 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
module/base/base.py
View File

@ -83,8 +83,8 @@ class ModuleBase:
time.sleep(0.01)
logger.info('early_ocr_import start')
from cnocr import CnOcr
_ = CnOcr
from module.ocr.al_ocr import AlOcr
_ = AlOcr
logger.info('early_ocr_import finish')
logger.info('early_ocr_import call')

10
module/ocr/al_ocr.py
View File

@ -17,10 +17,12 @@ from cnocr.hyperparams.cn_hyperparams import CnHyperparams as Hyperparams
def get_mxnet_context():
import re
import pkg_resources
for pkg in pkg_resources.working_set:
if re.match(r'^mxnet-cu\d+$', pkg.key):
logger.info(f'MXNet gpu package: {pkg.key}=={pkg.version} found, using it')
import cnocr
site_packages = os.path.abspath(os.path.join(cnocr.__file__, '../../'))
for file in os.listdir(site_packages):
# mxnet_cu101-1.6.0.dist-info
if re.match(r'^mxnet[-_]cu\d+', file):
logger.info(f'MXNet gpu package: {file} found, using it')
return 'gpu'
return 'cpu'