Coverage for haystack/components/converters/image/image_utils.py: 96%
128 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-21 13:53 +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 for page_number in resolved_page_range:
166 if page_number < 1 or page_number > num_pages:
167 logger.warning("Page {page_number} is out of range for the PDF file. Skipping it.", page_number=page_number)
168 continue
170 # Get dimensions of the page
171 page = pdf[max(page_number - 1, 0)] # Adjust for 0-based indexing
172 _, _, width, height = page.get_mediabox()
174 target_resolution_dpi = 300.0
176 # From pypdfium2 docs: scale (float) – A factor scaling the number of pixels per PDF canvas unit. This defines
177 # the resolution of the image. To convert a DPI value to a scale factor, multiply it by the size of 1 canvas
178 # unit in inches (usually 1/72in).
179 # https://pypdfium2.readthedocs.io/en/stable/python_api.html#pypdfium2._helpers.page.PdfPage.render
180 target_scale = target_resolution_dpi / 72.0
182 # Calculate potential pixels for target_dpi
183 pixels_for_target_scale = width * height * target_scale**2
185 pil_max_pixels = PILImage.MAX_IMAGE_PIXELS or int(1024 * 1024 * 1024 // 4 // 3)
186 # 90% of PIL's default limit to prevent borderline cases
187 pixel_limit = pil_max_pixels * 0.9
189 scale = target_scale
190 if pixels_for_target_scale > pixel_limit:
191 logger.info(
192 "Large PDF detected ({pixels:.2f} pixels). Resizing the image to fit the pixel limit.",
193 pixels=pixels_for_target_scale,
194 )
195 scale = (pixel_limit / (width * height)) ** 0.5
197 pdf_bitmap = page.render(scale=scale)
199 image: "Image" = pdf_bitmap.to_pil()
200 pdf_bitmap.close()
201 if size is not None:
202 # Set reducing_gap=None to disable multi-step shrink; better quality.
203 # https://pillow.readthedocs.io/en/latest/reference/Image.html#PIL.Image.Image.thumbnail
204 image.thumbnail(size=size, reducing_gap=None)
206 all_pdf_images.append((page_number, image))
208 pdf.close()
210 if return_base64:
211 return [
212 (page_number, _encode_pil_image_to_base64(image, mime_type="image/jpeg"))
213 for page_number, image in all_pdf_images
214 ]
216 return all_pdf_images
219class _ImageSourceInfo(TypedDict):
220 path: Path
221 mime_type: str | None
222 page_number: NotRequired[int] # Only present for PDF documents
225def _extract_image_sources_info(
226 documents: list[Document], file_path_meta_field: str, root_path: str
227) -> list[_ImageSourceInfo]:
228 """
229 Extracts the image source information from the documents.
231 :param documents: List of documents to extract image source information from.
232 :param file_path_meta_field: The metadata field in the Document that contains the file path to the image or PDF.
233 :param root_path: The root directory path where document files are located.
235 :returns:
236 A list of _ImageSourceInfo dictionaries, each containing the path and type of the image.
237 If the image is a PDF, the dictionary also contains the page number.
238 :raises ValueError: If the document is missing the file_path_meta_field key in its metadata, the file path is
239 invalid, the MIME type is not supported, or the page number is missing for a PDF document.
240 """
241 images_source_info: list[_ImageSourceInfo] = []
242 for doc in documents:
243 file_path = doc.meta.get(file_path_meta_field)
244 if file_path is None:
245 raise ValueError(
246 f"Document with ID '{doc.id}' is missing the '{file_path_meta_field}' key in its metadata."
247 f" Please ensure that the documents you are trying to convert have this key set."
248 )
250 resolved_file_path = Path(root_path, file_path)
252 # When root_path is set, ensure the resolved path stays within it to block path-traversal
253 # payloads (e.g. "../../etc/passwd") coming from document metadata. When root_path is unset,
254 # file paths are treated as absolute by design and no containment check is applied; callers that
255 # process untrusted metadata should configure root_path (see component docstrings).
256 if root_path:
257 resolved_file_path = resolved_file_path.resolve()
258 resolved_root = Path(root_path).resolve()
259 if not resolved_file_path.is_relative_to(resolved_root):
260 raise ValueError(
261 f"Document with ID '{doc.id}' has a file path '{file_path}' that escapes the "
262 f"configured root '{root_path}'. Resolved path: '{resolved_file_path}'."
263 )
265 if not resolved_file_path.is_file():
266 raise ValueError(
267 f"Document with ID '{doc.id}' has an invalid file path '{resolved_file_path}'. "
268 f"Please ensure that the documents you are trying to convert have valid file paths."
269 )
271 mime_type = doc.meta.get("mime_type") or mimetypes.guess_type(resolved_file_path)[0]
272 if mime_type not in IMAGE_MIME_TYPES:
273 raise ValueError(
274 f"Document with file path '{resolved_file_path}' has an unsupported MIME type '{mime_type}'. "
275 f"Please ensure that the documents you are trying to convert are of the supported "
276 f"types: {', '.join(IMAGE_MIME_TYPES)}."
277 )
279 image_info: _ImageSourceInfo = {"path": resolved_file_path, "mime_type": mime_type}
281 # If mimetype is PDF we also need the page number to be able to convert the right page
282 if mime_type == "application/pdf":
283 page_number = doc.meta.get("page_number")
284 if page_number is None:
285 raise ValueError(
286 f"Document with ID '{doc.id}' comes from the PDF file '{resolved_file_path}' but is missing "
287 f"the 'page_number' key in its metadata. Please ensure that PDF documents you are trying to "
288 f"convert have this key set."
289 )
290 image_info["page_number"] = page_number
292 images_source_info.append(image_info)
294 return images_source_info
297class _PDFPageInfo(TypedDict):
298 doc_idx: int
299 path: Path
300 page_number: int
303def _batch_convert_pdf_pages_to_images(
304 *, pdf_page_infos: list[_PDFPageInfo], return_base64: bool = False, size: tuple[int, int] | None = None
305) -> dict[int, str] | dict[int, "Image"]:
306 """
307 Converts selected PDF pages to images, returning a mapping from document indices to images (PIL or base64).
309 Pages are grouped by file path to ensure each PDF is opened and processed only once for efficiency.
311 :param pdf_page_infos: List of _PDFPageInfo dictionaries with doc_idx, path, and page_number.
312 :param size: Optional tuple of width and height to resize the images to.
313 :param return_base64: If True, return base64 encoded images instead of PIL images.
315 :returns: Dictionary mapping document indices to images (PIL.Image or base64 string).
316 """
317 if not pdf_page_infos:
318 return {}
320 page_infos_by_pdf_path = defaultdict(list)
321 for page_info in pdf_page_infos:
322 page_infos_by_pdf_path[page_info["path"]].append(page_info)
324 converted_images_by_doc_index = {}
326 for pdf_path, page_infos_for_pdf in page_infos_by_pdf_path.items():
327 page_numbers_to_convert = [info["page_number"] for info in page_infos_for_pdf]
328 bytestream = ByteStream.from_file_path(pdf_path)
330 converted_pages = _convert_pdf_to_images(
331 bytestream=bytestream, return_base64=return_base64, page_range=page_numbers_to_convert, size=size
332 )
334 # Map results back to document indices
335 page_number_to_image = dict(converted_pages)
336 for page_info in page_infos_for_pdf:
337 converted_images_by_doc_index[page_info["doc_idx"]] = page_number_to_image[page_info["page_number"]]
339 # mypy is not able to infer that we match the declared return type
340 return converted_images_by_doc_index # type: ignore[return-value]