# SPDX-FileCopyrightText: 2022-present deepset GmbH # # SPDX-License-Identifier: Apache-2.0 import base64 import mimetypes from collections import defaultdict from io import BytesIO from pathlib import Path from typing import TypedDict, Union from typing_extensions import NotRequired from haystack import logging from haystack.dataclasses import ByteStream, Document from haystack.dataclasses.image_content import IMAGE_MIME_TYPES, MIME_TO_FORMAT from haystack.lazy_imports import LazyImport with LazyImport("Run 'pip install pypdfium2'") as pypdfium2_import: from pypdfium2 import PdfDocument with LazyImport("Run 'pip install pillow'") as pillow_import: from PIL import Image as PILImage from PIL.Image import Image from PIL.ImageFile import ImageFile logger = logging.getLogger(__name__) def _encode_image_to_base64(bytestream: ByteStream, size: tuple[int, int] | None = None) -> tuple[str | None, str]: """ Encode an image from a ByteStream into a base64-encoded string. Optionally resize the image before encoding to improve performance for downstream processing. :param bytestream: ByteStream containing the image data. :param size: If provided, resizes the image to fit within the specified dimensions (width, height) while maintaining aspect ratio. This reduces file size, memory usage, and processing time, which is beneficial when working with models that have resolution constraints or when transmitting images to remote services. :returns: A tuple (mime_type, base64_str), where: - mime_type (Optional[str]): The mime type of the encoded image, determined from the original data or image content. Can be None if the mime type cannot be reliably identified. - base64_str (str): The base64-encoded string representation of the (optionally resized) image. """ if size is None: if bytestream.mime_type is None: logger.warning( "No mime type provided for the image. " "This may cause compatibility issues with downstream systems requiring a specific mime type. " "Please provide a mime type for the image." ) return bytestream.mime_type, base64.b64encode(bytestream.data).decode("utf-8") # Check the import pillow_import.check() # Load the image if bytestream.mime_type and bytestream.mime_type in MIME_TO_FORMAT: formats = [MIME_TO_FORMAT[bytestream.mime_type]] else: formats = None image: "ImageFile" = PILImage.open(BytesIO(bytestream.data), formats=formats) # NOTE: We prefer the format returned by PIL inferred_mime_type = image.get_format_mimetype() or bytestream.mime_type # Downsize the image in place if size is not None: # Set reducing_gap=None to disable multi-step shrink; better quality. # https://pillow.readthedocs.io/en/latest/reference/Image.html#PIL.Image.Image.thumbnail image.thumbnail(size=size, reducing_gap=None) # Convert the image to base64 string if not inferred_mime_type: logger.warning( "Could not determine mime type for image. Defaulting to 'image/jpeg'. " "Consider providing a mime_type parameter." ) inferred_mime_type = "image/jpeg" return inferred_mime_type, _encode_pil_image_to_base64(image=image, mime_type=inferred_mime_type) def _encode_pil_image_to_base64(image: Union["Image", "ImageFile"], mime_type: str = "image/jpeg") -> str: """ Convert a PIL Image object to a base64-encoded string. Automatically converts images with transparency to RGB if saving as JPEG. :param image: A PIL Image or ImageFile object to encode. :param mime_type: The MIME type to use when encoding the image. Defaults to "image/jpeg". :returns: Base64-encoded string representing the image. """ # Check the import pillow_import.check() # Convert image to RGB if it has an alpha channel and we are saving as JPEG if (mime_type == "image/jpeg" or mime_type == "image/jpg") and ( image.mode in ("RGBA", "LA") or (image.mode == "P" and "transparency" in image.info) ): image = image.convert("RGB") buffered = BytesIO() form = MIME_TO_FORMAT.get(mime_type) if form is None: logger.warning("Could not determine format for mime type {mime_type}. Defaulting to JPEG.", mime_type=mime_type) form = "JPEG" image.save(buffered, format=form) return base64.b64encode(buffered.getvalue()).decode("utf-8") def _convert_pdf_to_images( *, bytestream: ByteStream, return_base64: bool = False, page_range: list[int] | None = None, size: tuple[int, int] | None = None, ) -> list[tuple[int, "Image"]] | list[tuple[int, str]]: """ Convert a PDF file into a list of PIL Image objects or base64-encoded images. Checks PDF dimensions and adjusts size constraints based on aspect ratio. :param bytestream: ByteStream object containing the PDF data :param return_base64: If True, return base64-encoded images instead of PIL images. :param page_range: List of page numbers and/or page ranges to convert to images. Page numbers start at 1. If None, all pages in the PDF will be converted. Pages outside the valid range (1 to number of pages) will be skipped with a warning. For example, page_range=[1, 3] will convert only the first and third pages of the document. It also accepts printable range strings, e.g.: ['1-3', '5', '8', '10-12'] will convert pages 1, 2, 3, 5, 8, 10, 11, 12. :param size: If provided, resizes the image to fit within the specified dimensions (width, height) while maintaining aspect ratio. This reduces file size, memory usage, and processing time, which is beneficial when working with models that have resolution constraints or when transmitting images to remote services. :returns: A list of tuples, each tuple containing the page number and the PIL Image object or base64-encoded image string. """ pypdfium2_import.check() pillow_import.check() try: pdf = PdfDocument(BytesIO(bytestream.data)) except Exception as e: logger.warning( "Could not read PDF file {file_path}. Skipping it. Error: {error}", file_path=bytestream.meta.get("file_path"), error=e, ) return [] num_pages = len(pdf) if num_pages == 0: logger.warning("PDF file is empty: {file_path}", file_path=bytestream.meta.get("file_path")) pdf.close() return [] all_pdf_images = [] resolved_page_range = page_range or range(1, num_pages + 1) for page_number in resolved_page_range: if page_number < 1 or page_number > num_pages: logger.warning("Page {page_number} is out of range for the PDF file. Skipping it.", page_number=page_number) continue # Get dimensions of the page page = pdf[max(page_number - 1, 0)] # Adjust for 0-based indexing _, _, width, height = page.get_mediabox() target_resolution_dpi = 300.0 # From pypdfium2 docs: scale (float) – A factor scaling the number of pixels per PDF canvas unit. This defines # the resolution of the image. To convert a DPI value to a scale factor, multiply it by the size of 1 canvas # unit in inches (usually 1/72in). # https://pypdfium2.readthedocs.io/en/stable/python_api.html#pypdfium2._helpers.page.PdfPage.render target_scale = target_resolution_dpi / 72.0 # Calculate potential pixels for target_dpi pixels_for_target_scale = width * height * target_scale**2 pil_max_pixels = PILImage.MAX_IMAGE_PIXELS or int(1024 * 1024 * 1024 // 4 // 3) # 90% of PIL's default limit to prevent borderline cases pixel_limit = pil_max_pixels * 0.9 scale = target_scale if pixels_for_target_scale > pixel_limit: logger.info( "Large PDF detected ({pixels:.2f} pixels). Resizing the image to fit the pixel limit.", pixels=pixels_for_target_scale, ) scale = (pixel_limit / (width * height)) ** 0.5 pdf_bitmap = page.render(scale=scale) image: "Image" = pdf_bitmap.to_pil() pdf_bitmap.close() if size is not None: # Set reducing_gap=None to disable multi-step shrink; better quality. # https://pillow.readthedocs.io/en/latest/reference/Image.html#PIL.Image.Image.thumbnail image.thumbnail(size=size, reducing_gap=None) all_pdf_images.append((page_number, image)) pdf.close() if return_base64: return [ (page_number, _encode_pil_image_to_base64(image, mime_type="image/jpeg")) for page_number, image in all_pdf_images ] return all_pdf_images class _ImageSourceInfo(TypedDict): path: Path mime_type: str | None page_number: NotRequired[int] # Only present for PDF documents def _extract_image_sources_info( documents: list[Document], file_path_meta_field: str, root_path: str ) -> list[_ImageSourceInfo]: """ Extracts the image source information from the documents. :param documents: List of documents to extract image source information from. :param file_path_meta_field: The metadata field in the Document that contains the file path to the image or PDF. :param root_path: The root directory path where document files are located. :returns: A list of _ImageSourceInfo dictionaries, each containing the path and type of the image. If the image is a PDF, the dictionary also contains the page number. :raises ValueError: If the document is missing the file_path_meta_field key in its metadata, the file path is invalid, the MIME type is not supported, or the page number is missing for a PDF document. """ images_source_info: list[_ImageSourceInfo] = [] for doc in documents: file_path = doc.meta.get(file_path_meta_field) if file_path is None: raise ValueError( f"Document with ID '{doc.id}' is missing the '{file_path_meta_field}' key in its metadata." f" Please ensure that the documents you are trying to convert have this key set." ) resolved_file_path = Path(root_path, file_path) # When root_path is set, ensure the resolved path stays within it to block path-traversal # payloads (e.g. "../../etc/passwd") coming from document metadata. if root_path: resolved_file_path = resolved_file_path.resolve() resolved_root = Path(root_path).resolve() if not resolved_file_path.is_relative_to(resolved_root): raise ValueError( f"Document with ID '{doc.id}' has a file path '{file_path}' that escapes the " f"configured root '{root_path}'. Resolved path: '{resolved_file_path}'." ) if not resolved_file_path.is_file(): raise ValueError( f"Document with ID '{doc.id}' has an invalid file path '{resolved_file_path}'. " f"Please ensure that the documents you are trying to convert have valid file paths." ) mime_type = doc.meta.get("mime_type") or mimetypes.guess_type(resolved_file_path)[0] if mime_type not in IMAGE_MIME_TYPES: raise ValueError( f"Document with file path '{resolved_file_path}' has an unsupported MIME type '{mime_type}'. " f"Please ensure that the documents you are trying to convert are of the supported " f"types: {', '.join(IMAGE_MIME_TYPES)}." ) image_info: _ImageSourceInfo = {"path": resolved_file_path, "mime_type": mime_type} # If mimetype is PDF we also need the page number to be able to convert the right page if mime_type == "application/pdf": page_number = doc.meta.get("page_number") if page_number is None: raise ValueError( f"Document with ID '{doc.id}' comes from the PDF file '{resolved_file_path}' but is missing " f"the 'page_number' key in its metadata. Please ensure that PDF documents you are trying to " f"convert have this key set." ) image_info["page_number"] = page_number images_source_info.append(image_info) return images_source_info class _PDFPageInfo(TypedDict): doc_idx: int path: Path page_number: int def _batch_convert_pdf_pages_to_images( *, pdf_page_infos: list[_PDFPageInfo], return_base64: bool = False, size: tuple[int, int] | None = None ) -> dict[int, str] | dict[int, "Image"]: """ Converts selected PDF pages to images, returning a mapping from document indices to images (PIL or base64). Pages are grouped by file path to ensure each PDF is opened and processed only once for efficiency. :param pdf_page_infos: List of _PDFPageInfo dictionaries with doc_idx, path, and page_number. :param size: Optional tuple of width and height to resize the images to. :param return_base64: If True, return base64 encoded images instead of PIL images. :returns: Dictionary mapping document indices to images (PIL.Image or base64 string). """ if not pdf_page_infos: return {} page_infos_by_pdf_path = defaultdict(list) for page_info in pdf_page_infos: page_infos_by_pdf_path[page_info["path"]].append(page_info) converted_images_by_doc_index = {} for pdf_path, page_infos_for_pdf in page_infos_by_pdf_path.items(): page_numbers_to_convert = [info["page_number"] for info in page_infos_for_pdf] bytestream = ByteStream.from_file_path(pdf_path) converted_pages = _convert_pdf_to_images( bytestream=bytestream, return_base64=return_base64, page_range=page_numbers_to_convert, size=size ) # Map results back to document indices page_number_to_image = dict(converted_pages) for page_info in page_infos_for_pdf: converted_images_by_doc_index[page_info["doc_idx"]] = page_number_to_image[page_info["page_number"]] # mypy is not able to infer that we match the declared return type return converted_images_by_doc_index # type: ignore[return-value]