Add option for CPU only in EasyOCR

This commit is contained in:
AuroraWright
2024-02-12 01:56:07 -05:00
parent 6fb4aa4b12
commit eec4484217
2 changed files with 6 additions and 3 deletions

View File

@@ -488,12 +488,13 @@ class EasyOCR:
key = 'e' key = 'e'
available = False available = False
def __init__(self): def __init__(self, config={'gpu': True}):
if 'easyocr' not in sys.modules: if 'easyocr' not in sys.modules:
logger.warning('easyocr not available, EasyOCR will not work!') logger.warning('easyocr not available, EasyOCR will not work!')
else: else:
logger.info('Loading EasyOCR model') logger.info('Loading EasyOCR model')
self.model = easyocr.Reader(['ja','en']) logging.getLogger('easyocr.easyocr').disabled = True
self.model = easyocr.Reader(['ja','en'], gpu=config['gpu'])
self.available = True self.available = True
logger.info('EasyOCR ready') logger.info('EasyOCR ready')

View File

@@ -33,3 +33,5 @@
[mangaocr] [mangaocr]
pretrained_model_name_or_path = kha-white/manga-ocr-base pretrained_model_name_or_path = kha-white/manga-ocr-base
force_cpu = False force_cpu = False
[easyocr]
gpu = True