diff --git a/docling/datamodel/pipeline_options.py b/docling/datamodel/pipeline_options.py index 910b11cca4..b6e61883d9 100644 --- a/docling/datamodel/pipeline_options.py +++ b/docling/datamodel/pipeline_options.py @@ -662,6 +662,19 @@ class TesseractCliOcrOptions(OcrOptions): ) ), ] = None + orientation_confidence_threshold: Annotated[ + float, + Field( + description=( + "Minimum orientation confidence from Tesseract's orientation and script detection (OSD) for an " + "OCR region to be rotated before recognition. OSD reports it as the margin between its best and " + "second-best orientation scores, so values near 0 are close calls. Below the threshold the region " + "is read as rendered; `0` rotates on every detection." + ), + examples=[0.0, 1.0, 14.0], + ge=0.0, + ), + ] = 1.0 model_config = ConfigDict( extra="forbid", ) @@ -706,6 +719,19 @@ class TesseractOcrOptions(OcrOptions): ) ), ] = None + orientation_confidence_threshold: Annotated[ + float, + Field( + description=( + "Minimum orientation confidence from Tesseract's orientation and script detection (OSD) for an " + "OCR region to be rotated before recognition. OSD reports it as the margin between its best and " + "second-best orientation scores, so values near 0 are close calls. Below the threshold the region " + "is read as rendered; `0` rotates on every detection." + ), + examples=[0.0, 1.0, 14.0], + ge=0.0, + ), + ] = 1.0 model_config = ConfigDict( extra="forbid", ) diff --git a/docling/models/stages/ocr/tesseract_ocr_cli_model.py b/docling/models/stages/ocr/tesseract_ocr_cli_model.py index 2ea2c78137..3df79aaf62 100644 --- a/docling/models/stages/ocr/tesseract_ocr_cli_model.py +++ b/docling/models/stages/ocr/tesseract_ocr_cli_model.py @@ -18,7 +18,12 @@ from docling_core.types.doc.page import TextCell from docling.datamodel.accelerator_options import AcceleratorOptions -from docling.datamodel.base_models import Page +from docling.datamodel.base_models import ( + DoclingComponentType, + ErrorItem, + FailureCategory, + Page, +) from docling.datamodel.document import ConversionResult from docling.datamodel.pipeline_options import ( OcrOptions, @@ -67,6 +72,7 @@ def __init__( # multiplier for 72 dpi; the default 3.0 == 216 dpi. self.scale = self.options.scale + self._dpi = round(72 * self.scale) self._name: Optional[str] = None self._version: Optional[str] = None @@ -224,12 +230,13 @@ def _run_tesseract(self, ifilename: str, osd: Optional[pd.DataFrame]): if self.options.psm is not None: cmd.extend(["--psm", str(int(self.options.psm))]) + cmd.extend(["--dpi", str(self._dpi)]) cmd.append(self._sanitize_filename(ifilename)) cmd.extend(["stdout", "tsv"]) _log.info("command: {}".format(" ".join(cmd))) output = subprocess.run( - cmd, stdout=PIPE, stderr=DEVNULL, stdin=DEVNULL, check=True, shell=False + cmd, capture_output=True, stdin=DEVNULL, check=True, shell=False ) # _log.info(output) @@ -270,6 +277,8 @@ def _perform_osd(self, ifilename: str) -> pd.DataFrame: "0", "-l", "osd", + "--dpi", + str(self._dpi), self._sanitize_filename(ifilename), "stdout", ] @@ -341,7 +350,7 @@ def __call__( continue high_res_image = page._backend.get_page_image( scale=self.scale, cropbox=ocr_rect - ) + ).convert("L") try: with tempfile.NamedTemporaryFile( suffix=".png", mode="w+b", delete=False @@ -352,7 +361,10 @@ def __call__( df_osd: Optional[pd.DataFrame] = None try: df_osd = self._perform_osd(fname) - doc_orientation = _parse_orientation(df_osd) + doc_orientation = _parse_orientation( + df_osd, + self.options.orientation_confidence_threshold, + ) except subprocess.CalledProcessError as exc: _log.error( "OSD failed (doc %s, page: %s, " @@ -384,6 +396,18 @@ def __call__( image_file, exc.stderr, ) + conv_res.errors.append( + ErrorItem( + component_type=DoclingComponentType.MODEL, + module_name=type(self).__name__, + error_message=( + f"Tesseract failed on OCR rectangle {ocr_rect_i}: " + f"{_stderr_text(exc)}" + ), + category=FailureCategory.INFERENCE_FAILURE, + page_no=page.page_no, + ) + ) continue finally: if os.path.exists(fname): @@ -437,9 +461,16 @@ def get_options_type(cls) -> Type[OcrOptions]: return TesseractCliOcrOptions -def _parse_orientation(df_osd: pd.DataFrame) -> int: +def _parse_orientation(df_osd: pd.DataFrame, confidence_threshold: float) -> int: # For strictly optimal performance with invariant dataframe format: - mask = df_osd["key"].to_numpy() == "Orientation in degrees" - orientation_val = df_osd["value"].to_numpy()[mask][0] + keys = df_osd["key"].to_numpy() + values = df_osd["value"].to_numpy() + if float(values[keys == "Orientation confidence"][0]) < confidence_threshold: + return 0 + orientation_val = values[keys == "Orientation in degrees"][0] orientation = parse_tesseract_orientation(orientation_val.strip()) return orientation + + +def _stderr_text(exc: subprocess.CalledProcessError) -> str: + return " ".join((exc.stderr or b"").decode("utf-8", "replace").split()) diff --git a/docling/models/stages/ocr/tesseract_ocr_model.py b/docling/models/stages/ocr/tesseract_ocr_model.py index 19e0ea2d3c..d07ef27ee2 100644 --- a/docling/models/stages/ocr/tesseract_ocr_model.py +++ b/docling/models/stages/ocr/tesseract_ocr_model.py @@ -59,6 +59,7 @@ def __init__( self._auto_script: bool = not self.languages # multiplier for 72 dpi; the default 3.0 == 216 dpi. self.scale = self.options.scale + self._dpi = round(72 * self.scale) self._reader = None self._script_readers: dict[str, tesserocr.PyTessBaseAPI] = {} self._tesseract_vocabulary: list[str] = [] @@ -188,10 +189,11 @@ def __call__( continue high_res_image = page._backend.get_page_image( scale=self.scale, cropbox=ocr_rect - ) + ).convert("L") local_reader = self._reader self._osd_reader.SetImage(high_res_image) + self._osd_reader.SetSourceResolution(self._dpi) doc_orientation = 0 osd = self._osd_reader.DetectOrientationScript() @@ -209,7 +211,10 @@ def __call__( # to OCR in the hope OCR will succeed while OSD failed if self._auto_script: continue - else: + elif ( + osd["orient_conf"] + >= self.options.orientation_confidence_threshold + ): doc_orientation = parse_tesseract_orientation( osd["orient_deg"] ) @@ -244,6 +249,7 @@ def __call__( local_reader = self._script_readers[lang] local_reader.SetImage(high_res_image) + local_reader.SetSourceResolution(self._dpi) boxes = local_reader.GetComponentImages( self._reader_RIL.TEXTLINE, True ) diff --git a/docs/concepts/OCR.md b/docs/concepts/OCR.md index 8ddcdfad60..41cc31d6bc 100644 --- a/docs/concepts/OCR.md +++ b/docs/concepts/OCR.md @@ -301,6 +301,25 @@ tag where one exists, and the bare file name for everything else. An empty `lang` list runs Tesseract's per-page orientation and script detection. That requires the `osd` traineddata; without it, `lang=[]` raises with an install hint. +### Tesseract image input and orientation + +Both Tesseract engines receive a grayscale render of each OCR region, together with its resolution: +72 DPI times `scale`, so 216 DPI by default. Tesseract binarizes a colour image channel by channel, +which can blank out text on a highlighter mark or a coloured background; a grayscale image avoids +that. The cost is lower contrast for light text on a strongly coloured box and for light, saturated +text such as orange or yellow on white, which Tesseract can then miss. + +Before reading a region, Docling runs Tesseract's orientation and script detection (OSD) on it and +turns it upright. OSD reports its orientation confidence as the margin between its best and +second-best orientation, so a value near 0 is a close call. A detection below +`orientation_confidence_threshold` (default `1.0`) is ignored and the region is read as rendered. +Set it to `0` to rotate on every detection, or raise it when a wrong rotation costs you more than a +missed one; OCRmyPDF's page-rotation default, for example, is 14. + +When Tesseract fails on a region in the `tesseract` engine, the failure is reported in +`ConversionResult.errors` with its page number and Tesseract's message, and the document is marked +`PARTIAL_SUCCESS`. + ## OcrMac The engine's own vocabulary, in its own codes, is listed in diff --git a/tests/data/ocr/sources/highlighted_text.png b/tests/data/ocr/sources/highlighted_text.png new file mode 100644 index 0000000000..330dea587e Binary files /dev/null and b/tests/data/ocr/sources/highlighted_text.png differ diff --git a/tests/test_e2e_tesseract_highlighted_text.py b/tests/test_e2e_tesseract_highlighted_text.py new file mode 100644 index 0000000000..3ad55820ec --- /dev/null +++ b/tests/test_e2e_tesseract_highlighted_text.py @@ -0,0 +1,88 @@ +# SPDX-FileCopyrightText: The Docling Contributors +# SPDX-License-Identifier: MIT + +import importlib.util +import re +import shutil +from pathlib import Path + +import pytest + +from docling.datamodel.base_models import InputFormat +from docling.datamodel.pipeline_options import ( + OcrMode, + PdfPipelineOptions, + TesseractCliOcrOptions, + TesseractOcrOptions, +) +from docling.document_converter import DocumentConverter, ImageFormatOption + +pytestmark = pytest.mark.ml_ocr + +_FIXTURE = Path("./tests/data/ocr/sources/highlighted_text.png") + +_LINES = [ + "The parties agree that the landlord will repair the roof before winter.", + "Rent is due on the first day of each month at the office of the manager.", + "Late payments accrue a fee of five percent after a grace period of ten days.", + "The tenant must give sixty days written notice before moving out of the unit.", + "Pets are allowed only with the written consent of the landlord in advance.", + "The security deposit will be returned within thirty days after the lease ends.", + "Either party may terminate this agreement for a material breach of its terms.", + "This agreement is governed by the laws of the state where the unit is located.", +] +_HIGHLIGHTED_ROWS = (2, 3, 4) + +_SKIP_NO_TESSERACT = pytest.mark.skipif( + shutil.which("tesseract") is None, reason="tesseract not installed" +) +_SKIP_NO_TESSEROCR = pytest.mark.skipif( + importlib.util.find_spec("tesserocr") is None, reason="tesserocr not installed" +) + + +def _generate_fixture(): + from PIL import Image, ImageDraw, ImageFont + + dpi = 200 + image = Image.new("RGB", (round(8.5 * dpi), round(11 * dpi)), "white") + draw = ImageDraw.Draw(image) + font = ImageFont.load_default(size=28) + y = dpi + for row, line in enumerate(_LINES): + if row in _HIGHLIGHTED_ROWS: + left, top, right, bottom = draw.textbbox((dpi, y), line, font=font) + draw.rectangle((left - 4, top - 6, right + 4, bottom + 6), fill="yellow") + draw.text((dpi, y), line, fill="black", font=font) + y += 50 + image.save(_FIXTURE, dpi=(dpi, dpi)) + + +def _words(text): + return re.findall(r"[a-z]+", text.lower()) + + +@pytest.mark.parametrize( + "ocr_options_cls", + [ + pytest.param(TesseractCliOcrOptions, marks=_SKIP_NO_TESSERACT, id="tesseract"), + pytest.param(TesseractOcrOptions, marks=_SKIP_NO_TESSEROCR, id="tesserocr"), + ], +) +def test_full_page_ocr_reads_highlighted_lines(ocr_options_cls): + options = PdfPipelineOptions() + options.do_ocr = True + options.ocr_options = ocr_options_cls(lang=["eng"], mode=OcrMode.FULL_PAGE) + converter = DocumentConverter( + format_options={InputFormat.IMAGE: ImageFormatOption(pipeline_options=options)} + ) + + found = set(_words(converter.convert(_FIXTURE).document.export_to_markdown())) + + highlighted = [word for row in _HIGHLIGHTED_ROWS for word in _words(_LINES[row])] + missing = [word for word in highlighted if word not in found] + assert len(missing) <= len(highlighted) // 10, missing + + +if __name__ == "__main__": + _generate_fixture() diff --git a/tests/test_tesseract_ocr_cli_model.py b/tests/test_tesseract_ocr_cli_model.py index cff7377603..921b5d5115 100644 --- a/tests/test_tesseract_ocr_cli_model.py +++ b/tests/test_tesseract_ocr_cli_model.py @@ -1,15 +1,30 @@ # SPDX-FileCopyrightText: The Docling Contributors # SPDX-License-Identifier: MIT +import subprocess from unittest.mock import patch import pytest +from PIL import Image -from docling.datamodel.pipeline_options import TesseractCliOcrOptions +from docling.backend.image_backend import ImageDocumentBackend +from docling.datamodel.accelerator_options import AcceleratorOptions +from docling.datamodel.base_models import FailureCategory, InputFormat, Page +from docling.datamodel.document import ConversionResult, InputDocument +from docling.datamodel.pipeline_options import OcrMode, TesseractCliOcrOptions from docling.models.stages.ocr.tesseract_ocr_cli_model import TesseractOcrCliModel _MODULE = "docling.models.stages.ocr.tesseract_ocr_cli_model" +_OSD_STDOUT = ( + "Page number: 0\n" + "Orientation in degrees: 90\n" + "Rotate: 270\n" + "Orientation confidence: {confidence}\n" + "Script: Latin\n" + "Script confidence: 1.84\n" +) + _TSV_HEADER = ( "level\tpage_num\tblock_num\tpar_num\tline_num\tword_num" "\tleft\ttop\twidth\theight\tconf\ttext" @@ -44,6 +59,7 @@ def test_recognized_words_are_kept_verbatim(words: list[str]) -> None: model._safe_tessdata_path = None model._auto_script = False model._native_codes = ["eng"] + model._dpi = 216 with patch( f"{_MODULE}.subprocess.run", @@ -52,3 +68,92 @@ def test_recognized_words_are_kept_verbatim(words: list[str]) -> None: df = model._run_tesseract("page.png", osd=None) assert list(df["text"]) == words + + +class _FakeTesseract: + def __init__(self, osd_confidence=8.0, ocr_stderr=None): + self.osd_confidence = osd_confidence + self.ocr_stderr = ocr_stderr + self.commands: list[list[str]] = [] + self.ocr_images: list[tuple[str, tuple[int, int]]] = [] + + def __call__(self, cmd, **kwargs): + self.commands.append(cmd) + if "osd" in cmd: + stdout = _OSD_STDOUT.format(confidence=self.osd_confidence) + return _FakeCompletedProcess(stdout.encode("utf-8")) + with Image.open(next(arg for arg in cmd if arg.endswith(".png"))) as image: + self.ocr_images.append((image.mode, image.size)) + if self.ocr_stderr is not None: + raise subprocess.CalledProcessError(1, cmd, b"", self.ocr_stderr) + return _FakeCompletedProcess(_tsv(["Hello"]).encode("utf-8")) + + +def _convert_page(monkeypatch, tmp_path, tesseract, **options): + monkeypatch.setattr(f"{_MODULE}.subprocess.run", tesseract) + monkeypatch.setattr( + TesseractOcrCliModel, + "_get_name_and_version", + lambda self: ("tesseract", "5.5.0"), + ) + monkeypatch.setattr( + TesseractOcrCliModel, + "_set_languages", + lambda self: setattr(self, "_tesseract_vocabulary", ["eng", "osd"]), + ) + model = TesseractOcrCliModel( + enabled=True, + artifacts_path=None, + options=TesseractCliOcrOptions(mode=OcrMode.FULL_PAGE, **options), + accelerator_options=AcceleratorOptions(), + ) + image_path = tmp_path / "page.png" + Image.new("RGB", (200, 100), "white").save(image_path) + in_doc = InputDocument( + path_or_stream=image_path, + format=InputFormat.IMAGE, + backend=ImageDocumentBackend, + ) + conv_res = ConversionResult(input=in_doc) + page_backend = in_doc._backend.load_page(0) + page = Page(page_no=1, size=page_backend.get_size()) + page._backend = page_backend + assert list(model(conv_res, [page])) == [page] + return conv_res + + +@pytest.mark.parametrize( + ("confidence", "ocr_size"), + [(0.4, (600, 300)), (8.0, (300, 600))], + ids=["close-call", "confident"], +) +def test_region_is_rotated_only_on_a_confident_orientation( + monkeypatch, tmp_path, confidence, ocr_size +): + tesseract = _FakeTesseract(osd_confidence=confidence) + + _convert_page(monkeypatch, tmp_path, tesseract, lang=["eng"]) + + assert tesseract.ocr_images == [("L", ocr_size)] + + +def test_tesseract_is_told_the_render_resolution(monkeypatch, tmp_path): + tesseract = _FakeTesseract() + + _convert_page(monkeypatch, tmp_path, tesseract, lang=["eng"], scale=2.0) + + assert [cmd[cmd.index("--dpi") + 1] for cmd in tesseract.commands] == [ + "144", + "144", + ] + + +def test_failed_region_is_recorded_on_the_page(monkeypatch, tmp_path): + tesseract = _FakeTesseract(ocr_stderr=b"Error in pixReadStream: simulated") + + conv_res = _convert_page(monkeypatch, tmp_path, tesseract, lang=["eng"]) + + assert [ + (error.page_no, error.category, error.module_name) for error in conv_res.errors + ] == [(1, FailureCategory.INFERENCE_FAILURE, "TesseractOcrCliModel")] + assert "pixReadStream: simulated" in conv_res.errors[0].error_message diff --git a/tests/test_tesseract_ocr_model.py b/tests/test_tesseract_ocr_model.py index 7a9c3135d5..616d40388c 100644 --- a/tests/test_tesseract_ocr_model.py +++ b/tests/test_tesseract_ocr_model.py @@ -26,6 +26,9 @@ class _OcrReader(_Reader): def SetImage(self, image): pass + def SetSourceResolution(self, ppi): + pass + def DetectOrientationScript(self): return {"orient_deg": 0, "orient_conf": 10.0, "script_name": "Latin"} @@ -106,6 +109,16 @@ def test_cell_confidence_is_scaled_to_unit_range(monkeypatch, tmp_path): accelerator_options=AcceleratorOptions(), ) + conv_res, page = _ocr_blank_page(model, tmp_path, page_no=0) + + assert page.parsed_page is not None + cells = page.parsed_page.textline_cells + assert [cell.confidence for cell in cells] == [pytest.approx(0.42)] + assert conv_res.confidence.pages[0].ocr_score == pytest.approx(0.42) + del model + + +def _ocr_blank_page(model, tmp_path, page_no=1): image_path = tmp_path / "page.png" Image.new("RGB", (200, 100), "white").save(image_path) in_doc = InputDocument( @@ -115,13 +128,59 @@ def test_cell_confidence_is_scaled_to_unit_range(monkeypatch, tmp_path): ) conv_res = ConversionResult(input=in_doc) page_backend = in_doc._backend.load_page(0) - page = Page(page_no=0, size=page_backend.get_size()) + page = Page(page_no=page_no, size=page_backend.get_size()) page._backend = page_backend - list(model(conv_res, [page])) + assert list(model(conv_res, [page])) == [page] + return conv_res, page - assert page.parsed_page is not None - cells = page.parsed_page.textline_cells - assert [cell.confidence for cell in cells] == [pytest.approx(0.42)] - assert conv_res.confidence.pages[0].ocr_score == pytest.approx(0.42) + +def _recording_reader(osd): + class _RecordingReader(_OcrReader): + images: list[tuple[str, tuple[int, int]]] = [] + resolutions: list[int] = [] + + def SetImage(self, image): + _RecordingReader.images.append((image.mode, image.size)) + + def SetSourceResolution(self, ppi): + _RecordingReader.resolutions.append(ppi) + + def DetectOrientationScript(self): + return osd + + return _RecordingReader + + +def _model_with_reader(monkeypatch, reader, **options): + monkeypatch.setitem( + sys.modules, + "tesserocr", + _fake_tesserocr(lambda **kwargs: ("/tessdata", ["eng", "osd"]), reader), + ) + return TesseractOcrModel( + enabled=True, + artifacts_path=None, + options=TesseractOcrOptions(mode=OcrMode.FULL_PAGE, **options), + accelerator_options=AcceleratorOptions(), + ) + + +@pytest.mark.parametrize( + ("orient_conf", "ocr_size"), + [(0.4, (600, 300)), (8.0, (300, 600))], + ids=["close-call", "confident"], +) +def test_region_is_rotated_only_on_a_confident_orientation( + monkeypatch, tmp_path, orient_conf, ocr_size +): + reader = _recording_reader( + {"orient_deg": 90, "orient_conf": orient_conf, "script_name": "Latin"} + ) + model = _model_with_reader(monkeypatch, reader, lang=["eng"]) + + _ocr_blank_page(model, tmp_path) + + assert reader.images == [("L", (600, 300)), ("L", ocr_size)] + assert reader.resolutions == [216, 216] del model