diff --git a/packages/markitdown-ocr/src/markitdown_ocr/_docx_converter_with_ocr.py b/packages/markitdown-ocr/src/markitdown_ocr/_docx_converter_with_ocr.py index f2463de11..88e5a6735 100644 --- a/packages/markitdown-ocr/src/markitdown_ocr/_docx_converter_with_ocr.py +++ b/packages/markitdown-ocr/src/markitdown_ocr/_docx_converter_with_ocr.py @@ -3,8 +3,6 @@ Extracts images from Word documents and performs OCR while maintaining context. """ -import io -import re import sys from typing import Any, BinaryIO, Optional @@ -21,12 +19,11 @@ _dependency_exc_info = None try: import mammoth - from docx import Document except ImportError: _dependency_exc_info = sys.exc_info() -# Placeholder injected into HTML so that mammoth never sees the OCR markers. -# Must be a single token with no special markdown characters. +# Placeholder carried through Mammoth and HTML-to-Markdown conversion. +# It must be a single token with no special markdown characters. _PLACEHOLDER = "MARKITDOWNOCRBLOCK{}" @@ -83,35 +80,56 @@ def convert( ) if ocr_service: - # 1. Extract and OCR images — returns raw text per image - file_stream.seek(0) - image_ocr_map = self._extract_and_ocr_images(file_stream, ocr_service) + ocr_text_by_placeholder: dict[str, str] = {} + image_occurrence = 0 + + def convert_image(image: Any) -> dict[str, str]: + nonlocal image_occurrence - # 2. Convert DOCX → HTML via mammoth + placeholder = _PLACEHOLDER.format(image_occurrence) + image_occurrence += 1 + + try: + with image.open() as image_stream: + ocr_result = ocr_service.extract_text(image_stream) + text = ocr_result.text.strip() + if text: + ocr_text_by_placeholder[placeholder] = text + except Exception: + pass + + # Keep every occurrence in the converted output, including empty + # or failed OCR, so later images cannot shift to an earlier slot. + return {"src": placeholder, "alt": placeholder} + + # Convert DOCX → HTML while Mammoth visits every image occurrence. file_stream.seek(0) pre_process_stream = pre_process_docx(file_stream) html_result = mammoth.convert_to_html( - pre_process_stream, style_map=kwargs.get("style_map") + pre_process_stream, + style_map=kwargs.get("style_map"), + convert_image=mammoth.images.img_element(convert_image), ).value - # 3. Replace tags with plain placeholder tokens so that - # mammoth's HTML→markdown step never escapes our OCR markers. - html_with_placeholders, ocr_texts = self._inject_placeholders( - html_result, image_ocr_map - ) - - # 4. Convert HTML → markdown + # Convert HTML → markdown before adding OCR text so formatting in OCR + # output is not escaped by the markdown converter. md_result = self._html_converter.convert_string( - html_with_placeholders, **kwargs + html_result, **kwargs ) md = md_result.markdown - # 5. Swap placeholders for the actual OCR blocks (post-conversion - # so * and _ are never escaped by the markdown converter). - for i, raw_text in enumerate(ocr_texts): - placeholder = _PLACEHOLDER.format(i) + for placeholder, raw_text in ocr_text_by_placeholder.items(): ocr_block = f"*[Image OCR]\n{raw_text}\n[End OCR]*" - md = md.replace(placeholder, ocr_block) + md = md.replace( + f"![{placeholder}]({placeholder})", + ocr_block, + ) + + # Remove failed or empty OCR occurrences only after successful + # placeholders have been substituted at their original positions. + for occurrence in range(image_occurrence): + placeholder = _PLACEHOLDER.format(occurrence) + md = md.replace(f"![{placeholder}]({placeholder})", "") return DocumentConverterResult(markdown=md) else: @@ -122,68 +140,3 @@ def convert( mammoth.convert_to_html(pre_process_stream, style_map=style_map).value, **kwargs, ) - - def _extract_and_ocr_images( - self, file_stream: BinaryIO, ocr_service: LLMVisionOCRService - ) -> dict[str, str]: - """ - Extract images from DOCX and OCR them. - - Returns: - Dict mapping image relationship IDs to raw OCR text (no markers). - """ - ocr_map = {} - - try: - file_stream.seek(0) - doc = Document(file_stream) - - for rel in doc.part.rels.values(): - if "image" in rel.target_ref.lower(): - try: - image_bytes = rel.target_part.blob - image_stream = io.BytesIO(image_bytes) - ocr_result = ocr_service.extract_text(image_stream) - - if ocr_result.text.strip(): - # Store raw text only — markers added later - ocr_map[rel.rId] = ocr_result.text.strip() - - except Exception: - continue - - except Exception: - pass - - return ocr_map - - def _inject_placeholders( - self, html: str, ocr_map: dict[str, str] - ) -> tuple[str, list[str]]: - """ - Replace tags with numbered placeholder tokens. - - Returns: - (html_with_placeholders, ordered list of raw OCR texts) - """ - if not ocr_map: - return html, [] - - ocr_texts = list(ocr_map.values()) - used: list[int] = [] - - def replace_img(match: re.Match) -> str: # type: ignore[type-arg] - for i in range(len(ocr_texts)): - if i not in used: - used.append(i) - return f"

{_PLACEHOLDER.format(i)}

" - return "" # remove image if all OCR texts already used - - result = re.sub(r"]*>", replace_img, html) - - # Any OCR texts that had no matching tag go at the end - for i in range(len(ocr_texts)): - if i not in used: - result += f"

{_PLACEHOLDER.format(i)}

" - - return result, ocr_texts diff --git a/packages/markitdown-ocr/src/markitdown_ocr/_pdf_converter_with_ocr.py b/packages/markitdown-ocr/src/markitdown_ocr/_pdf_converter_with_ocr.py index c1dc0f613..86c88a7e6 100644 --- a/packages/markitdown-ocr/src/markitdown_ocr/_pdf_converter_with_ocr.py +++ b/packages/markitdown-ocr/src/markitdown_ocr/_pdf_converter_with_ocr.py @@ -191,43 +191,63 @@ def convert( # If OCR is enabled, interleave text and images by position if ocr_service: images_on_page = self._extract_page_images(pdf_bytes, page_num) + has_native_text = False + has_image_ocr_text = False if images_on_page: # Extract text lines with Y positions - chars = page.chars + try: + chars = page.chars + except Exception: + chars = [] + if chars: - # Group chars into lines based on Y position - lines_with_y = [] - current_line = [] - current_y = None - - for char in sorted( - chars, key=lambda c: (c["top"], c["x0"]) - ): - y = char["top"] - if current_y is None: - current_y = y - elif abs(y - current_y) > 2: # New line threshold - if current_line: - text = "".join( - [c["text"] for c in current_line] - ) - lines_with_y.append( - {"y": current_y, "text": text.strip()} - ) - current_line = [] - current_y = y - current_line.append(char) - - # Add last line - if current_line: - text = "".join([c["text"] for c in current_line]) - lines_with_y.append( - {"y": current_y, "text": text.strip()} - ) + try: + # Group chars into lines based on Y position + lines_with_y = [] + current_line = [] + current_y = None + + for char in sorted( + chars, key=lambda c: (c["top"], c["x0"]) + ): + y = char["top"] + if current_y is None: + current_y = y + elif abs(y - current_y) > 2: # New line threshold + if current_line: + text = "".join( + [c["text"] for c in current_line] + ) + lines_with_y.append( + { + "y": current_y, + "text": text.strip(), + } + ) + current_line = [] + current_y = y + current_line.append(char) + + # Add last line + if current_line: + text = "".join( + [c["text"] for c in current_line] + ) + lines_with_y.append( + {"y": current_y, "text": text.strip()} + ) + except Exception: + lines_with_y = [] else: # Fallback: use simple text extraction - text_content = page.extract_text() or "" + lines_with_y = [] + + if not lines_with_y: + try: + text_content = page.extract_text() or "" + except Exception: + text_content = "" lines_with_y = [ {"y": i * 10, "text": line} for i, line in enumerate(text_content.split("\n")) @@ -236,19 +256,22 @@ def convert( # OCR all images image_data = [] for img_info in images_on_page: - ocr_result = ocr_service.extract_text( - img_info["stream"] - ) - if ocr_result.text.strip(): - image_data.append( - { - "y_pos": img_info["y_pos"], - "name": img_info["name"], - "ocr_text": ocr_result.text, - "backend": ocr_result.backend_used, - "type": "image", - } + try: + ocr_result = ocr_service.extract_text( + img_info["stream"] ) + if ocr_result.text.strip(): + image_data.append( + { + "y_pos": img_info["y_pos"], + "name": img_info["name"], + "ocr_text": ocr_result.text, + "backend": ocr_result.backend_used, + "type": "image", + } + ) + except Exception: + continue # Add text items content_items = [ @@ -261,6 +284,10 @@ def convert( if item["text"] ] content_items.extend(image_data) + has_native_text = any( + item["type"] == "text" for item in content_items + ) + has_image_ocr_text = bool(image_data) # Sort all items by Y position (top to bottom) content_items.sort(key=lambda x: x["y_pos"]) @@ -277,12 +304,27 @@ def convert( markdown_content.append(img_marker) else: # No images detected - just extract regular text - text_content = page.extract_text() or "" + try: + text_content = page.extract_text() or "" + except Exception: + text_content = "" if text_content.strip(): markdown_content.append(text_content.strip()) + has_native_text = True + + # A page heading alone is not extracted content. OCR this + # individual page only when neither native nor inline OCR + # text was available. + if not has_native_text and not has_image_ocr_text: + markdown_content.append( + self._ocr_rendered_page(page, page_num, ocr_service) + ) else: # No OCR, just extract text - text_content = page.extract_text() or "" + try: + text_content = page.extract_text() or "" + except Exception: + text_content = "" if text_content.strip(): markdown_content.append(text_content.strip()) @@ -357,31 +399,10 @@ def _ocr_full_pages( pdf_bytes.seek(0) with pdfplumber.open(pdf_bytes) as pdf: for page_num, page in enumerate(pdf.pages, 1): - try: - markdown_parts.append(f"\n## Page {page_num}\n") - - # Render page to image - page_img = page.to_image(resolution=300) - img_stream = io.BytesIO() - page_img.original.save(img_stream, format="PNG") - img_stream.seek(0) - - # Run OCR - ocr_result = ocr_service.extract_text(img_stream) - - if ocr_result.text.strip(): - text = ocr_result.text.strip() - markdown_parts.append(f"*[Image OCR]\n{text}\n[End OCR]*") - else: - markdown_parts.append( - "*[No text could be extracted from this page]*" - ) - - except Exception as e: - markdown_parts.append( - f"*[Error processing page {page_num}: {str(e)}]*" - ) - continue + markdown_parts.append(f"\n## Page {page_num}\n") + markdown_parts.append( + self._ocr_rendered_page(page, page_num, ocr_service) + ) except Exception: # pdfplumber failed (e.g. malformed EOF) — try PyMuPDF for rendering @@ -420,3 +441,22 @@ def _ocr_full_pages( return "*[Error: Could not process scanned PDF]*" return "\n\n".join(markdown_parts).strip() + + def _ocr_rendered_page( + self, page: Any, page_num: int, ocr_service: LLMVisionOCRService + ) -> str: + """Render and OCR one pdfplumber page without affecting other pages.""" + try: + page_img = page.to_image(resolution=300) + img_stream = io.BytesIO() + page_img.original.save(img_stream, format="PNG") + img_stream.seek(0) + + ocr_result = ocr_service.extract_text(img_stream) + if ocr_result.text.strip(): + text = ocr_result.text.strip() + return f"*[Image OCR]\n{text}\n[End OCR]*" + + return "*[No text could be extracted from this page]*" + except Exception as e: + return f"*[Error processing page {page_num}: {str(e)}]*" diff --git a/packages/markitdown-ocr/tests/test_docx_converter.py b/packages/markitdown-ocr/tests/test_docx_converter.py index 0fb666504..9f0959db4 100644 --- a/packages/markitdown-ocr/tests/test_docx_converter.py +++ b/packages/markitdown-ocr/tests/test_docx_converter.py @@ -10,9 +10,12 @@ [End OCR]* """ +import io import sys from pathlib import Path +from types import SimpleNamespace from typing import Any +from unittest.mock import patch import pytest @@ -163,48 +166,66 @@ def test_docx_complex_layout(svc: MockOCRService) -> None: assert _convert("docx_complex_layout.docx", svc) == expected -# --------------------------------------------------------------------------- -# _inject_placeholders — internal unit tests (no file I/O) -# --------------------------------------------------------------------------- +class SequencedOCRService: + def extract_text(self, image_stream: Any, **kwargs: Any) -> OCRResult: + image_name = image_stream.read().decode() + if image_name == "empty": + return OCRResult(text="", backend_used="mock") + if image_name == "error": + raise RuntimeError("OCR unavailable") + return OCRResult(text=f"OCR {image_name}", backend_used="mock") -def test_inject_placeholders_single_image() -> None: - converter = DocxConverterWithOCR() - html = "

Before

After

" - result_html, texts = converter._inject_placeholders(html, {"rId1": "TEXT"}) - assert " None: - converter = DocxConverterWithOCR() - html = "

Mid

" - result_html, texts = converter._inject_placeholders( - html, {"rId1": "FIRST", "rId2": "SECOND"} - ) - assert "MARKITDOWNOCRBLOCK0" in result_html - assert "MARKITDOWNOCRBLOCK1" in result_html - assert result_html.index("MARKITDOWNOCRBLOCK0") < result_html.index( - "MARKITDOWNOCRBLOCK1" - ) - assert len(texts) == 2 + def open(self) -> io.BytesIO: + return io.BytesIO(self._name.encode()) -def test_inject_placeholders_no_img_tag_appends_at_end() -> None: - converter = DocxConverterWithOCR() - html = "

No images

" - result_html, texts = converter._inject_placeholders(html, {"rId1": "ORPHAN"}) - assert "MARKITDOWNOCRBLOCK0" in result_html - assert texts == ["ORPHAN"] - +def test_docx_ocr_failures_do_not_shift_later_occurrences() -> None: + def convert_to_html( + _stream: Any, *, convert_image: Any, **kwargs: Any + ) -> SimpleNamespace: + image_html = [] + for image_name in ("empty", "error", "later"): + element = convert_image(FakeMammothImage(image_name))[0] + placeholder = element.attributes["src"] + image_html.append( + f'

{placeholder}

' + ) + return SimpleNamespace( + value=f"

Before

{image_html[0]}

Middle

" + f"{image_html[1]}

After

{image_html[2]}

End

" + ) -def test_inject_placeholders_empty_map_leaves_html_unchanged() -> None: converter = DocxConverterWithOCR() - html = "

Content

" - result_html, texts = converter._inject_placeholders(html, {}) - assert result_html == html - assert texts == [] + with ( + patch( + "markitdown_ocr._docx_converter_with_ocr.pre_process_docx", + return_value=io.BytesIO(b"docx"), + ), + patch( + "markitdown_ocr._docx_converter_with_ocr.mammoth.convert_to_html", + side_effect=convert_to_html, + ), + ): + markdown = converter.convert( + io.BytesIO(b"docx"), + StreamInfo(extension=".docx"), + ocr_service=SequencedOCRService(), + ).text_content + + assert "MARKITDOWNOCRBLOCK" not in markdown + assert "OCR later" in markdown + assert "OCR empty" not in markdown + assert markdown.index("Middle") < markdown.index("After") < markdown.index( + "*[Image OCR]\nOCR later\n[End OCR]*" + ) # --------------------------------------------------------------------------- diff --git a/packages/markitdown-ocr/tests/test_pdf_converter.py b/packages/markitdown-ocr/tests/test_pdf_converter.py index 5d4adcc5e..feedd9f23 100644 --- a/packages/markitdown-ocr/tests/test_pdf_converter.py +++ b/packages/markitdown-ocr/tests/test_pdf_converter.py @@ -144,17 +144,30 @@ def test_pdf_complex_layout(svc: MockOCRService) -> None: # --------------------------------------------------------------------------- -# pdf_multipage.pdf — pdfplumber/pdfminer fail (EOF); PyMuPDF fallback used +# pdf_multipage.pdf — native text and inline images remain in document order # --------------------------------------------------------------------------- def test_pdf_multipage(svc: MockOCRService) -> None: - # pdfplumber cannot open this file (Unexpected EOF), so _ocr_full_pages - # falls back to PyMuPDF for page rendering. Each page becomes one OCR block. expected = ( - f"## Page 1\n\n\n{_OCR_BLOCK}\n\n\n" - f"## Page 2\n\n\n{_OCR_BLOCK}\n\n\n" - f"## Page 3\n\n\n{_OCR_BLOCK}" + "## Page 1\n\n\n" + "Page 1 - Content before image\n\n" + "This is important text that appears BEFORE the image.\n\n\n\n" + f"{_OCR_BLOCK}\n\n\n" + "This text appears AFTER the image on page 1.\n\n" + "More content follows here.\n\n\n" + "## Page 2\n\n\n" + "Page 2 - Content with image at end\n\n" + "Main content of page 2 starts here.\n\n" + "This is paragraph 1.\n\n" + "This is paragraph 2.\n\n" + "Final paragraph before image.\n\n\n\n" + f"{_OCR_BLOCK}\n\n\n\n" + "## Page 3\n\n\n" + "Page 3 - Image at top\n\n\n\n" + f"{_OCR_BLOCK}\n\n\n" + "Content that follows the image.\n\n" + "This text is AFTER the image." ) assert _convert("pdf_multipage.pdf", svc) == expected @@ -218,6 +231,74 @@ def test_pdf_scanned_fallback_format(svc: MockOCRService) -> None: ), f"_ocr_full_pages must produce:\n{expected!r}\nActual:\n{md!r}" +# --------------------------------------------------------------------------- +# Per-page fallback — mixed native/scanned pages use OCR only where needed +# --------------------------------------------------------------------------- + + +def test_pdf_page_fallback_skips_native_text_and_ocrs_empty_page( + svc: MockOCRService, +) -> None: + converter = PdfConverterWithOCR() + native_page = MagicMock() + native_page.extract_text.return_value = "Native page text" + scanned_page = MagicMock() + scanned_page.extract_text.return_value = "" + + mock_pdf = MagicMock() + mock_pdf.pages = [native_page, scanned_page] + mock_pdf.__enter__.return_value = mock_pdf + + with ( + patch("pdfplumber.open", return_value=mock_pdf), + patch.object(converter, "_extract_page_images", return_value=[]), + ): + markdown = converter.convert( + io.BytesIO(b"pdf"), StreamInfo(extension=".pdf"), ocr_service=svc + ).text_content + + assert "Native page text" in markdown + assert markdown.count(_OCR_BLOCK) == 1 + native_page.to_image.assert_not_called() + scanned_page.to_image.assert_called_once_with(resolution=300) + + +def test_pdf_page_fallback_runs_when_inline_image_ocr_is_empty() -> None: + converter = PdfConverterWithOCR() + page = MagicMock() + page.chars = [] + page.extract_text.return_value = "" + + mock_pdf = MagicMock() + mock_pdf.pages = [page] + mock_pdf.__enter__.return_value = mock_pdf + + empty_service = MagicMock() + empty_service.extract_text.side_effect = [ + OCRResult(text="", backend_used="mock"), + OCRResult(text="Full page OCR", backend_used="mock"), + ] + + with ( + patch("pdfplumber.open", return_value=mock_pdf), + patch.object( + converter, + "_extract_page_images", + return_value=[ + {"stream": io.BytesIO(b"image"), "y_pos": 0, "name": "image"} + ], + ), + ): + markdown = converter.convert( + io.BytesIO(b"pdf"), + StreamInfo(extension=".pdf"), + ocr_service=empty_service, + ).text_content + + assert "Full page OCR" in markdown + page.to_image.assert_called_once_with(resolution=300) + + # --------------------------------------------------------------------------- # No OCR service — no OCR tags emitted # ---------------------------------------------------------------------------