mirrors/docling
1
0
Fork 0
mirror of https://github.com/DS4SD/docling.git synced 2025-08-02 07:22:14 +00:00
Code Issues Packages Projects Releases Wiki Activity

Modify "auto" language in TesseractOcr to initialize the script readers lazily

Browse Source 
Signed-off-by: Pavel Denisov <pavel.denisov@iais.fraunhofer.de>
This commit is contained in:
Pavel Denisov 2025-01-23 10:03:53 +01:00
parent 570a1a560a
commit eed50d46e3
1 changed files with 19 additions and 22 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

41
docling/models/tesseract_ocr_model.py
View File

@ -20,7 +20,6 @@ class TesseractOcrModel(BaseOcrModel):
self.scale = 3 # multiplier for 72 dpi == 216 dpi. self.scale = 3 # multiplier for 72 dpi == 216 dpi.
self.reader = None self.reader = None
self.script_readers = None
if self.enabled: if self.enabled:
install_errmsg = ( install_errmsg = (
@ -56,6 +55,8 @@ class TesseractOcrModel(BaseOcrModel):
_log.debug("Initializing TesserOCR: %s", tesseract_version) _log.debug("Initializing TesserOCR: %s", tesseract_version)
lang = "+".join(self.options.lang) lang = "+".join(self.options.lang)
self.script_readers: dict[str, tesserocr.PyTessBaseAPI] = {}
tesserocr_kwargs = { tesserocr_kwargs = {
"psm": tesserocr.PSM.AUTO, "psm": tesserocr.PSM.AUTO,
"init": True, "init": True,
@ -69,14 +70,6 @@ class TesseractOcrModel(BaseOcrModel):
self.reader = tesserocr.PyTessBaseAPI( self.reader = tesserocr.PyTessBaseAPI(
**{"lang": "osd", "psm": tesserocr.PSM.OSD_ONLY} | tesserocr_kwargs **{"lang": "osd", "psm": tesserocr.PSM.OSD_ONLY} | tesserocr_kwargs
) )
self.script_readers = {}
scripts = [l for l in tesserocr_languages if l.startswith("script")]
for script in scripts:
self.script_readers[script] = tesserocr.PyTessBaseAPI(
**{"lang": script} | tesserocr_kwargs,
)
else: else:
self.reader = tesserocr.PyTessBaseAPI( self.reader = tesserocr.PyTessBaseAPI(
**{"lang": lang} | tesserocr_kwargs, **{"lang": lang} | tesserocr_kwargs,
@ -92,18 +85,18 @@ class TesseractOcrModel(BaseOcrModel):
def __call__( def __call__(
self, conv_res: ConversionResult, page_batch: Iterable[Page] self, conv_res: ConversionResult, page_batch: Iterable[Page]
) -> Iterable[Page]: ) -> Iterable[Page]:
if not self.enabled: if not self.enabled:
yield from page_batch yield from page_batch
return return
import tesserocr
for page in page_batch: for page in page_batch:
assert page._backend is not None assert page._backend is not None
if not page._backend.is_valid(): if not page._backend.is_valid():
yield page yield page
else: else:
with TimeRecorder(conv_res, "ocr"): with TimeRecorder(conv_res, "ocr"):
assert self.reader is not None assert self.reader is not None
ocr_rects = self.get_ocr_rects(page) ocr_rects = self.get_ocr_rects(page)
@ -120,7 +113,7 @@ class TesseractOcrModel(BaseOcrModel):
# Retrieve text snippets with their bounding boxes # Retrieve text snippets with their bounding boxes
self.reader.SetImage(high_res_image) self.reader.SetImage(high_res_image)
if self.script_readers is not None: if self.options.lang == ["auto"]:
osd = self.reader.DetectOrientationScript() osd = self.reader.DetectOrientationScript()
# No text, probably # No text, probably
@ -136,17 +129,21 @@ class TesseractOcrModel(BaseOcrModel):
elif script == "Korean": elif script == "Korean":
script = "Hangul" script = "Hangul"
if f"script/{script}" in self.script_readers: _log.debug(
_log.debug( f'Using model for the detected script "{script}"'
f'Using model for the detected script "{script}"' )
if script not in self.script_readers:
self.script_readers[script] = tesserocr.PyTessBaseAPI(
path=self.reader.GetDatapath(),
lang=f"script/{script}",
psm=tesserocr.PSM.AUTO,
init=True,
oem=tesserocr.OEM.DEFAULT,
) )
local_reader = self.script_readers[f"script/{script}"]
local_reader.SetImage(high_res_image) local_reader = self.script_readers[script]
else: local_reader.SetImage(high_res_image)
_log.warning(
f'No model for the detected script "{script}"'
)
continue
else: else:
local_reader = self.reader local_reader = self.reader