sui_feng/AzurLaneAutoScript
1
0
Fork 0
mirror of https://gitee.com/sui-feng-cb/AzurLaneAutoScript1 synced 2026-03-12 07:28:21 +08:00
Code 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
Download Patch File Download Diff File Expand all files Collapse all files

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'