【langchain/pymupdf】pdfを画像に変換してLLMにわたすサンプルコード

概要

RAGなどでPDFを画像にしてLLMに渡したくなることがよくありますが、実装方法をよく忘れるので、メモしました。

準備

langchainとpymupdfを使います。pymupdfのライセンスはAGPLですので、自作アプリに組み込む場合は注意してください。
仮想環境はuvで構築します。uv以外を使用の方はpipやrunのコマンドを適宜読み替えてください。

uv init pdfimagerag
cd pdfimagerag
uv add langchain langchain-openai pymupdf python-dotenv

.envを作成しOPENAI_API_KEYを記入します。環境変数にすでにOPENAI_API_KEYがある場合は不要です。

サンプルとして以下のPDFを使いました。
https://www.nihon-kankou.or.jp/home/userfiles/files/js04report.pdf

コード

以下のような感じです。

pdfimagetest.py

import base64

import dotenv
import fitz  # pymupdf
from langchain_core.messages import HumanMessage
from langchain_openai import ChatOpenAI

dotenv.load_dotenv()


def convert_pdf_to_base64_images(pdf_path: str) -> list[str]:
    doc = fitz.open(pdf_path)
    base64_images_by_page = []
    for page_num in range(len(doc)):
        page = doc.load_page(page_num)
        pix = page.get_pixmap()
        base64_image = base64.b64encode(pix.tobytes()).decode("utf-8")
        base64_images_by_page.append(base64_image)
    return base64_images_by_page


def ask_gpt(text: str, base64_image: str) -> str:
    model = ChatOpenAI(model="gpt-4o-mini")
    message = HumanMessage(
        content=[
            {"type": "text", "text": text},
            {
                "type": "image_url",
                "image_url": {"url": f"data:image/jpeg;base64,{base64_image}"},
            },
        ],
    )
    response = model.invoke([message])
    return response.content


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description="PDFのパスを取得")
    parser.add_argument("pdf_path", help="PDFファイルのパス")
    parser.add_argument("--page_num", type=int, default=0, help="PDFのページ番号")
    parser.add_argument(
        "--question",
        type=str,
        default="この画像は全体的に何色ですか",
        help="GPTに尋ねる質問",
    )
    args = parser.parse_args()
    pdf_path = args.pdf_path

    base64_images_by_page = convert_pdf_to_base64_images(pdf_path)
    image_data = base64_images_by_page[args.page_num]
    response = ask_gpt(args.question, image_data)
    print(response)

コマンドラインで実行してみます。
最初のページの色味について質問してみます。

uv run pdfimagetest.py js04report.pdf --page_num 0 --question この画像は全体的に何色ですか

この画像は主に赤色のグラデーションで構成されており、背景は明るい赤から薄いピンクにかけての色合いです。

1ページ目は赤っぽい色をしているのでうまく読み取れていそうです。
2ページ目の色味についても質問してみます。

uv run pdfimagetest.py js04report.pdf --page_num 1 --question この画像は全体的に何色ですか

この画像は主に白色で、文字は黒色で印刷されています。

2ページ目は白っぽい色をしているので読み取れていそうです。

解説

pdfをpymupdf（fitz）で読み込んだあと、page.get_pixmapを使うことで、画像マップを取得できます。
pdfの画像変換はpdf2imageも有名ですがpdf2imageではpopplerをOSにインストールしておく必要があるため、今回はpythonだけで完結できるpymupdfを使いました。ただ、pymupdfもライセンスがAGPLで使いにくさがあるため、用途に応じて使い分けるとよいかと思います。
pymupdfで取得したpixmapをopenaiに渡すためにbase64に変換します。

def convert_pdf_to_base64_images(pdf_path: str) -> list[str]:
    doc = fitz.open(pdf_path)
    base64_images_by_page = []
    for page_num in range(len(doc)):
        page = doc.load_page(page_num)
        pix = page.get_pixmap()
        base64_image = base64.b64encode(pix.tobytes()).decode("utf-8")
        base64_images_by_page.append(base64_image)
    return base64_images_by_page

LLMの呼び出しはlangchainを使用しています。
ChatOpenAIで画像を渡すときにはHumanMessageでtype=image_urlを指定すればよいです。

def ask_gpt(text: str, base64_image: str) -> str:
    model = ChatOpenAI(model="gpt-4o-mini")
    message = HumanMessage(
        content=[
            {"type": "text", "text": text},
            {
                "type": "image_url",
                "image_url": {"url": f"data:image/jpeg;base64,{base64_image}"},
            },
        ],
    )
    response = model.invoke([message])
    return response.content