Coverage for haystack/components/converters/image/pdf_to_image.py: 93%

1# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai> 

2# 

3# SPDX-License-Identifier: Apache-2.0 

4 

5from pathlib import Path 

6from typing import Any, Literal 

7 

8from haystack import component, logging 

9from haystack.components.converters.image.image_utils import _convert_pdf_to_images, pillow_import, pypdfium2_import 

10from haystack.components.converters.utils import get_bytestream_from_source, normalize_metadata 

11from haystack.dataclasses import ByteStream 

12from haystack.dataclasses.image_content import ImageContent 

13from haystack.utils import expand_page_range 

14 

15logger = logging.getLogger(__name__) 

16 

17 

18@component 

19class PDFToImageContent: 

20 """ 

21 Converts PDF files to ImageContent objects. 

22 

23 ### Usage example 

24 ```python 

25 from haystack.components.converters.image import PDFToImageContent 

26 

27 converter = PDFToImageContent() 

28 

29 sources = ["file.pdf", "another_file.pdf"] 

30 

31 image_contents = converter.run(sources=sources)["image_contents"] 

32 print(image_contents) 

33 

34 # [ImageContent(base64_image='...', 

35 # mime_type='application/pdf', 

36 # detail=None, 

37 # meta={'file_path': 'file.pdf', 'page_number': 1}), 

38 # ...] 

39 ``` 

40 """ 

41 

42 def __init__( 

43 self, 

44 *, 

45 detail: Literal["auto", "high", "low"] | None = None, 

46 size: tuple[int, int] | None = None, 

47 page_range: list[str | int] | None = None, 

48 ) -> None: 

49 """ 

50 Create the PDFToImageContent component. 

51 

52 :param detail: Optional detail level of the image (only supported by OpenAI). One of "auto", "high", or "low". 

53 This will be passed to the created ImageContent objects. 

54 :param size: If provided, resizes the image to fit within the specified dimensions (width, height) while 

55 maintaining aspect ratio. This reduces file size, memory usage, and processing time, which is beneficial 

56 when working with models that have resolution constraints or when transmitting images to remote services. 

57 :param page_range: List of page numbers and/or page ranges to convert to images. Page numbers start at 1. 

58 If None, all pages in the PDF will be converted. Pages outside the valid range (1 to number of pages) 

59 will be skipped with a warning. For example, page_range=[1, 3] will convert only the first and third 

60 pages of the document. It also accepts printable range strings, e.g.: ['1-3', '5', '8', '10-12'] 

61 will convert pages 1, 2, 3, 5, 8, 10, 11, 12. 

62 """ 

63 self.detail = detail 

64 self.size = size 

65 self.page_range = page_range 

66 pypdfium2_import.check() 

67 pillow_import.check() 

68 

69 @component.output_types(image_contents=list[ImageContent]) 

70 def run( 

71 self, 

72 sources: list[str | Path | ByteStream], 

73 meta: dict[str, Any] | list[dict[str, Any]] | None = None, 

74 *, 

75 detail: Literal["auto", "high", "low"] | None = None, 

76 size: tuple[int, int] | None = None, 

77 page_range: list[str | int] | None = None, 

78 ) -> dict[str, list[ImageContent]]: 

79 """ 

80 Converts files to ImageContent objects. 

81 

82 :param sources: 

83 List of file paths or ByteStream objects to convert. 

84 :param meta: 

85 Optional metadata to attach to the ImageContent objects. 

86 This value can be a list of dictionaries or a single dictionary. 

87 If it's a single dictionary, its content is added to the metadata of all produced ImageContent objects. 

88 If it's a list, its length must match the number of sources as they're zipped together. 

89 For ByteStream objects, their `meta` is added to the output ImageContent objects. 

90 :param detail: 

91 Optional detail level of the image (only supported by OpenAI). One of "auto", "high", or "low". 

92 This will be passed to the created ImageContent objects. 

93 If not provided, the detail level will be the one set in the constructor. 

94 :param size: 

95 If provided, resizes the image to fit within the specified dimensions (width, height) while 

96 maintaining aspect ratio. This reduces file size, memory usage, and processing time, which is beneficial 

97 when working with models that have resolution constraints or when transmitting images to remote services. 

98 If not provided, the size value will be the one set in the constructor. 

99 :param page_range: 

100 List of page numbers and/or page ranges to convert to images. Page numbers start at 1. 

101 If None, all pages in the PDF will be converted. Pages outside the valid range (1 to number of pages) 

102 will be skipped with a warning. For example, page_range=[1, 3] will convert only the first and third 

103 pages of the document. It also accepts printable range strings, e.g.: ['1-3', '5', '8', '10-12'] 

104 will convert pages 1, 2, 3, 5, 8, 10, 11, 12. 

105 If not provided, the page_range value will be the one set in the constructor. 

106 

107 :returns: 

108 A dictionary with the following keys: 

109 - `image_contents`: A list of ImageContent objects. 

110 """ 

111 if not sources: 

112 return {"image_contents": []} 

113 

114 resolved_detail = detail or self.detail 

115 resolved_size = size or self.size 

116 resolved_page_range = page_range or self.page_range 

117 

118 expanded_page_range = expand_page_range(resolved_page_range) if resolved_page_range else None 

119 

120 image_contents = [] 

121 

122 meta_list = normalize_metadata(meta, sources_count=len(sources)) 

123 

124 for source, metadata in zip(sources, meta_list, strict=True): 

125 if isinstance(source, str): 

126 source = Path(source) 

127 

128 try: 

129 bytestream = get_bytestream_from_source(source) 

130 except Exception as e: 

131 logger.warning("Could not read {source}. Skipping it. Error: {error}", source=source, error=e) 

132 continue 

133 try: 

134 page_num_and_base64_images = _convert_pdf_to_images( 

135 bytestream=bytestream, page_range=expanded_page_range, size=resolved_size, return_base64=True 

136 ) 

137 except Exception as e: 

138 logger.warning( 

139 "Could not convert file {source}. Skipping it. Error message: {error}", source=source, error=e 

140 ) 

141 continue 

142 

143 merged_metadata = {**bytestream.meta, **metadata} 

144 

145 for page_number, image in page_num_and_base64_images: 

146 per_page_metadata = {**merged_metadata, "page_number": page_number} 

147 # we already know that image is a string because we set return_base64=True but mypy doesn't know that 

148 assert isinstance(image, str) 

149 image_contents.append( 

150 ImageContent( 

151 base64_image=image, mime_type="image/jpeg", meta=per_page_metadata, detail=resolved_detail 

152 ) 

153 ) 

154 

155 return {"image_contents": image_contents}