Coverage for haystack/components/converters/image/image_utils.py: 96%
130 statements
« prev ^ index » next coverage.py v7.16.1, created at 2026-10-09 14:34 +0000
1# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
2#
3# SPDX-License-Identifier: Apache-2.0
5import base64
6import mimetypes
7from collections import defaultdict
8from io import BytesIO
9from pathlib import Path
10from typing import TypedDict, Union
12from typing_extensions import NotRequired
14from haystack import logging
15from haystack.dataclasses import ByteStream, Document
16from haystack.dataclasses.image_content import IMAGE_MIME_TYPES, MIME_TO_FORMAT
17from haystack.lazy_imports import LazyImport
19with LazyImport("Run 'pip install pypdfium2'") as pypdfium2_import:
20 from pypdfium2 import PdfDocument
22with LazyImport("Run 'pip install pillow'") as pillow_import:
23 from PIL import Image as PILImage
24 from PIL.Image import Image
25 from PIL.ImageFile import ImageFile
28logger = logging.getLogger(__name__)
31def _encode_image_to_base64(bytestream: ByteStream, size: tuple[int, int] | None = None) -> tuple[str | None, str]:
32 """
33 Encode an image from a ByteStream into a base64-encoded string.
35 Optionally resize the image before encoding to improve performance for downstream processing.
37 :param bytestream: ByteStream containing the image data.
38 :param size: If provided, resizes the image to fit within the specified dimensions (width, height) while
39 maintaining aspect ratio. This reduces file size, memory usage, and processing time, which is beneficial
40 when working with models that have resolution constraints or when transmitting images to remote services.
42 :returns:
43 A tuple (mime_type, base64_str), where:
44 - mime_type (Optional[str]): The mime type of the encoded image, determined from the original data or image
45 content. Can be None if the mime type cannot be reliably identified.
46 - base64_str (str): The base64-encoded string representation of the (optionally resized) image.
47 """
48 if size is None:
49 if bytestream.mime_type is None:
50 logger.warning(
51 "No mime type provided for the image. "
52 "This may cause compatibility issues with downstream systems requiring a specific mime type. "
53 "Please provide a mime type for the image."
54 )
55 return bytestream.mime_type, base64.b64encode(bytestream.data).decode("utf-8")
57 # Check the import
58 pillow_import.check()
60 # Load the image
61 if bytestream.mime_type and bytestream.mime_type in MIME_TO_FORMAT:
62 formats = [MIME_TO_FORMAT[bytestream.mime_type]]
63 else:
64 formats = None
65 image: "ImageFile" = PILImage.open(BytesIO(bytestream.data), formats=formats)
67 # NOTE: We prefer the format returned by PIL
68 inferred_mime_type = image.get_format_mimetype() or bytestream.mime_type
70 # Downsize the image in place
71 if size is not None:
72 # Set reducing_gap=None to disable multi-step shrink; better quality.
73 # https://pillow.readthedocs.io/en/latest/reference/Image.html#PIL.Image.Image.thumbnail
74 image.thumbnail(size=size, reducing_gap=None)
76 # Convert the image to base64 string
77 if not inferred_mime_type:
78 logger.warning(
79 "Could not determine mime type for image. Defaulting to 'image/jpeg'. "
80 "Consider providing a mime_type parameter."
81 )
82 inferred_mime_type = "image/jpeg"
83 return inferred_mime_type, _encode_pil_image_to_base64(image=image, mime_type=inferred_mime_type)
86def _encode_pil_image_to_base64(image: Union["Image", "ImageFile"], mime_type: str = "image/jpeg") -> str:
87 """
88 Convert a PIL Image object to a base64-encoded string.
90 Automatically converts images with transparency to RGB if saving as JPEG.
92 :param image: A PIL Image or ImageFile object to encode.
93 :param mime_type: The MIME type to use when encoding the image. Defaults to "image/jpeg".
94 :returns:
95 Base64-encoded string representing the image.
96 """
97 # Check the import
98 pillow_import.check()
100 # Convert image to RGB if it has an alpha channel and we are saving as JPEG
101 if (mime_type == "image/jpeg" or mime_type == "image/jpg") and (
102 image.mode in ("RGBA", "LA") or (image.mode == "P" and "transparency" in image.info)
103 ):
104 image = image.convert("RGB")
106 buffered = BytesIO()
107 form = MIME_TO_FORMAT.get(mime_type)
108 if form is None:
109 logger.warning("Could not determine format for mime type {mime_type}. Defaulting to JPEG.", mime_type=mime_type)
110 form = "JPEG"
111 image.save(buffered, format=form)
112 return base64.b64encode(buffered.getvalue()).decode("utf-8")
115def _convert_pdf_to_images(
116 *,
117 bytestream: ByteStream,
118 return_base64: bool = False,
119 page_range: list[int] | None = None,
120 size: tuple[int, int] | None = None,
121) -> list[tuple[int, "Image"]] | list[tuple[int, str]]:
122 """
123 Convert a PDF file into a list of PIL Image objects or base64-encoded images.
125 Checks PDF dimensions and adjusts size constraints based on aspect ratio.
127 :param bytestream: ByteStream object containing the PDF data
128 :param return_base64: If True, return base64-encoded images instead of PIL images.
129 :param page_range: List of page numbers and/or page ranges to convert to images. Page numbers start at 1.
130 If None, all pages in the PDF will be converted. Pages outside the valid range (1 to number of pages)
131 will be skipped with a warning. For example, page_range=[1, 3] will convert only the first and third
132 pages of the document. It also accepts printable range strings, e.g.: ['1-3', '5', '8', '10-12']
133 will convert pages 1, 2, 3, 5, 8, 10, 11, 12.
134 :param size: If provided, resizes the image to fit within the specified dimensions (width, height) while
135 maintaining aspect ratio. This reduces file size, memory usage, and processing time, which is beneficial
136 when working with models that have resolution constraints or when transmitting images to remote services.
138 :returns:
139 A list of tuples, each tuple containing the page number and the PIL Image object or base64-encoded image string.
140 """
142 pypdfium2_import.check()
143 pillow_import.check()
145 try:
146 pdf = PdfDocument(BytesIO(bytestream.data))
147 except Exception as e:
148 logger.warning(
149 "Could not read PDF file {file_path}. Skipping it. Error: {error}",
150 file_path=bytestream.meta.get("file_path"),
151 error=e,
152 )
153 return []
155 num_pages = len(pdf)
156 if num_pages == 0:
157 logger.warning("PDF file is empty: {file_path}", file_path=bytestream.meta.get("file_path"))
158 pdf.close()
159 return []
161 all_pdf_images = []
163 resolved_page_range = page_range or range(1, num_pages + 1)
165 try:
166 for page_number in resolved_page_range:
167 if page_number < 1 or page_number > num_pages:
168 logger.warning(
169 "Page {page_number} is out of range for the PDF file {file_path}. Skipping it.",
170 page_number=page_number,
171 file_path=bytestream.meta.get("file_path"),
172 )
173 continue
175 # Get dimensions of the page
176 page = pdf[max(page_number - 1, 0)] # Adjust for 0-based indexing
177 _, _, width, height = page.get_mediabox()
179 target_resolution_dpi = 300.0
181 # From pypdfium2 docs: scale (float) – A factor scaling the number of pixels per PDF canvas unit. This
182 # defines the resolution of the image. To convert a DPI value to a scale factor, multiply it by the size of
183 # 1 canvas unit in inches (usually 1/72in).
184 # https://pypdfium2.readthedocs.io/en/stable/python_api.html#pypdfium2._helpers.page.PdfPage.render
185 target_scale = target_resolution_dpi / 72.0
187 # Calculate potential pixels for target_dpi
188 pixels_for_target_scale = width * height * target_scale**2
190 pil_max_pixels = PILImage.MAX_IMAGE_PIXELS or int(1024 * 1024 * 1024 // 4 // 3)
191 # 90% of PIL's default limit to prevent borderline cases
192 pixel_limit = pil_max_pixels * 0.9
194 scale = target_scale
195 if pixels_for_target_scale > pixel_limit:
196 logger.info(
197 "Large PDF detected ({pixels:.2f} pixels). Resizing the image to fit the pixel limit.",
198 pixels=pixels_for_target_scale,
199 )
200 scale = (pixel_limit / (width * height)) ** 0.5
202 pdf_bitmap = page.render(scale=scale)
204 image: "Image" = pdf_bitmap.to_pil()
205 pdf_bitmap.close()
206 if size is not None:
207 # Set reducing_gap=None to disable multi-step shrink; better quality.
208 # https://pillow.readthedocs.io/en/latest/reference/Image.html#PIL.Image.Image.thumbnail
209 image.thumbnail(size=size, reducing_gap=None)
211 all_pdf_images.append((page_number, image))
212 finally:
213 pdf.close()
215 if return_base64:
216 return [
217 (page_number, _encode_pil_image_to_base64(image, mime_type="image/jpeg"))
218 for page_number, image in all_pdf_images
219 ]
221 return all_pdf_images
224class _ImageSourceInfo(TypedDict):
225 path: Path
226 mime_type: str | None
227 page_number: NotRequired[int] # Only present for PDF documents
230def _extract_image_sources_info(
231 documents: list[Document], file_path_meta_field: str, root_path: str
232) -> list[_ImageSourceInfo]:
233 """
234 Extracts the image source information from the documents.
236 :param documents: List of documents to extract image source information from.
237 :param file_path_meta_field: The metadata field in the Document that contains the file path to the image or PDF.
238 :param root_path: The root directory path where document files are located.
240 :returns:
241 A list of _ImageSourceInfo dictionaries, each containing the path and type of the image.
242 If the image is a PDF, the dictionary also contains the page number.
243 :raises ValueError: If the document is missing the file_path_meta_field key in its metadata, the file path is
244 invalid, the MIME type is not supported, or the page number is missing for a PDF document.
245 """
246 images_source_info: list[_ImageSourceInfo] = []
247 for doc in documents:
248 file_path = doc.meta.get(file_path_meta_field)
249 if file_path is None:
250 raise ValueError(
251 f"Document with ID '{doc.id}' is missing the '{file_path_meta_field}' key in its metadata."
252 f" Please ensure that the documents you are trying to convert have this key set."
253 )
255 resolved_file_path = Path(root_path, file_path)
257 # When root_path is set, ensure the resolved path stays within it to block path-traversal
258 # payloads (e.g. "../../etc/passwd") coming from document metadata. When root_path is unset,
259 # file paths are treated as absolute by design and no containment check is applied; callers that
260 # process untrusted metadata should configure root_path (see component docstrings).
261 if root_path:
262 resolved_file_path = resolved_file_path.resolve()
263 resolved_root = Path(root_path).resolve()
264 if not resolved_file_path.is_relative_to(resolved_root):
265 raise ValueError(
266 f"Document with ID '{doc.id}' has a file path '{file_path}' that escapes the "
267 f"configured root '{root_path}'. Resolved path: '{resolved_file_path}'."
268 )
270 if not resolved_file_path.is_file():
271 raise ValueError(
272 f"Document with ID '{doc.id}' has an invalid file path '{resolved_file_path}'. "
273 f"Please ensure that the documents you are trying to convert have valid file paths."
274 )
276 mime_type = doc.meta.get("mime_type") or mimetypes.guess_type(resolved_file_path)[0]
277 if mime_type not in IMAGE_MIME_TYPES:
278 raise ValueError(
279 f"Document with file path '{resolved_file_path}' has an unsupported MIME type '{mime_type}'. "
280 f"Please ensure that the documents you are trying to convert are of the supported "
281 f"types: {', '.join(IMAGE_MIME_TYPES)}."
282 )
284 image_info: _ImageSourceInfo = {"path": resolved_file_path, "mime_type": mime_type}
286 # If mimetype is PDF we also need the page number to be able to convert the right page
287 if mime_type == "application/pdf":
288 page_number = doc.meta.get("page_number")
289 if page_number is None:
290 raise ValueError(
291 f"Document with ID '{doc.id}' comes from the PDF file '{resolved_file_path}' but is missing "
292 f"the 'page_number' key in its metadata. Please ensure that PDF documents you are trying to "
293 f"convert have this key set."
294 )
295 image_info["page_number"] = page_number
297 images_source_info.append(image_info)
299 return images_source_info
302class _PDFPageInfo(TypedDict):
303 doc_idx: int
304 path: Path
305 page_number: int
308def _batch_convert_pdf_pages_to_images(
309 *, pdf_page_infos: list[_PDFPageInfo], return_base64: bool = False, size: tuple[int, int] | None = None
310) -> dict[int, str] | dict[int, "Image"]:
311 """
312 Converts selected PDF pages to images, returning a mapping from document indices to images (PIL or base64).
314 Pages are grouped by file path to ensure each PDF is opened and processed only once for efficiency.
316 :param pdf_page_infos: List of _PDFPageInfo dictionaries with doc_idx, path, and page_number.
317 :param size: Optional tuple of width and height to resize the images to.
318 :param return_base64: If True, return base64 encoded images instead of PIL images.
320 :returns: Dictionary mapping document indices to images (PIL.Image or base64 string). Documents whose page could
321 not be converted are not included.
322 """
323 if not pdf_page_infos:
324 return {}
326 page_infos_by_pdf_path = defaultdict(list)
327 for page_info in pdf_page_infos:
328 page_infos_by_pdf_path[page_info["path"]].append(page_info)
330 converted_images_by_doc_index = {}
332 for pdf_path, page_infos_for_pdf in page_infos_by_pdf_path.items():
333 page_numbers_to_convert = [info["page_number"] for info in page_infos_for_pdf]
334 bytestream = ByteStream.from_file_path(pdf_path, meta={"file_path": str(pdf_path)})
336 converted_pages = _convert_pdf_to_images(
337 bytestream=bytestream, return_base64=return_base64, page_range=page_numbers_to_convert, size=size
338 )
340 # Map results back to document indices. Pages that could not be converted (out of range, or the PDF
341 # could not be read) are missing from the results, so their documents are left out.
342 page_number_to_image = dict(converted_pages)
343 for page_info in page_infos_for_pdf:
344 if page_info["page_number"] in page_number_to_image:
345 converted_images_by_doc_index[page_info["doc_idx"]] = page_number_to_image[page_info["page_number"]]
347 # mypy is not able to infer that we match the declared return type
348 return converted_images_by_doc_index # type: ignore[return-value]