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..c67093b33 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 @@ -127,31 +127,39 @@ def _extract_and_ocr_images( self, file_stream: BinaryIO, ocr_service: LLMVisionOCRService ) -> dict[str, str]: """ - Extract images from DOCX and OCR them. + Extract images from DOCX and OCR them in parallel. Returns: Dict mapping image relationship IDs to raw OCR text (no markers). """ - ocr_map = {} + ocr_map: dict[str, str] = {} try: file_stream.seek(0) doc = Document(file_stream) + # Phase 1: Collect all image streams with their rIds + image_specs: list[tuple[str, BinaryIO]] = [] # (rId, 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() - + image_specs.append((rel.rId, image_stream)) except Exception: continue + # Phase 2: Batch OCR all images in parallel + if image_specs: + rids = [spec[0] for spec in image_specs] + streams = [spec[1] for spec in image_specs] + results = ocr_service.extract_text_batch( + [(s, None) for s in streams] + ) + for rId, result in zip(rids, results): + if result.text.strip(): + ocr_map[rId] = result.text.strip() + except Exception: pass diff --git a/packages/markitdown-ocr/src/markitdown_ocr/_ocr_service.py b/packages/markitdown-ocr/src/markitdown_ocr/_ocr_service.py index 2885e1f47..879e75c5c 100644 --- a/packages/markitdown-ocr/src/markitdown_ocr/_ocr_service.py +++ b/packages/markitdown-ocr/src/markitdown_ocr/_ocr_service.py @@ -4,6 +4,7 @@ """ import base64 +from concurrent.futures import ThreadPoolExecutor, as_completed from typing import Any, BinaryIO from dataclasses import dataclass @@ -28,6 +29,7 @@ def __init__( client: Any, model: str, default_prompt: str | None = None, + max_workers: int = 5, ) -> None: """ Initialize LLM Vision OCR service. @@ -36,9 +38,11 @@ def __init__( client: OpenAI-compatible client model: Model name (e.g., 'gpt-4o', 'gemini-2.0-flash') default_prompt: Default prompt for OCR extraction + max_workers: Maximum number of parallel OCR workers (default 5) """ self.client = client self.model = model + self.max_workers = max_workers self.default_prompt = default_prompt or ( "Extract all text from this image. " "Return ONLY the extracted text, maintaining the original " @@ -108,3 +112,59 @@ def extract_text( return OCRResult(text="", backend_used="llm_vision", error=str(e)) finally: image_stream.seek(0) + + def extract_text_batch( + self, + images: list[tuple[BinaryIO, StreamInfo | None]], + prompt: str | None = None, + **kwargs: Any, + ) -> list[OCRResult]: + """ + Extract text from multiple images in parallel using a thread pool. + + Args: + images: List of (image_stream, stream_info) tuples + prompt: Optional prompt override (shared across all images) + **kwargs: Additional arguments (max_workers override, etc.) + + Returns: + List of OCRResult, one per input image (same order as input). + Individual failures are captured in the OCRResult's error field + rather than raising an exception. + """ + if not images: + return [] + + max_workers = kwargs.get("max_workers", self.max_workers) + + def _process_one( + idx: int, img_stream: BinaryIO, stream_info: StreamInfo | None + ) -> tuple[int, OCRResult]: + """Wrapper that catches all exceptions so one failure doesn't + kill the entire batch.""" + try: + result = self.extract_text( + img_stream, prompt=prompt, stream_info=stream_info + ) + except Exception as e: + result = OCRResult( + text="", backend_used="llm_vision", error=str(e) + ) + return idx, result + + results: list[OCRResult | None] = [None] * len(images) + + with ThreadPoolExecutor(max_workers=max_workers) as executor: + futures = { + executor.submit(_process_one, i, img, info): i + for i, (img, info) in enumerate(images) + } + for future in as_completed(futures): + idx, result = future.result() + results[idx] = result + + # Defensive: any slot still None → empty result + return [ + r if r is not None else OCRResult(text="", backend_used="llm_vision") + for r in results + ] 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..ae6d29a9d 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 @@ -126,6 +126,52 @@ def _extract_images_from_page(page: Any) -> list[dict]: return images_info +def _extract_text_lines_from_page(page: Any) -> list[dict]: + """ + Extract text lines with Y positions from a PDF page. + + Uses pdfplumber char-level data to group characters into lines + based on their vertical position. + + Args: + page: pdfplumber page object + + Returns: + List of dicts with 'y' (float) and 'text' (str) keys, sorted top-to-bottom + """ + chars = page.chars + if not chars: + # Fallback: use simple text extraction + text_content = page.extract_text() or "" + return [ + {"y": i * 10, "text": line} + for i, line in enumerate(text_content.split("\n")) + ] + + lines_with_y: list[dict] = [] + current_line: list[Any] = [] + current_y: float | None = 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 or 0, "text": text.strip()}) + + return lines_with_y + + class PdfConverterWithOCR(DocumentConverter): """ Enhanced PDF Converter with OCR support for embedded images. @@ -181,107 +227,81 @@ def convert( file_stream.seek(0) pdf_bytes = io.BytesIO(file_stream.read()) - markdown_content = [] + markdown_content: list[str] = [] try: with pdfplumber.open(pdf_bytes) as pdf: - for page_num, page in enumerate(pdf.pages, 1): - markdown_content.append(f"\n## Page {page_num}\n") - - # If OCR is enabled, interleave text and images by position - if ocr_service: - images_on_page = self._extract_page_images(pdf_bytes, page_num) - - if images_on_page: - # Extract text lines with Y positions - chars = page.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()} - ) - else: - # Fallback: use simple text extraction - text_content = page.extract_text() or "" - lines_with_y = [ - {"y": i * 10, "text": line} - for i, line in enumerate(text_content.split("\n")) - ] - - # OCR all images - image_data = [] - for img_info in images_on_page: - ocr_result = ocr_service.extract_text( - img_info["stream"] + if ocr_service: + # ── Phase 1: Collect all images from all pages ── + # all_images: (image_stream, y_pos, page_num, img_name) + all_images: list[tuple[BinaryIO, float, int, str]] = [] + # page_text_lines: page_num -> list of {y, text} dicts + page_text_lines: dict[int, list[dict]] = {} + + for page_num, page in enumerate(pdf.pages, 1): + page_text_lines[page_num] = _extract_text_lines_from_page(page) + for img_info in _extract_images_from_page(page): + all_images.append( + ( + img_info["stream"], + img_info["y_pos"], + page_num, + img_info["name"], ) - 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", - } - ) + ) - # Add text items - content_items = [ - { - "y_pos": item["y"], - "text": item["text"], - "type": "text", - } - for item in lines_with_y + # ── Phase 2: Batch OCR all images in parallel ── + ocr_by_page: dict[int, list[tuple[float, str]]] = {} + if all_images: + ocr_results = ocr_service.extract_text_batch( + [(stream, None) for stream, _, _, _ in all_images] + ) + for i, (_, y_pos, pg, name) in enumerate(all_images): + text = ocr_results[i].text.strip() + if text: + ocr_by_page.setdefault(pg, []).append((y_pos, text)) + + # ── Phase 3: Build output per page (interleave text + OCR) ── + for page_num, page in enumerate(pdf.pages, 1): + markdown_content.append(f"\n## Page {page_num}\n") + + page_ocr = ocr_by_page.get(page_num, []) + + if page_ocr: + # Build items: text lines + OCR blocks + content_items: list[dict] = [ + {"y_pos": item["y"], "text": item["text"], "type": "text"} + for item in page_text_lines[page_num] if item["text"] ] - content_items.extend(image_data) + for y_pos, ocr_text in page_ocr: + content_items.append( + { + "y_pos": y_pos, + "ocr_text": ocr_text, + "type": "image", + } + ) # Sort all items by Y position (top to bottom) content_items.sort(key=lambda x: x["y_pos"]) - # Build markdown by interleaving text and images for item in content_items: if item["type"] == "text": markdown_content.append(item["text"]) - else: # image - ocr_text = item["ocr_text"] - img_marker = ( - f"\n\n*[Image OCR]\n{ocr_text}\n[End OCR]*\n" + else: + markdown_content.append( + f"\n\n*[Image OCR]\n{item['ocr_text']}\n[End OCR]*\n" ) - markdown_content.append(img_marker) else: - # No images detected - just extract regular text + # No images on this page — just extract text text_content = page.extract_text() or "" if text_content.strip(): markdown_content.append(text_content.strip()) - else: - # No OCR, just extract text + else: + # No OCR — simple text extraction + for page_num, page in enumerate(pdf.pages, 1): + markdown_content.append(f"\n## Page {page_num}\n") text_content = page.extract_text() or "" if text_content.strip(): markdown_content.append(text_content.strip()) @@ -302,120 +322,195 @@ def convert( except Exception: markdown = "" - # Final fallback: If still empty/whitespace and OCR is available, - # treat as scanned PDF and OCR full pages - if ocr_service and (not markdown or not markdown.strip()): - pdf_bytes.seek(0) - markdown = self._ocr_full_pages(pdf_bytes, ocr_service) - - return DocumentConverterResult(markdown=markdown) + # Final fallback: If output is empty or contains only page-number + # headers (scanned PDF with no extractable text), use full-page OCR. + if ocr_service: + # Strip out page-header boilerplate to test for real content + import re as _re # local import — only used here - def _extract_page_images(self, pdf_bytes: io.BytesIO, page_num: int) -> list[dict]: - """ - Extract images from a PDF page using pdfplumber. - - Args: - pdf_bytes: PDF file as BytesIO - page_num: Page number (1-indexed) - - Returns: - List of image info dicts with 'stream', 'bbox', 'name', 'y_pos' - """ - images = [] - - try: - pdf_bytes.seek(0) - with pdfplumber.open(pdf_bytes) as pdf: - if page_num <= len(pdf.pages): - page = pdf.pages[page_num - 1] # 0-indexed - images = _extract_images_from_page(page) - except Exception: - pass - - # Sort by vertical position (top to bottom) - images.sort(key=lambda x: x["y_pos"]) + _real_content = _re.sub( + r"\s*## Page \d+\s*", "", markdown + ).strip() + if not _real_content: + pdf_bytes.seek(0) + ocr_dpi = kwargs.get("ocr_dpi", 300) + markdown = self._ocr_full_pages(pdf_bytes, ocr_service, ocr_dpi=ocr_dpi) - return images + return DocumentConverterResult(markdown=markdown) def _ocr_full_pages( - self, pdf_bytes: io.BytesIO, ocr_service: LLMVisionOCRService + self, pdf_bytes: io.BytesIO, ocr_service: LLMVisionOCRService, + ocr_dpi: int = 300, ) -> str: """ - Fallback for scanned PDFs: Convert entire pages to images and OCR them. - Used when text extraction returns empty/whitespace results. - - Args: - pdf_bytes: PDF file as BytesIO - ocr_service: OCR service to use + Fallback for scanned PDFs: render pages to images and OCR them. - Returns: - Markdown text extracted from OCR of full pages + Uses a streaming pipeline: each page is submitted for OCR the moment + its render finishes. Rendering and OCR run concurrently, so total + time ≈ max(render_all, ocr_all) instead of render_all + ocr_all. """ - markdown_parts = [] + from concurrent.futures import ThreadPoolExecutor as _TPE, as_completed as _ac + from ._ocr_service import OCRResult + + markdown_parts: list[str] = [] + ocr_results: dict[int, str] = {} # page_num -> text try: 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]*" - ) + n_pages = len(pdf.pages) + if n_pages == 0: + return "" - except Exception as e: - markdown_parts.append( - f"*[Error processing page {page_num}: {str(e)}]*" - ) - continue + # ── Single-page fast path (no thread overhead) ── + if n_pages == 1: + try: + pg = pdf.pages[0] + pg_img = pg.to_image(resolution=ocr_dpi) + buf = io.BytesIO() + pg_img.original.save(buf, format="PNG") + buf.seek(0) + result = ocr_service.extract_text(buf) + if result.text.strip(): + ocr_results[1] = result.text.strip() + except Exception: + pass + else: + # ── Multi-page streaming pipeline ── + render_workers = min(4, n_pages) + ocr_workers = ocr_service.max_workers + + def _render_page(pg_num: int) -> tuple[int, BinaryIO | None]: + try: + pg = pdf.pages[pg_num - 1] + pg_img = pg.to_image(resolution=ocr_dpi) + buf = io.BytesIO() + pg_img.original.save(buf, format="PNG") + buf.seek(0) + return pg_num, buf + except Exception: + return pg_num, None + + def _ocr_one(img_stream: BinaryIO) -> OCRResult: + try: + return ocr_service.extract_text(img_stream) + except Exception as e: + return OCRResult(text="", error=str(e)) + + with _TPE(max_workers=render_workers) as render_pool: + ocr_futures: dict = {} + render_futures = { + render_pool.submit(_render_page, i): i + for i in range(1, n_pages + 1) + } + with _TPE(max_workers=ocr_workers) as ocr_pool: + for render_future in _ac(render_futures): + pg_num, buf = render_future.result() + if buf is not None: + ocr_futures[ + ocr_pool.submit(_ocr_one, buf) + ] = pg_num + for ocr_future in _ac(ocr_futures): + pg_num = ocr_futures[ocr_future] + result = ocr_future.result() + if result.text.strip(): + ocr_results[pg_num] = result.text.strip() + + # ── Assemble output in page order ── + for pg_num in sorted(ocr_results): + markdown_parts.append(f"\n## Page {pg_num}\n") + markdown_parts.append( + f"*[Image OCR]\n{ocr_results[pg_num]}\n[End OCR]*" + ) + # Pages that failed to produce text + for pg_num in range(1, n_pages + 1): + if pg_num not in ocr_results: + markdown_parts.append(f"\n## Page {pg_num}\n") + markdown_parts.append( + "*[No text could be extracted from this page]*" + ) except Exception: - # pdfplumber failed (e.g. malformed EOF) — try PyMuPDF for rendering + # pdfplumber failed — try PyMuPDF with same streaming pipeline markdown_parts = [] + ocr_results = {} try: import fitz # PyMuPDF pdf_bytes.seek(0) doc = fitz.open(stream=pdf_bytes.read(), filetype="pdf") - for page_num in range(1, doc.page_count + 1): - try: - markdown_parts.append(f"\n## Page {page_num}\n") - page = doc[page_num - 1] - mat = fitz.Matrix(300 / 72, 300 / 72) # 300 DPI - pix = page.get_pixmap(matrix=mat) - img_stream = io.BytesIO(pix.tobytes("png")) - img_stream.seek(0) - - ocr_result = ocr_service.extract_text(img_stream) + n_pages = doc.page_count + if n_pages == 0: + doc.close() + return "" - 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]*" - ) + # ── Single-page fast path ── + if n_pages == 1: + try: + pg = doc[0] + mat = fitz.Matrix(ocr_dpi / 72, ocr_dpi / 72) + pix = pg.get_pixmap(matrix=mat) + buf = io.BytesIO(pix.tobytes("png")) + buf.seek(0) + result = ocr_service.extract_text(buf) + if result.text.strip(): + ocr_results[1] = result.text.strip() + except Exception: + pass + else: + # ── Multi-page streaming pipeline ── + render_workers = min(4, n_pages) + ocr_workers = ocr_service.max_workers + + def _render_mupdf(pg_num: int) -> tuple[int, BinaryIO | None]: + try: + pg = doc[pg_num - 1] + mat = fitz.Matrix(ocr_dpi / 72, ocr_dpi / 72) + pix = pg.get_pixmap(matrix=mat) + buf = io.BytesIO(pix.tobytes("png")) + buf.seek(0) + return pg_num, buf + except Exception: + return pg_num, None + + def _ocr_one_mupdf(img_stream: BinaryIO) -> OCRResult: + try: + return ocr_service.extract_text(img_stream) + except Exception as e: + return OCRResult(text="", error=str(e)) + + with _TPE(max_workers=render_workers) as render_pool: + ocr_futures = {} + render_futures = { + render_pool.submit(_render_mupdf, i): i + for i in range(1, n_pages + 1) + } + with _TPE(max_workers=ocr_workers) as ocr_pool: + for render_future in _ac(render_futures): + pg_num, buf = render_future.result() + if buf is not None: + ocr_futures[ + ocr_pool.submit(_ocr_one_mupdf, buf) + ] = pg_num + for ocr_future in _ac(ocr_futures): + pg_num = ocr_futures[ocr_future] + result = ocr_future.result() + if result.text.strip(): + ocr_results[pg_num] = result.text.strip() + doc.close() - except Exception as e: + for pg_num in sorted(ocr_results): + markdown_parts.append(f"\n## Page {pg_num}\n") + markdown_parts.append( + f"*[Image OCR]\n{ocr_results[pg_num]}\n[End OCR]*" + ) + for pg_num in range(1, n_pages + 1): + if pg_num not in ocr_results: + markdown_parts.append(f"\n## Page {pg_num}\n") markdown_parts.append( - f"*[Error processing page {page_num}: {str(e)}]*" + "*[No text could be extracted from this page]*" ) - continue - doc.close() + except Exception: return "*[Error: Could not process scanned PDF]*" diff --git a/packages/markitdown-ocr/src/markitdown_ocr/_plugin.py b/packages/markitdown-ocr/src/markitdown_ocr/_plugin.py index f4d7bcf5a..c0308a45a 100644 --- a/packages/markitdown-ocr/src/markitdown_ocr/_plugin.py +++ b/packages/markitdown-ocr/src/markitdown_ocr/_plugin.py @@ -37,6 +37,7 @@ def register_converters(markitdown: MarkItDown, **kwargs: Any) -> None: llm_client = kwargs.get("llm_client") llm_model = kwargs.get("llm_model") llm_prompt = kwargs.get("llm_prompt") + max_workers = kwargs.get("max_workers", 5) ocr_service: LLMVisionOCRService | None = None if llm_client and llm_model: @@ -44,6 +45,7 @@ def register_converters(markitdown: MarkItDown, **kwargs: Any) -> None: client=llm_client, model=llm_model, default_prompt=llm_prompt, + max_workers=max_workers, ) # Register converters with priority -1.0 (before built-ins at 0.0) diff --git a/packages/markitdown-ocr/src/markitdown_ocr/_pptx_converter_with_ocr.py b/packages/markitdown-ocr/src/markitdown_ocr/_pptx_converter_with_ocr.py index 7e91ed6b4..47d4fc1ce 100644 --- a/packages/markitdown-ocr/src/markitdown_ocr/_pptx_converter_with_ocr.py +++ b/packages/markitdown-ocr/src/markitdown_ocr/_pptx_converter_with_ocr.py @@ -5,10 +5,9 @@ import io import sys +from concurrent.futures import ThreadPoolExecutor, as_completed from typing import Any, BinaryIO, Optional -from typing import BinaryIO, Any, Optional - from markitdown.converters import HtmlConverter from markitdown import DocumentConverter, DocumentConverterResult, StreamInfo from markitdown._exceptions import ( @@ -73,73 +72,123 @@ def convert( kwargs.get("ocr_service") or self.ocr_service ) llm_client = kwargs.get("llm_client") + llm_model = kwargs.get("llm_model") + llm_prompt = kwargs.get("llm_prompt") presentation = pptx.Presentation(file_stream) + + # ── Pre-scan: collect all image shapes across all slides ── + # Each entry: (image_stream, content_type, filename) + image_entries: list[tuple[BinaryIO, str | None, str | None]] = [] + + for slide in presentation.slides: + sorted_shapes = sorted( + slide.shapes, + key=lambda x: ( + float("-inf") if not x.top else x.top, + float("-inf") if not x.left else x.left, + ), + ) + self._collect_image_shapes(sorted_shapes, image_entries) + + # ── Parallel image processing ── + # image_results[i] = final text for the i-th image shape + image_results: list[str] = [""] * len(image_entries) + + if image_entries: + # Round 1: Parallel LLM caption for all images + if llm_client and llm_model: + from markitdown.converters._llm_caption import llm_caption + import os + + def _caption_one( + idx: int, + img_stream: BinaryIO, + content_type: str | None, + filename: str | None, + ) -> tuple[int, str]: + try: + image_extension = None + if filename: + image_extension = os.path.splitext(filename)[1] + s_info = StreamInfo( + mimetype=content_type, + extension=image_extension, + filename=filename, + ) + result = llm_caption( + img_stream, + s_info, + client=llm_client, + model=llm_model, + prompt=llm_prompt, + ) + return idx, result or "" + except Exception: + return idx, "" + + max_w = ocr_service.max_workers if ocr_service else 5 + with ThreadPoolExecutor(max_workers=max_w) as executor: + futures = { + executor.submit( + _caption_one, i, stream, ct, fn + ): i + for i, (stream, ct, fn) in enumerate(image_entries) + } + for future in as_completed(futures): + idx, text = future.result() + image_results[idx] = text + + # Round 2: Parallel OCR for images where LLM caption failed + if ocr_service: + ocr_indices: list[int] = [] + ocr_streams: list[tuple[BinaryIO, StreamInfo | None]] = [] + for i in range(len(image_entries)): + if not image_results[i].strip(): + ocr_indices.append(i) + stream, _ct, _fn = image_entries[i] + stream.seek(0) + ocr_streams.append((stream, None)) + + if ocr_streams: + ocr_results = ocr_service.extract_text_batch(ocr_streams) + for j, result in enumerate(ocr_results): + if result.text.strip(): + image_results[ocr_indices[j]] = result.text.strip() + + # ── Render slides with cached image results ── md_content = "" slide_num = 0 + img_cursor = [0] # mutable counter for consuming image_results for slide in presentation.slides: slide_num += 1 - md_content += f"\\n\\n\\n" + md_content += f"\n\n\n" title = slide.shapes.title def get_shape_content(shape, **kwargs): nonlocal md_content - # Pictures + # Pictures — use pre-computed result if self._is_picture(shape): - # Get image data - image_stream = io.BytesIO(shape.image.blob) - - # Try LLM description first if available - llm_description = "" - if llm_client and kwargs.get("llm_model"): - try: - from ._llm_caption import llm_caption - - image_filename = shape.image.filename - image_extension = None - if image_filename: - import os - - image_extension = os.path.splitext(image_filename)[1] - - image_stream_info = StreamInfo( - mimetype=shape.image.content_type, - extension=image_extension, - filename=image_filename, - ) - - llm_description = llm_caption( - image_stream, - image_stream_info, - client=llm_client, - model=kwargs.get("llm_model"), - prompt=kwargs.get("llm_prompt"), - ) - except Exception: - pass - - # Try OCR if LLM failed or not available - ocr_text = "" - if not llm_description and ocr_service: - try: - image_stream.seek(0) - ocr_result = ocr_service.extract_text(image_stream) - if ocr_result.text.strip(): - ocr_text = ocr_result.text.strip() - except Exception: - pass - - # Format extracted content using unified OCR block format - content = (llm_description or ocr_text or "").strip() + idx = img_cursor[0] + img_cursor[0] += 1 + content = ( + image_results[idx].strip() + if idx < len(image_results) + else "" + ) if content: - md_content += f"\n*[Image OCR]\n{content}\n[End OCR]*\n" + md_content += ( + f"\n*[Image OCR]\n{content}\n[End OCR]*\n" + ) # Tables if self._is_table(shape): - md_content += self._convert_table_to_markdown(shape.table, **kwargs) + md_content += self._convert_table_to_markdown( + shape.table, **kwargs + ) # Charts if shape.has_chart: @@ -148,9 +197,9 @@ def get_shape_content(shape, **kwargs): # Text areas elif shape.has_text_frame: if shape == title: - md_content += "# " + shape.text.lstrip() + "\\n" + md_content += "# " + shape.text.lstrip() + "\n" else: - md_content += shape.text + "\\n" + md_content += shape.text + "\n" # Group Shapes if shape.shape_type == pptx.enum.shapes.MSO_SHAPE_TYPE.GROUP: @@ -177,7 +226,7 @@ def get_shape_content(shape, **kwargs): md_content = md_content.strip() if slide.has_notes_slide: - md_content += "\\n\\n### Notes:\\n" + md_content += "\n\n### Notes:\n" notes_frame = slide.notes_slide.notes_text_frame if notes_frame is not None: md_content += notes_frame.text @@ -185,6 +234,35 @@ def get_shape_content(shape, **kwargs): return DocumentConverterResult(markdown=md_content.strip()) + def _collect_image_shapes( + self, + shapes: Any, + images: list[tuple[BinaryIO, str | None, str | None]], + ) -> None: + """Recursively walk shapes and collect image data into the images list. + The traversal order MUST match the rendering order so that image_results + indices align between the pre-scan and render phases. + """ + for shape in shapes: + if self._is_picture(shape): + image_stream = io.BytesIO(shape.image.blob) + images.append( + ( + image_stream, + shape.image.content_type, + shape.image.filename, + ) + ) + if shape.shape_type == pptx.enum.shapes.MSO_SHAPE_TYPE.GROUP: + sorted_sub = sorted( + shape.shapes, + key=lambda x: ( + float("-inf") if not x.top else x.top, + float("-inf") if not x.left else x.left, + ), + ) + self._collect_image_shapes(sorted_sub, images) + def _is_picture(self, shape): if shape.shape_type == pptx.enum.shapes.MSO_SHAPE_TYPE.PICTURE: return True @@ -216,15 +294,15 @@ def _convert_table_to_markdown(self, table, **kwargs): return ( self._html_converter.convert_string(html_table, **kwargs).markdown.strip() - + "\\n" + + "\n" ) def _convert_chart_to_markdown(self, chart): try: - md = "\\n\\n### Chart" + md = "\n\n### Chart" if chart.has_title: md += f": {chart.chart_title.text_frame.text}" - md += "\\n\\n" + md += "\n\n" data = [] category_names = [c.label for c in chart.plots[0].categories] series_names = [s.name for s in chart.series] @@ -241,9 +319,9 @@ def _convert_chart_to_markdown(self, chart): markdown_table.append("| " + " | ".join(map(str, row)) + " |") header = markdown_table[0] separator = "|" + "|".join(["---"] * len(data[0])) + "|" - return md + "\\n".join([header, separator] + markdown_table[1:]) + return md + "\n".join([header, separator] + markdown_table[1:]) except ValueError as e: if "unsupported plot type" in str(e): - return "\\n\\n[unsupported chart]\\n\\n" + return "\n\n[unsupported chart]\n\n" except Exception: - return "\\n\\n[unsupported chart]\\n\\n" + return "\n\n[unsupported chart]\n\n" diff --git a/packages/markitdown-ocr/src/markitdown_ocr/_xlsx_converter_with_ocr.py b/packages/markitdown-ocr/src/markitdown_ocr/_xlsx_converter_with_ocr.py index 481e07195..23af90f1e 100644 --- a/packages/markitdown-ocr/src/markitdown_ocr/_xlsx_converter_with_ocr.py +++ b/packages/markitdown-ocr/src/markitdown_ocr/_xlsx_converter_with_ocr.py @@ -108,61 +108,88 @@ def _convert_standard( def _convert_with_ocr( self, file_stream: BinaryIO, ocr_service: LLMVisionOCRService, **kwargs: Any ) -> DocumentConverterResult: - """Convert XLSX with image OCR.""" + """Convert XLSX with image OCR (batch-parallel).""" file_stream.seek(0) wb = load_workbook(file_stream) - md_content = "" + # ── Phase 1: Collect all images from all sheets ── + # all_images: (image_stream, sheet_name, cell_ref) + all_images: list[tuple[BinaryIO, str, str]] = [] for sheet_name in wb.sheetnames: sheet = wb[sheet_name] - md_content += f"## {sheet_name}\n\n" + for img_stream, cell_ref in self._extract_sheet_images(sheet): + all_images.append((img_stream, sheet_name, cell_ref)) + + # ── Phase 2: Batch OCR all images in parallel ── + ocr_by_sheet: dict[str, list[dict]] = {} + if all_images: + ocr_results = ocr_service.extract_text_batch( + [(stream, None) for stream, _, _ in all_images] + ) + for i, (_, sheet_name, cell_ref) in enumerate(all_images): + text = ocr_results[i].text.strip() + if text: + ocr_by_sheet.setdefault(sheet_name, []).append( + {"cell_ref": cell_ref, "ocr_text": text} + ) + + # ── Phase 3: Render sheets ── + # Read every sheet in a single pass (sheet_name=None) instead of + # re-parsing the whole file once per sheet. + file_stream.seek(0) + try: + all_dfs = pd.read_excel( + file_stream, sheet_name=None, engine="openpyxl" + ) + except Exception: + all_dfs = {} - # Convert sheet data to markdown table - file_stream.seek(0) - try: - df = pd.read_excel( - file_stream, sheet_name=sheet_name, engine="openpyxl" - ) - html_content = df.to_html(index=False) - md_content += ( - self._html_converter.convert_string( - html_content, **kwargs - ).markdown.strip() - + "\n\n" - ) - except Exception: - # If pandas fails, just skip the table - pass + md_content = "" - # Extract and OCR images in this sheet - images_with_ocr = self._extract_and_ocr_sheet_images(sheet, ocr_service) + for sheet_name in wb.sheetnames: + md_content += f"## {sheet_name}\n\n" + # Convert sheet data to markdown table + df = all_dfs.get(sheet_name) + if df is not None: + try: + html_content = df.to_html(index=False) + md_content += ( + self._html_converter.convert_string( + html_content, **kwargs + ).markdown.strip() + + "\n\n" + ) + except Exception: + pass + + # Append pre-computed OCR results for this sheet + images_with_ocr = ocr_by_sheet.get(sheet_name, []) if images_with_ocr: md_content += "### Images in this sheet:\n\n" for img_info in images_with_ocr: - ocr_text = img_info["ocr_text"] - md_content += f"*[Image OCR]\n{ocr_text}\n[End OCR]*\n\n" + md_content += ( + f"*[Image OCR]\n{img_info['ocr_text']}\n[End OCR]*\n\n" + ) return DocumentConverterResult(markdown=md_content.strip()) - def _extract_and_ocr_sheet_images( - self, sheet: Any, ocr_service: LLMVisionOCRService - ) -> list[dict]: + def _extract_sheet_images( + self, sheet: Any + ) -> list[tuple[BinaryIO, str]]: """ - Extract and OCR images from an Excel sheet. + Extract images from an Excel sheet (no OCR — just collection). Args: sheet: openpyxl worksheet - ocr_service: OCR service - Returns: - List of dicts with 'cell_ref' and 'ocr_text' + Yields: + (image_stream, cell_ref) tuples, one per image in the sheet """ - results = [] + results: list[tuple[BinaryIO, str]] = [] try: - # Check if sheet has images if hasattr(sheet, "_images"): for img in sheet._images: try: @@ -170,12 +197,10 @@ def _extract_and_ocr_sheet_images( if hasattr(img, "_data"): image_data = img._data() elif hasattr(img, "image"): - # Some versions store it differently image_data = img.image else: continue - # Create image stream image_stream = io.BytesIO(image_data) # Get cell reference @@ -187,27 +212,15 @@ def _extract_and_ocr_sheet_images( if hasattr(from_cell, "col") and hasattr( from_cell, "row" ): - # Convert column number to letter col_letter = self._column_number_to_letter( from_cell.col ) cell_ref = f"{col_letter}{from_cell.row + 1}" - # Perform OCR - ocr_result = ocr_service.extract_text(image_stream) - - if ocr_result.text.strip(): - results.append( - { - "cell_ref": cell_ref, - "ocr_text": ocr_result.text.strip(), - "backend": ocr_result.backend_used, - } - ) + results.append((image_stream, cell_ref)) except Exception: continue - except Exception: pass diff --git a/packages/markitdown-ocr/tests/test_docx_converter.py b/packages/markitdown-ocr/tests/test_docx_converter.py index 0fb666504..8731205e2 100644 --- a/packages/markitdown-ocr/tests/test_docx_converter.py +++ b/packages/markitdown-ocr/tests/test_docx_converter.py @@ -35,6 +35,13 @@ def extract_text( # noqa: ANN101 ) -> OCRResult: return OCRResult(text=_MOCK_TEXT, backend_used="mock") + def extract_text_batch( + self, + images: list[tuple[Any, Any]], + **kwargs: Any, + ) -> list[OCRResult]: + return [self.extract_text(stream) for stream, _ in images] + @pytest.fixture(scope="module") def svc() -> MockOCRService: diff --git a/packages/markitdown-ocr/tests/test_pdf_converter.py b/packages/markitdown-ocr/tests/test_pdf_converter.py index 5d4adcc5e..73c9d134e 100644 --- a/packages/markitdown-ocr/tests/test_pdf_converter.py +++ b/packages/markitdown-ocr/tests/test_pdf_converter.py @@ -41,6 +41,13 @@ def extract_text( ) -> OCRResult: return OCRResult(text=_MOCK_TEXT, backend_used="mock") + def extract_text_batch( + self, + images: list[tuple[Any, Any]], + **kwargs: Any, + ) -> list[OCRResult]: + return [self.extract_text(stream) for stream, _ in images] + @pytest.fixture(scope="module") def svc() -> MockOCRService: @@ -144,17 +151,32 @@ def test_pdf_complex_layout(svc: MockOCRService) -> None: # --------------------------------------------------------------------------- -# pdf_multipage.pdf — pdfplumber/pdfminer fail (EOF); PyMuPDF fallback used +# pdf_multipage.pdf — 3-page PDF with text + embedded images +# The batch-parallel path opens the PDF once and extracts text + images +# in a single pass, correctly interleaving them by Y position. # --------------------------------------------------------------------------- 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 diff --git a/packages/markitdown-ocr/tests/test_pptx_converter.py b/packages/markitdown-ocr/tests/test_pptx_converter.py index 724f1039c..c303bf933 100644 --- a/packages/markitdown-ocr/tests/test_pptx_converter.py +++ b/packages/markitdown-ocr/tests/test_pptx_converter.py @@ -9,7 +9,7 @@ MOCK_OCR_TEXT_12345 [End OCR]* -Note: PPTX slide text uses literal backslash-n (\\n) sequences from the +Note: PPTX slide text uses literal backslash-n (\n) sequences from the underlying PPTX converter template; OCR blocks use real newlines. """ @@ -41,6 +41,13 @@ def extract_text( ) -> OCRResult: return OCRResult(text=_MOCK_TEXT, backend_used="mock") + def extract_text_batch( + self, + images: list[tuple[Any, Any]], + **kwargs: Any, + ) -> list[OCRResult]: + return [self.extract_text(stream) for stream, _ in images] + @pytest.fixture(scope="module") def svc() -> MockOCRService: @@ -66,7 +73,7 @@ def _convert(filename: str, ocr_service: MockOCRService) -> str: def test_pptx_image_start(svc: MockOCRService) -> None: # Slide 1: title "Welcome" followed by an image expected = ( - "\\n\\n\\n# Welcome\\n\\n" + "\n# Welcome\n\n" "\n*[Image OCR]\nMOCK_OCR_TEXT_12345\n[End OCR]*" ) assert _convert("pptx_image_start.pptx", svc) == expected @@ -80,10 +87,10 @@ def test_pptx_image_start(svc: MockOCRService) -> None: def test_pptx_image_middle(svc: MockOCRService) -> None: # Slide 1: Introduction | Slide 2: Architecture + image | Slide 3: Conclusion # noqa: E501 expected = ( - "\\n\\n\\n# Introduction" - "\\n\\n\\n\\n\\n# Architecture\\n\\n" + "\n# Introduction" + "\n\n\n# Architecture\n\n" "\n*[Image OCR]\nMOCK_OCR_TEXT_12345\n[End OCR]*" - "\\n\\n\\n# Conclusion\\n\\n" + "\n\n\n# Conclusion" ) assert _convert("pptx_image_middle.pptx", svc) == expected @@ -96,8 +103,8 @@ def test_pptx_image_middle(svc: MockOCRService) -> None: def test_pptx_image_end(svc: MockOCRService) -> None: # Slide 1: Presentation | Slide 2: Thank You + image expected = ( - "\\n\\n\\n# Presentation" - "\\n\\n\\n\\n\\n# Thank You\\n\\n" + "\n# Presentation" + "\n\n\n# Thank You\n\n" "\n*[Image OCR]\nMOCK_OCR_TEXT_12345\n[End OCR]*" ) assert _convert("pptx_image_end.pptx", svc) == expected @@ -111,8 +118,8 @@ def test_pptx_image_end(svc: MockOCRService) -> None: def test_pptx_multiple_images(svc: MockOCRService) -> None: # Slide 1: two images, no title text expected = ( - "\\n\\n\\n# \\n" - "\n*[Image OCR]\nMOCK_OCR_TEXT_12345\n[End OCR]*" + "\n# " + "\n\n*[Image OCR]\nMOCK_OCR_TEXT_12345\n[End OCR]*" "\n\n*[Image OCR]\nMOCK_OCR_TEXT_12345\n[End OCR]*" ) assert _convert("pptx_multiple_images.pptx", svc) == expected @@ -125,9 +132,9 @@ def test_pptx_multiple_images(svc: MockOCRService) -> None: def test_pptx_complex_layout(svc: MockOCRService) -> None: expected = ( - "\\n\\n\\n# Product Comparison" - "\\n\\nOur products lead the market\\n" - "\n*[Image OCR]\nMOCK_OCR_TEXT_12345\n[End OCR]*" + "\n# Product Comparison" + "\n\nOur products lead the market" + "\n\n*[Image OCR]\nMOCK_OCR_TEXT_12345\n[End OCR]*" ) assert _convert("pptx_complex_layout.pptx", svc) == expected diff --git a/packages/markitdown-ocr/tests/test_xlsx_converter.py b/packages/markitdown-ocr/tests/test_xlsx_converter.py index 4ab30c600..c3903e89e 100644 --- a/packages/markitdown-ocr/tests/test_xlsx_converter.py +++ b/packages/markitdown-ocr/tests/test_xlsx_converter.py @@ -42,6 +42,13 @@ def extract_text( ) -> OCRResult: return OCRResult(text=_MOCK_TEXT, backend_used="mock") + def extract_text_batch( + self, + images: list[tuple[Any, Any]], + **kwargs: Any, + ) -> list[OCRResult]: + return [self.extract_text(stream) for stream, _ in images] + @pytest.fixture(scope="module") def svc() -> MockOCRService: