|
1 | | -""" |
2 | | -Enhanced DOCX Converter with OCR support for embedded images. |
3 | | -Extracts images from Word documents and performs OCR while maintaining context. |
4 | | -""" |
| 1 | +"""DOCX image OCR using the core document conversion pipeline.""" |
5 | 2 |
|
6 | | -import io |
7 | | -import re |
8 | | -import sys |
| 3 | +import hashlib |
| 4 | +import html |
9 | 5 | from typing import Any, BinaryIO, Optional |
| 6 | +from warnings import warn |
10 | 7 |
|
11 | | -from markitdown.converters import HtmlConverter |
12 | | -from markitdown.converter_utils.docx.pre_process import pre_process_docx |
13 | 8 | from markitdown import DocumentConverterResult, StreamInfo |
14 | | -from markitdown._exceptions import ( |
15 | | - MissingDependencyException, |
16 | | - MISSING_DEPENDENCY_MESSAGE, |
17 | | -) |
18 | | -from ._ocr_service import LLMVisionOCRService |
| 9 | +from markitdown.converters import DocxConverter |
19 | 10 |
|
20 | | -# Try loading dependencies |
21 | | -_dependency_exc_info = None |
22 | | -try: |
23 | | - import mammoth |
24 | | - from docx import Document |
25 | | -except ImportError: |
26 | | - _dependency_exc_info = sys.exc_info() |
| 11 | +from ._ocr_service import LLMVisionOCRService, _extract_text_with_metadata |
27 | 12 |
|
28 | | -# Placeholder injected into HTML so that mammoth never sees the OCR markers. |
29 | | -# Must be a single token with no special markdown characters. |
30 | | -_PLACEHOLDER = "MARKITDOWNOCRBLOCK{}" |
31 | 13 |
|
32 | | -_UNDERLINE_STYLE_MAP = "u => u" |
33 | | - |
34 | | - |
35 | | -def _read_embedded_style_map(file_stream: BinaryIO) -> Optional[str]: |
36 | | - """Read the style map embedded in a .docx, if it has one.""" |
37 | | - position = file_stream.tell() |
38 | | - file_stream.seek(0) |
39 | | - try: |
40 | | - return mammoth.read_embedded_style_map(file_stream) |
41 | | - finally: |
42 | | - file_stream.seek(position) |
43 | | - |
44 | | - |
45 | | -class DocxConverterWithOCR(HtmlConverter): |
46 | | - """ |
47 | | - Enhanced DOCX Converter with OCR support for embedded images. |
48 | | - Maintains document flow while extracting text from images inline. |
49 | | - """ |
| 14 | +class DocxConverterWithOCR(DocxConverter): |
| 15 | + """Recognize embedded images while inheriting native DOCX conversion.""" |
50 | 16 |
|
51 | 17 | def __init__(self, ocr_service: Optional[LLMVisionOCRService] = None): |
52 | 18 | super().__init__() |
53 | | - self._html_converter = HtmlConverter() |
| 19 | + if not hasattr(DocxConverter, "_image_to_html"): |
| 20 | + raise RuntimeError( |
| 21 | + "DOCX OCR requires markitdown>=0.1.8b3 for the " |
| 22 | + "DocxConverter._image_to_html hook. " |
| 23 | + "Upgrade with: pip install --upgrade 'markitdown>=0.1.8b3'." |
| 24 | + ) |
54 | 25 | self.ocr_service = ocr_service |
55 | 26 |
|
56 | | - def accepts( |
| 27 | + def convert( |
57 | 28 | self, |
58 | 29 | file_stream: BinaryIO, |
59 | 30 | stream_info: StreamInfo, |
60 | 31 | **kwargs: Any, |
61 | | - ) -> bool: |
62 | | - mimetype = (stream_info.mimetype or "").lower() |
63 | | - extension = (stream_info.extension or "").lower() |
64 | | - |
65 | | - if extension == ".docx": |
66 | | - return True |
67 | | - |
68 | | - if mimetype.startswith( |
69 | | - "application/vnd.openxmlformats-officedocument.wordprocessingml" |
70 | | - ): |
71 | | - return True |
72 | | - |
73 | | - return False |
| 32 | + ) -> DocumentConverterResult: |
| 33 | + # Keep repeated-image recognition local to this document, not the instance. |
| 34 | + kwargs["_docx_ocr_cache"] = {} |
| 35 | + return super().convert(file_stream, stream_info, **kwargs) |
74 | 36 |
|
75 | | - def convert( |
| 37 | + def _image_to_html( |
76 | 38 | self, |
77 | | - file_stream: BinaryIO, |
| 39 | + image_stream: BinaryIO, |
78 | 40 | stream_info: StreamInfo, |
79 | 41 | **kwargs: Any, |
80 | | - ) -> DocumentConverterResult: |
81 | | - if _dependency_exc_info is not None: |
82 | | - raise MissingDependencyException( |
83 | | - MISSING_DEPENDENCY_MESSAGE.format( |
84 | | - converter=type(self).__name__, |
85 | | - extension=".docx", |
86 | | - feature="docx", |
87 | | - ) |
88 | | - ) from _dependency_exc_info[1].with_traceback( |
89 | | - _dependency_exc_info[2] |
90 | | - ) # type: ignore[union-attr] |
91 | | - |
92 | | - # Get OCR service if available (from kwargs or instance) |
93 | | - ocr_service: Optional[LLMVisionOCRService] = ( |
94 | | - kwargs.get("ocr_service") or self.ocr_service |
95 | | - ) |
96 | | - |
97 | | - # Pre-process once, up front, so that every subsequent read sees the |
98 | | - # repaired archive. pre_process_docx() also fixes .docx files whose ZIP |
99 | | - # local file headers disagree with the central directory; reading the |
100 | | - # original stream first would make those raise, or silently yield no |
101 | | - # images and drop the OCR output. |
102 | | - pre_process_stream = pre_process_docx(file_stream) |
103 | | - |
104 | | - # Read the embedded style map and combine with any provided style map |
105 | | - caller_style_map = kwargs.get("style_map") |
106 | | - embedded_style_map = _read_embedded_style_map(pre_process_stream) |
107 | | - |
108 | | - style_map = "\n".join( |
109 | | - part |
110 | | - for part in ( |
111 | | - caller_style_map, |
112 | | - embedded_style_map, |
113 | | - _UNDERLINE_STYLE_MAP, |
114 | | - ) |
115 | | - if part |
116 | | - ) |
117 | | - |
118 | | - if ocr_service: |
119 | | - # 1. Extract and OCR images — returns raw text per image |
120 | | - pre_process_stream.seek(0) |
121 | | - image_ocr_map = self._extract_and_ocr_images( |
122 | | - pre_process_stream, ocr_service |
123 | | - ) |
124 | | - |
125 | | - # 2. Convert DOCX → HTML via mammoth |
126 | | - pre_process_stream.seek(0) |
127 | | - html_result = mammoth.convert_to_html( |
128 | | - pre_process_stream, |
129 | | - style_map=style_map, |
130 | | - include_embedded_style_map=False, |
131 | | - ).value |
132 | | - |
133 | | - # 3. Replace <img> tags with plain placeholder tokens so that |
134 | | - # mammoth's HTML→markdown step never escapes our OCR markers. |
135 | | - html_with_placeholders, ocr_texts = self._inject_placeholders( |
136 | | - html_result, image_ocr_map |
137 | | - ) |
138 | | - |
139 | | - # 4. Convert HTML → markdown |
140 | | - md_result = self._html_converter.convert_string( |
141 | | - html_with_placeholders, **kwargs |
142 | | - ) |
143 | | - md = md_result.markdown |
144 | | - |
145 | | - # 5. Swap placeholders for the actual OCR blocks (post-conversion |
146 | | - # so * and _ are never escaped by the markdown converter). |
147 | | - for i, raw_text in enumerate(ocr_texts): |
148 | | - placeholder = _PLACEHOLDER.format(i) |
149 | | - ocr_block = f"*[Image OCR]\n{raw_text}\n[End OCR]*" |
150 | | - md = md.replace(placeholder, ocr_block) |
151 | | - |
152 | | - return DocumentConverterResult(markdown=md) |
153 | | - else: |
154 | | - # Standard conversion without OCR |
155 | | - pre_process_stream.seek(0) |
156 | | - return self._html_converter.convert_string( |
157 | | - mammoth.convert_to_html( |
158 | | - pre_process_stream, |
159 | | - style_map=style_map, |
160 | | - include_embedded_style_map=False, |
161 | | - ).value, |
162 | | - **kwargs, |
| 42 | + ) -> Optional[str]: |
| 43 | + ocr_service = kwargs.get("ocr_service") or self.ocr_service |
| 44 | + if ocr_service is None: |
| 45 | + return None |
| 46 | + |
| 47 | + cache: dict[bytes, Optional[str]] = kwargs.get("_docx_ocr_cache", {}) |
| 48 | + key = hashlib.sha256(image_stream.read()).digest() |
| 49 | + image_stream.seek(0) |
| 50 | + if key in cache: |
| 51 | + return cache[key] |
| 52 | + |
| 53 | + result = _extract_text_with_metadata(ocr_service, image_stream, stream_info) |
| 54 | + if result.error: |
| 55 | + warn( |
| 56 | + f"DOCX image OCR failed: {result.error}. Keeping the native image.", |
| 57 | + RuntimeWarning, |
| 58 | + stacklevel=2, |
163 | 59 | ) |
164 | | - |
165 | | - def _extract_and_ocr_images( |
166 | | - self, file_stream: BinaryIO, ocr_service: LLMVisionOCRService |
167 | | - ) -> dict[str, str]: |
168 | | - """ |
169 | | - Extract images from DOCX and OCR them. |
170 | | -
|
171 | | - Returns: |
172 | | - Dict mapping image relationship IDs to raw OCR text (no markers). |
173 | | - """ |
174 | | - ocr_map = {} |
175 | | - |
176 | | - try: |
177 | | - file_stream.seek(0) |
178 | | - doc = Document(file_stream) |
179 | | - |
180 | | - for rel in doc.part.rels.values(): |
181 | | - if "image" in rel.target_ref.lower(): |
182 | | - try: |
183 | | - image_bytes = rel.target_part.blob |
184 | | - image_stream = io.BytesIO(image_bytes) |
185 | | - ocr_result = ocr_service.extract_text(image_stream) |
186 | | - |
187 | | - if ocr_result.text.strip(): |
188 | | - # Store raw text only — markers added later |
189 | | - ocr_map[rel.rId] = ocr_result.text.strip() |
190 | | - |
191 | | - except Exception: |
192 | | - continue |
193 | | - |
194 | | - except Exception: |
195 | | - pass |
196 | | - |
197 | | - return ocr_map |
198 | | - |
199 | | - def _inject_placeholders( |
200 | | - self, html: str, ocr_map: dict[str, str] |
201 | | - ) -> tuple[str, list[str]]: |
202 | | - """ |
203 | | - Replace <img> tags with numbered placeholder tokens. |
204 | | -
|
205 | | - Returns: |
206 | | - (html_with_placeholders, ordered list of raw OCR texts) |
207 | | - """ |
208 | | - if not ocr_map: |
209 | | - return html, [] |
210 | | - |
211 | | - ocr_texts = list(ocr_map.values()) |
212 | | - used: list[int] = [] |
213 | | - |
214 | | - def replace_img(match: re.Match) -> str: # type: ignore[type-arg] |
215 | | - for i in range(len(ocr_texts)): |
216 | | - if i not in used: |
217 | | - used.append(i) |
218 | | - return f"<p>{_PLACEHOLDER.format(i)}</p>" |
219 | | - return "" # remove image if all OCR texts already used |
220 | | - |
221 | | - result = re.sub(r"<img[^>]*>", replace_img, html) |
222 | | - |
223 | | - # Any OCR texts that had no matching <img> tag go at the end |
224 | | - for i in range(len(ocr_texts)): |
225 | | - if i not in used: |
226 | | - result += f"<p>{_PLACEHOLDER.format(i)}</p>" |
227 | | - |
228 | | - return result, ocr_texts |
| 60 | + cache[key] = None |
| 61 | + return None |
| 62 | + text = result.text.strip() |
| 63 | + if not text: |
| 64 | + cache[key] = None |
| 65 | + return None |
| 66 | + |
| 67 | + text = text.replace("\r\n", "\n").replace("\r", "\n") |
| 68 | + content = html.escape(text).replace("\n", "<br>") |
| 69 | + fragment = f"<p><em>[Image OCR]<br>{content}<br>[End OCR]</em></p>" |
| 70 | + cache[key] = fragment |
| 71 | + return fragment |
0 commit comments