feat: Introduce automatic language detection in TesseractOcrCliModel (#800)

* feat: Introduce automatic language detection in tesseract_ocr_cli model. Extend unit tests.

Signed-off-by: Nikos Livathinos <nli@zurich.ibm.com>

* docs: Add example how to use "auto" language with tesseract OCR engines

Signed-off-by: Nikos Livathinos <nli@zurich.ibm.com>

* fix: Refactor the TesseractOcrModel and TesseractOcrCliModel to validate if the auto-detected
language is installed in the system and if not fall back to a default option without language.

Signed-off-by: Nikos Livathinos <nli@zurich.ibm.com>

---------

Signed-off-by: Nikos Livathinos <nli@zurich.ibm.com>

docling/models/tesseract_ocr_cli_model.py

@@ -4,7 +4,7 @@ import logging
 import os
 import tempfile
 from subprocess import DEVNULL, PIPE, Popen
-from typing import Iterable, Optional, Tuple
+from typing import Iterable, List, Optional, Tuple
 
 import pandas as pd
 from docling_core.types.doc import BoundingBox, CoordOrigin
@@ -14,6 +14,7 @@ from docling.datamodel.document import ConversionResult
 from docling.datamodel.pipeline_options import TesseractCliOcrOptions
 from docling.datamodel.settings import settings
 from docling.models.base_ocr_model import BaseOcrModel
+from docling.utils.ocr_utils import map_tesseract_script
 from docling.utils.profiling import TimeRecorder
 
 _log = logging.getLogger(__name__)
@@ -28,10 +29,13 @@ class TesseractOcrCliModel(BaseOcrModel):
 
         self._name: Optional[str] = None
         self._version: Optional[str] = None
+        self._tesseract_languages: Optional[List[str]] = None
+        self._script_prefix: Optional[str] = None
 
         if self.enabled:
             try:
                 self._get_name_and_version()
+                self._set_languages_and_prefix()
 
             except Exception as exc:
                 raise RuntimeError(
@@ -73,12 +77,20 @@ class TesseractOcrCliModel(BaseOcrModel):
         return name, version
 
     def _run_tesseract(self, ifilename: str):
-
+        r"""
+        Run tesseract CLI
+        """
         cmd = [self.options.tesseract_cmd]
 
-        if self.options.lang is not None and len(self.options.lang) > 0:
+        if "auto" in self.options.lang:
+            lang = self._detect_language(ifilename)
+            if lang is not None:
+                cmd.append("-l")
+                cmd.append(lang)
+        elif self.options.lang is not None and len(self.options.lang) > 0:
             cmd.append("-l")
             cmd.append("+".join(self.options.lang))
+
         if self.options.path is not None:
             cmd.append("--tessdata-dir")
             cmd.append(self.options.path)
@@ -106,6 +118,63 @@ class TesseractOcrCliModel(BaseOcrModel):
 
         return df_filtered
 
+    def _detect_language(self, ifilename: str):
+        r"""
+        Run tesseract in PSM 0 mode to detect the language
+        """
+        assert self._tesseract_languages is not None
+
+        cmd = [self.options.tesseract_cmd]
+        cmd.extend(["--psm", "0", "-l", "osd", ifilename, "stdout"])
+        _log.info("command: {}".format(" ".join(cmd)))
+        proc = Popen(cmd, stdout=PIPE, stderr=DEVNULL)
+        output, _ = proc.communicate()
+        decoded_data = output.decode("utf-8")
+        df = pd.read_csv(
+            io.StringIO(decoded_data), sep=":", header=None, names=["key", "value"]
+        )
+        scripts = df.loc[df["key"] == "Script"].value.tolist()
+        if len(scripts) == 0:
+            _log.warning("Tesseract cannot detect the script of the page")
+            return None
+
+        script = map_tesseract_script(scripts[0].strip())
+        lang = f"{self._script_prefix}{script}"
+
+        # Check if the detected language has been installed
+        if lang not in self._tesseract_languages:
+            msg = f"Tesseract detected the script '{script}' and language '{lang}'."
+            msg += " However this language is not installed in your system and will be ignored."
+            _log.warning(msg)
+            return None
+
+        _log.debug(
+            f"Using tesseract model for the detected script '{script}' and language '{lang}'"
+        )
+        return lang
+
+    def _set_languages_and_prefix(self):
+        r"""
+        Read and set the languages installed in tesseract and decide the script prefix
+        """
+        # Get all languages
+        cmd = [self.options.tesseract_cmd]
+        cmd.append("--list-langs")
+        _log.info("command: {}".format(" ".join(cmd)))
+        proc = Popen(cmd, stdout=PIPE, stderr=DEVNULL)
+        output, _ = proc.communicate()
+        decoded_data = output.decode("utf-8")
+        df = pd.read_csv(io.StringIO(decoded_data), header=None)
+        self._tesseract_languages = df[0].tolist()[1:]
+
+        # Decide the script prefix
+        if any([l.startswith("script/") for l in self._tesseract_languages]):
+            script_prefix = "script/"
+        else:
+            script_prefix = ""
+
+        self._script_prefix = script_prefix
+
     def __call__(
         self, conv_res: ConversionResult, page_batch: Iterable[Page]
     ) -> Iterable[Page]:
@@ -120,7 +189,6 @@ class TesseractOcrCliModel(BaseOcrModel):
                 yield page
             else:
                 with TimeRecorder(conv_res, "ocr"):
-
                     ocr_rects = self.get_ocr_rects(page)
 
                     all_ocr_cells = []

docling/models/tesseract_ocr_model.py

@@ -8,6 +8,7 @@ from docling.datamodel.document import ConversionResult
 from docling.datamodel.pipeline_options import TesseractOcrOptions
 from docling.datamodel.settings import settings
 from docling.models.base_ocr_model import BaseOcrModel
+from docling.utils.ocr_utils import map_tesseract_script
 from docling.utils.profiling import TimeRecorder
 
 _log = logging.getLogger(__name__)
@@ -20,6 +21,7 @@ class TesseractOcrModel(BaseOcrModel):
 
         self.scale = 3  # multiplier for 72 dpi == 216 dpi.
         self.reader = None
+        self.osd_reader = None
 
         if self.enabled:
             install_errmsg = (
@@ -47,8 +49,8 @@ class TesseractOcrModel(BaseOcrModel):
             except:
                 raise ImportError(install_errmsg)
 
-            _, tesserocr_languages = tesserocr.get_languages()
+            _, self._tesserocr_languages = tesserocr.get_languages()
-            if not tesserocr_languages:
+            if not self._tesserocr_languages:
                 raise ImportError(missing_langs_errmsg)
 
             # Initialize the tesseractAPI
@@ -57,7 +59,7 @@ class TesseractOcrModel(BaseOcrModel):
 
             self.script_readers: dict[str, tesserocr.PyTessBaseAPI] = {}
 
-            if any([l.startswith("script/") for l in tesserocr_languages]):
+            if any([l.startswith("script/") for l in self._tesserocr_languages]):
                 self.script_prefix = "script/"
             else:
                 self.script_prefix = ""
@@ -72,14 +74,14 @@ class TesseractOcrModel(BaseOcrModel):
                 tesserocr_kwargs["path"] = self.options.path
 
             if lang == "auto":
-                self.reader = tesserocr.PyTessBaseAPI(
+                self.reader = tesserocr.PyTessBaseAPI(**tesserocr_kwargs)
+                self.osd_reader = tesserocr.PyTessBaseAPI(
                     **{"lang": "osd", "psm": tesserocr.PSM.OSD_ONLY} | tesserocr_kwargs
                 )
             else:
                 self.reader = tesserocr.PyTessBaseAPI(
                     **{"lang": lang} | tesserocr_kwargs,
                 )
-
             self.reader_RIL = tesserocr.RIL
 
     def __del__(self):
@@ -96,8 +98,6 @@ class TesseractOcrModel(BaseOcrModel):
             yield from page_batch
             return
 
-        import tesserocr
-
         for page in page_batch:
             assert page._backend is not None
             if not page._backend.is_valid():
@@ -105,6 +105,7 @@ class TesseractOcrModel(BaseOcrModel):
             else:
                 with TimeRecorder(conv_res, "ocr"):
                     assert self.reader is not None
+                    assert self._tesserocr_languages is not None
 
                     ocr_rects = self.get_ocr_rects(page)
 
@@ -117,43 +118,42 @@ class TesseractOcrModel(BaseOcrModel):
                             scale=self.scale, cropbox=ocr_rect
                         )
 
-                        # Retrieve text snippets with their bounding boxes
+                        local_reader = self.reader
-                        self.reader.SetImage(high_res_image)
+                        if "auto" in self.options.lang:
+                            assert self.osd_reader is not None
 
-                        if self.options.lang == ["auto"]:
+                            self.osd_reader.SetImage(high_res_image)
-                            osd = self.reader.DetectOrientationScript()
+                            osd = self.osd_reader.DetectOrientationScript()
 
                             # No text, probably
                             if osd is None:
                                 continue
 
                             script = osd["script_name"]
+                            script = map_tesseract_script(script)
+                            lang = f"{self.script_prefix}{script}"
 
-                            if script == "Katakana" or script == "Hiragana":
+                            # Check if the detected languge is present in the system
-                                script = "Japanese"
+                            if lang not in self._tesserocr_languages:
-                            elif script == "Han":
+                                msg = f"Tesseract detected the script '{script}' and language '{lang}'."
-                                script = "HanS"
+                                msg += " However this language is not installed in your system and will be ignored."
-                            elif script == "Korean":
+                                _log.warning(msg)
-                                script = "Hangul"
+                            else:
+                                if script not in self.script_readers:
+                                    import tesserocr
 
-                            _log.debug(
+                                    self.script_readers[script] = (
-                                f'Using model for the detected script "{script}"'
+                                        tesserocr.PyTessBaseAPI(
-                            )
+                                            path=self.reader.GetDatapath(),
-
+                                            lang=lang,
-                            if script not in self.script_readers:
+                                            psm=tesserocr.PSM.AUTO,
-                                self.script_readers[script] = tesserocr.PyTessBaseAPI(
+                                            init=True,
-                                    path=self.reader.GetDatapath(),
+                                            oem=tesserocr.OEM.DEFAULT,
-                                    lang=f"{self.script_prefix}{script}",
+                                        )
-                                    psm=tesserocr.PSM.AUTO,
+                                    )
-                                    init=True,
+                                local_reader = self.script_readers[script]
-                                    oem=tesserocr.OEM.DEFAULT,
-                                )
-
-                            local_reader = self.script_readers[script]
-                            local_reader.SetImage(high_res_image)
-                        else:
-                            local_reader = self.reader
 
+                        local_reader.SetImage(high_res_image)
                         boxes = local_reader.GetComponentImages(
                             self.reader_RIL.TEXTLINE, True
                         )

docling/utils/ocr_utils.py

@@ -0,0 +1,9 @@
+def map_tesseract_script(script: str) -> str:
+    r""" """
+    if script == "Katakana" or script == "Hiragana":
+        script = "Japanese"
+    elif script == "Han":
+        script = "HanS"
+    elif script == "Korean":
+        script = "Hangul"
+    return script

docs/examples/tesseract_lang_detection.py

@@ -0,0 +1,37 @@
+from pathlib import Path
+
+from docling.datamodel.base_models import InputFormat
+from docling.datamodel.pipeline_options import (
+    PdfPipelineOptions,
+    TesseractCliOcrOptions,
+    TesseractOcrOptions,
+)
+from docling.document_converter import DocumentConverter, PdfFormatOption
+
+
+def main():
+    input_doc = Path("./tests/data/2206.01062.pdf")
+
+    # Set lang=["auto"] with a tesseract OCR engine: TesseractOcrOptions, TesseractCliOcrOptions
+    # ocr_options = TesseractOcrOptions(lang=["auto"])
+    ocr_options = TesseractCliOcrOptions(lang=["auto"])
+
+    pipeline_options = PdfPipelineOptions(
+        do_ocr=True, force_full_page_ocr=True, ocr_options=ocr_options
+    )
+
+    converter = DocumentConverter(
+        format_options={
+            InputFormat.PDF: PdfFormatOption(
+                pipeline_options=pipeline_options,
+            )
+        }
+    )
+
+    doc = converter.convert(input_doc).document
+    md = doc.export_to_markdown()
+    print(md)
+
+
+if __name__ == "__main__":
+    main()

mkdocs.yml

@@ -75,6 +75,7 @@ nav:
       - "Table export": examples/export_tables.py
       - "Multimodal export": examples/export_multimodal.py
       - "Force full page OCR": examples/full_page_ocr.py
+      - "Automatic OCR language detection with tesseract": examples/tesseract_lang_detection.py
       - "Accelerator options": examples/run_with_accelerator.py
       - "Simple translation": examples/translate.py
     - ✂️ Chunking:

tests/test_e2e_ocr_conversion.py

@@ -62,6 +62,7 @@ def test_e2e_conversions():
         TesseractOcrOptions(force_full_page_ocr=True),
         TesseractOcrOptions(force_full_page_ocr=True, lang=["auto"]),
         TesseractCliOcrOptions(force_full_page_ocr=True),
+        TesseractCliOcrOptions(force_full_page_ocr=True, lang=["auto"]),
         RapidOcrOptions(force_full_page_ocr=True),
     ]