Coverage for pdf_utils.py: 100%

1import io 

2import base64 

3import fitz 

4 

5from PIL import Image 

6 

7from typing import List 

8from typing import Tuple 

9 

10 

11def parse_pdf( 

12 pdf_path: str 

13) -> Tuple[List[str], List[str]]: 

14 """ 

15 Parse the PDF document to extract text and images. 

16 """ 

17 pdf_text = [] 

18 pdf_images = [] 

19 

20 with fitz.open(pdf_path) as doc: 

21 for page in doc: 

22 text = page.get_text() 

23 images = page.get_images(full=True) 

24 pdf_text.append(text) 

25 for img in images: 

26 xref = img[0] 

27 base_image = doc.extract_image(xref) 

28 image_bytes = base_image["image"] 

29 image = Image.open(io.BytesIO(image_bytes)).convert("RGB") 

30 encoded_image = encode_image(image) 

31 pdf_images.append(encoded_image) 

32 for page in doc: 

33 image = page_to_image(page) 

34 encoded_image = encode_image(image) 

35 pdf_images.append(encoded_image) 

36 

37 return pdf_text, pdf_images 

38 

39 

40def page_to_image( 

41 page: fitz.Page 

42) -> Image.Image: 

43 """ 

44 Render the full page to a pixel map (as RGB image). 

45 """ 

46 ZOOM = 2 # Increase resolution (1 = 72 DPI, 2 = 144 DPI) 

47 mat = fitz.Matrix(ZOOM, ZOOM) 

48 pix = page.get_pixmap(matrix=mat, colorspace=fitz.csRGB) 

49 image = Image.frombytes("RGB", (pix.width, pix.height), pix.samples) 

50 return image 

51 

52 

53def encode_image( 

54 image: Image.Image, 

55 format: str = "JPEG" 

56) -> str: 

57 """ 

58 Encode a PIL Image to a base64 data URL. 

59 """ 

60 buffered = io.BytesIO() 

61 image.save(buffered, format=format) 

62 val = buffered.getvalue() 

63 encoded = base64.b64encode(val).decode("utf-8") 

64 return f"data:image/jpeg;base64,{encoded}"