qurator-spk/eynollah
1
0
Fork 0
mirror of https://github.com/qurator-spk/eynollah.git synced 2026-07-26 05:29:16 +02:00
Code Issues Projects Releases Wiki Activity

cnn-rnn-ocr: increase default batch size and VRAM limit

Browse source
This commit is contained in:
Robert Sachunsky 2026-07-18 00:59:57 +02:00
parent be8b161607
commit 79a9bb0128
3 changed files with 5 additions and 5 deletions
Download patch file Download diff file Expand all files Collapse all files

2
src/eynollah/eynollah_ocr.py
View file

@ -83,7 +83,7 @@ class Eynollah_ocr(Eynollah):
self.logger = logger if logger else logging.getLogger('eynollah.ocr') self.logger = logger if logger else logging.getLogger('eynollah.ocr')
self.min_conf_value_of_textline_text = min_conf_value_of_textline_text self.min_conf_value_of_textline_text = min_conf_value_of_textline_text
self.b_s = batch_size or 2 if tr_ocr else 8 self.b_s = batch_size or (2 if tr_ocr else 64)
self.model_zoo = model_zoo self.model_zoo = model_zoo
self.setup_models(device=device) self.setup_models(device=device)

2
src/eynollah/model_zoo/model_zoo.py
View file

@ -24,7 +24,7 @@ MODEL_VRAM_LIMITS = {
"region_fl_np": 1756, "region_fl_np": 1756,
"table": 1818, "table": 1818,
"reading_order": 632, "reading_order": 632,
"ocr": 850, "ocr": 2400, # 850 for bs 8
} }
class EynollahModelZoo: class EynollahModelZoo:

6
src/eynollah/predictor.py
View file

@ -1,5 +1,5 @@
from contextlib import ExitStack from contextlib import ExitStack
from typing import List, Dict, Tuple, Union from typing import List, Dict, Sequence, Tuple, Union
import logging import logging
import logging.handlers import logging.handlers
import multiprocessing as mp import multiprocessing as mp
@ -41,10 +41,10 @@ class Predictor(mp.context.SpawnProcess):
def input_shape(self): def input_shape(self):
return self({}) return self({})
def predict(self, data: ArrayT, verbose=0) -> ArrayT: def predict(self, data: Union[Sequence[ArrayT], ArrayT], verbose=0) -> Union[Sequence[ArrayT], ArrayT]:
return self(data) return self(data)
def __call__(self, data: Union[ArrayT, Dict]) -> Union[ArrayT, Tuple]: def __call__(self, data: Union[Sequence[ArrayT], ArrayT, Dict]) -> Union[ArrayT, Tuple]:
# unusable as per python/cpython#79967 # unusable as per python/cpython#79967
#with self.jobid.get_lock(): #with self.jobid.get_lock():
# would work, but not public: # would work, but not public: