Skip to content

Pdf Loader

PDFThumbnailReader

Bases: PDFReader

PDF parser with thumbnail for each page.

Source code in libs/kotaemon/kotaemon/loaders/pdf_loader.py 
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
class PDFThumbnailReader(PDFReader):
    """PDF parser with thumbnail for each page."""

    def __init__(self) -> None:
        """
        Initialize PDFReader.
        """
        super().__init__(return_full_document=False)

    def load_data(
        self,
        file: Path,
        extra_info: Optional[Dict] = None,
        fs: Optional[AbstractFileSystem] = None,
    ) -> List[Document]:
        """Parse file."""
        documents = super().load_data(file, extra_info, fs)

        page_numbers_str = []
        filtered_docs = []
        is_int_page_number: dict[str, bool] = {}

        for doc in documents:
            if "page_label" in doc.metadata:
                page_num_str = doc.metadata["page_label"]
                page_numbers_str.append(page_num_str)
                try:
                    _ = int(page_num_str)
                    is_int_page_number[page_num_str] = True
                    filtered_docs.append(doc)
                except ValueError:
                    is_int_page_number[page_num_str] = False
                    continue

        documents = filtered_docs
        page_numbers = list(range(len(page_numbers_str)))

        print("Page numbers:", len(page_numbers))
        page_thumbnails = get_page_thumbnails(file, page_numbers)

        documents.extend(
            [
                Document(
                    text="Page thumbnail",
                    metadata={
                        "image_origin": page_thumbnail,
                        "type": "thumbnail",
                        "page_label": page_number,
                        **(extra_info if extra_info is not None else {}),
                    },
                )
                for (page_thumbnail, page_number) in zip(
                    page_thumbnails, page_numbers_str
                )
                if is_int_page_number[page_number]
            ]
        )

        return documents

load_data 

1
load_data(file, extra_info=None, fs=None)

Parse file.

Source code in libs/kotaemon/kotaemon/loaders/pdf_loader.py 
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
def load_data(
    self,
    file: Path,
    extra_info: Optional[Dict] = None,
    fs: Optional[AbstractFileSystem] = None,
) -> List[Document]:
    """Parse file."""
    documents = super().load_data(file, extra_info, fs)

    page_numbers_str = []
    filtered_docs = []
    is_int_page_number: dict[str, bool] = {}

    for doc in documents:
        if "page_label" in doc.metadata:
            page_num_str = doc.metadata["page_label"]
            page_numbers_str.append(page_num_str)
            try:
                _ = int(page_num_str)
                is_int_page_number[page_num_str] = True
                filtered_docs.append(doc)
            except ValueError:
                is_int_page_number[page_num_str] = False
                continue

    documents = filtered_docs
    page_numbers = list(range(len(page_numbers_str)))

    print("Page numbers:", len(page_numbers))
    page_thumbnails = get_page_thumbnails(file, page_numbers)

    documents.extend(
        [
            Document(
                text="Page thumbnail",
                metadata={
                    "image_origin": page_thumbnail,
                    "type": "thumbnail",
                    "page_label": page_number,
                    **(extra_info if extra_info is not None else {}),
                },
            )
            for (page_thumbnail, page_number) in zip(
                page_thumbnails, page_numbers_str
            )
            if is_int_page_number[page_number]
        ]
    )

    return documents

get_page_thumbnails 

1
get_page_thumbnails(file_path, pages, dpi=80)

Get image thumbnails of the pages in the PDF file.

Parameters:

Name Type Description Default
file_path Path

path to the image file

 required
page_number list[int]

list of page numbers to extract

 required

Returns:

Type Description
List[Image]

list[Image.Image]: list of page thumbnails
Source code in libs/kotaemon/kotaemon/loaders/pdf_loader.py 
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
def get_page_thumbnails(
    file_path: Path, pages: list[int], dpi: int = 80
) -> List[Image.Image]:
    """Get image thumbnails of the pages in the PDF file.

    Args:
        file_path (Path): path to the image file
        page_number (list[int]): list of page numbers to extract

    Returns:
        list[Image.Image]: list of page thumbnails
    """

    img: Image.Image
    suffix = file_path.suffix.lower()
    assert suffix == ".pdf", "This function only supports PDF files."
    try:
        import fitz
    except ImportError:
        raise ImportError("Please install PyMuPDF: 'pip install PyMuPDF'")

    doc = fitz.open(file_path)

    output_imgs = []
    for page_number in pages:
        page = doc.load_page(page_number)
        pm = page.get_pixmap(dpi=dpi)
        img = Image.frombytes("RGB", [pm.width, pm.height], pm.samples)
        output_imgs.append(convert_image_to_base64(img))

    return output_imgs