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

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

View file

@ -83,7 +83,7 @@ class Eynollah_ocr(Eynollah):
self.logger = logger if logger else logging.getLogger('eynollah.ocr')
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.setup_models(device=device)

View file

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

View file

@ -1,5 +1,5 @@
from contextlib import ExitStack
from typing import List, Dict, Tuple, Union
from typing import List, Dict, Sequence, Tuple, Union
import logging
import logging.handlers
import multiprocessing as mp
@ -41,10 +41,10 @@ class Predictor(mp.context.SpawnProcess):
def input_shape(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)
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
#with self.jobid.get_lock():
# would work, but not public: