見出し画像

PDFから図表を自動抽出してファイルとして保存するAIプログラム(ソースコード付き)

こんにちは!
ノーリーです。ChatGPT使ってますか?

本にある図表、挿絵や写真だけを取り出したい、って時ありませんか?
生成AIのClaude先生に手伝ってもらって、さくっとつくってみました。
画像ファイルは中身がわかるようなファイル名になっています。

この記事は、大阪のIT専門学校「清風情報工科学院」の校長・平岡憲人(ノーリー)がお送りします。

PDFから図表を自動抽出するAIプログラム

このプログラムは、PDFファイルから図表や画像を自動的に抽出し、整理するための高度なツールです。最新のAI技術と画像処理技術を組み合わせることで、PDFドキュメントからの情報抽出を効率化します。

主な機能

  1. 図表の自動抽出: PDFファイルから図や表を自動的に検出し、個別の画像ファイルとして保存します。

  2. AIによる自動キャプション生成: 抽出された各画像に対して、Google Gemini AIを使用して自動的にキャプションを生成します。これにより、抽出された画像の内容を簡単に理解できます。

  3. 表のCSV変換: PDFから抽出された表を、編集や分析が容易なCSVファイルに変換します。

  4. ページ番号の自動修正: AIを使用してページ番号を自動的に抽出し、必要に応じて修正を行います。これにより、抽出された情報の正確な位置を把握できます。

  5. ユーザーフレンドリーなインターフェース: 直感的なGUIを使用してPDFファイルを選択でき、複雑な操作を必要としません。

特徴

  • 高度な画像処理: OpenCVを使用して、PDFから図表を正確に検出し抽出します。

  • 柔軟なファイル管理: 抽出された画像やCSVファイルに対して、自動的にユニークで意味のあるファイル名を生成します。

  • エラー処理: ファイル名の文字化けを回避し、ユニーク化処理を行うことで、データの整合性を保ちます。

  • デバッグモード: 開発者向けに、ページ番号抽出とAI修正をテストするためのモードを搭載しています。

このプログラムは、研究者、ビジネスアナリスト、学生など、PDFドキュメントから効率的に情報を抽出する必要がある方々に特に有用です。AIと自動化技術を活用することで、手動での図表抽出作業にかかる時間と労力を大幅に削減します。

ソースコード

必要なPythonライブラリ

このプログラムを実行するには、Pythonに以下のライブラリをインストールする必要があります。pipを使用して以下のコマンドを実行してください:

pip install PyMuPDF
pip install Pillow
pip install google-generativeai
pip install camelot-py[cv]
pip install opencv-python
pip install numpy
pip install tk
  • PyMuPDF: PDF処理のためのライブラリ

  • Pillow: Python Imaging Library (PIL) のフォーク。画像処理に使用

  • google-generativeai: Google の Gemini AI API を使用するためのライブラリ

  • camelot-py[cv]: PDF から表を抽出するためのライブラリ。[cv] オプションは OpenCV サポートを含む

  • opencv-python: OpenCV ライブラリ。画像処理に使用

  • numpy: 数値計算のためのライブラリ

  • tk: Tkinter GUI ツールキット。ファイル選択ダイアログに使用

必要なアプリケーション

camelot-py ライブラリが PDF から表を抽出する際に Ghostscript を使用します。多くのシステムでは別途インストールが必要です。

Windows

  1. Ghostscriptの公式ウェブサイト(https://www.ghostscript.com/releases/gsdnld.html )にアクセスします。

  2. お使いのWindows版(32ビットまたは64ビット)に合わせて、最新版のインストーラーをダウンロードします。

  3. ダウンロードしたインストーラーを実行し、画面の指示に従ってインストールを完了します。

  4. インストール後、Windowsの環境変数のPATHにGhostscriptのbinディレクトリを追加します。

macOS

MacPortsを使用する場合:

sudo port install ghostscript

Homebrewを使用する場合:

brew install ghostscript

Gemini API キー

プログラムは Google の Gemini AI を使用しているため、Gemini API キーが必要です。これは環境変数 GEMINI_API_KEY として設定する必要があります。現在は、制限付きで無料で利用可能です。具体的には、Gemini 1.5 Flashなら毎分15回、Gemini 1.5 Proなら毎分2回まで無料利用できます。このプログラムでは、Gemini 1.5 Flashを利用しています。

これらの準備が整えば、プログラムを実行する環境が整います。

プログラム

Pythonで書かれています。
Windowsにて動作確認しました。

# 必要なライブラリのインポート
import fitz  # PyMuPDF
import io
import os
import tkinter as tk
from tkinter import filedialog
import google.generativeai as genai
from PIL import Image
import time
import re
import camelot
import cv2
import numpy as np
import unicodedata

# デバッグモードの設定
# MODE = "debug"  # "debug" または "normal"
MODE = "normal"  # "debug" または "normal"

# Gemini APIの設定
# 環境変数からAPIキーを取得
api_key = os.environ.get('GEMINI_API_KEY')
if not api_key:
    raise ValueError("環境変数 GEMINI_API_KEY が設定されていません。")

# Geminiモデルの初期化
genai.configure(api_key=api_key)
model = genai.GenerativeModel('gemini-1.5-flash')

# PDFファイル選択用の関数
def select_pdf_file():
    # Tkinterのルートウィンドウを作成し、非表示に
    root = tk.Tk()
    root.withdraw()
    # ファイル選択ダイアログを表示し、選択されたファイルのパスを返す
    file_path = filedialog.askopenfilename(filetypes=[("PDF files", "*.pdf")])
    return file_path

# 画像のキャプション生成関数
def generate_caption(image_bytes):
    try:
        # バイトデータからPIL Imageオブジェクトを作成
        image = Image.open(io.BytesIO(image_bytes))
        # Geminiモデルを使用してキャプションを生成
        response = model.generate_content([
            "画像の内容を簡潔に説明し体言止めでキャプションをつけて。最後に読点をつけないこと。",
            image
        ])
        return sanitize_filename(response.text.strip().rstrip('。'))
    except Exception as e:
        print(f"キャプション生成エラー: {e}")
        return "未分類画像"

# ユニークなファイル名を生成する関数
def get_unique_filename(output_path):
    base, extension = os.path.splitext(output_path)
    counter = 1
    # 同名のファイルが存在する場合、カウンターを増やしてユニークな名前を生成
    while os.path.exists(output_path):
        output_path = f"{base}({counter}){extension}"
        counter += 1
    return output_path

# PDFから画像を抽出する関数
def extract_images_from_pdf(pdf_path, output_folder, page_numbers):
    # PDFファイルを開く
    doc = fitz.open(pdf_path)
    
    # 各ページを処理
    for page_num in range(len(doc)):
        page = doc.load_page(page_num)
        image_list = page.get_images(full=True)
        
        # ページ内の各画像を処理
        for img_index, img in enumerate(image_list):
            print(f"画像情報: {img}")  # デバッグ用
            print(f"img[1] の型: {type(img[1])}")  # デバッグ用
            xref = img[0]
            base_image = doc.extract_image(xref)
            image_bytes = base_image["image"]
            
            # キャプションを生成
            caption = generate_caption(image_bytes)
            time.sleep(4)  # APIの制限を考慮して遅延を入れる
            
            # ファイル名の生成
            safe_caption = sanitize_filename(caption)
            file_name = f"P{page_numbers[page_num]:03d}_{safe_caption[:50]}.png"
            output_path = os.path.join(output_folder, file_name)
            
            # ユニークなファイル名を取得
            unique_output_path = get_unique_filename(output_path)
            
            # 画像を保存
            with open(unique_output_path, "wb") as image_file:
                image_file.write(image_bytes)
            
            print(f"画像を保存しました: {os.path.basename(unique_output_path)}")

    # PDFファイルを閉じる
    doc.close()

# PDFから表を抽出しCSVとして保存する関数
def extract_tables_from_pdf(pdf_path, output_folder, image_file_names):
    # PDFから表を抽出
    tables = camelot.read_pdf(pdf_path, pages='all')
    
    # 各表を処理
    for i, table in enumerate(tables):
        if i < len(image_file_names):
            base_name = os.path.splitext(image_file_names[i])[0]
            csv_file_name = f"{sanitize_filename(base_name)}.csv"
        else:
            csv_file_name = f'table_{i+1}.csv'
        
        output_path = os.path.join(output_folder, csv_file_name)
        output_path = get_unique_filename(output_path)
        
        # CSVとして保存
        table.df.to_csv(output_path, encoding='utf-8-sig', index=False)
        print(f"表をCSVとして保存しました: {os.path.basename(output_path)}")

def extract_page_number(page):
    # ページの下部2.5cmの領域のテキストを抽出
    rect = page.rect
    dpi = 72  # PDFのデフォルトDPI
    bottom_height = 2.5 / 2.54 * dpi  # 2.5cmをポイントに変換
    bottom_area = fitz.Rect(rect.x0, rect.y1 - bottom_height, rect.x1, rect.y1)
    text = page.get_text("text", clip=bottom_area)
    
    # 数字のみを抽出
    numbers = re.findall(r'\d+', text)
    print(f"抽出された数字: {numbers}")
    
    # 最後の数字を返す(ページ番号は通常最後に配置されるため)
    return numbers[-1] if numbers else None

# PDFから表の画像を抽出する関数
def extract_table_images_from_pdf(pdf_path, output_folder, page_numbers):
    doc = fitz.open(pdf_path)
    image_file_names = []
    
    for page_num in range(len(doc)):
        page = doc[page_num]
        
        # ページ番号を抽出
        page_number = page_numbers[page_num]
        
        pix = page.get_pixmap()
        img = np.frombuffer(pix.samples, dtype=np.uint8).reshape(pix.h, pix.w, pix.n)
        img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)

        # 画像の前処理
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        edges = cv2.Canny(gray, 50, 150, apertureSize=3)
        
        # 輪郭の検出
        contours, _ = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

        for i, contour in enumerate(contours):
            x, y, w, h = cv2.boundingRect(contour)
            aspect_ratio = w / h
            area = w * h
            page_area = img.shape[0] * img.shape[1]
            
            # 図表らしい輪郭のみを抽出
            if 0.5 < aspect_ratio < 2 and 0.01 * page_area < area < 0.9 * page_area:
                table_img = img[y:y+h, x:x+w]
                
                # 画像をバイト列に変換
                _, img_bytes = cv2.imencode('.png', table_img)
                img_bytes = img_bytes.tobytes()
                
                # Geminiを使用してキャプションを生成
                caption = generate_caption(img_bytes)
                time.sleep(4)  # APIの制限を考慮して遅延を入れる
                
                # キャプションをサニタイズ
                safe_caption = sanitize_filename(caption)
                if safe_caption == "未分類画像":
                    safe_caption = sanitize_filename(f"table_{i+1}")
                file_name = f"P{page_number}_{safe_caption[:50]}.png"
                output_path = os.path.join(output_folder, file_name)
                
                unique_output_path = get_unique_filename(output_path)
                
                # 画像を保存
                with open(unique_output_path, "wb") as image_file:
                    image_file.write(cv2.imencode('.png', table_img)[1].tobytes())
                print(f"画像を保存しました: {os.path.basename(unique_output_path)}")
                
                image_file_names.append(os.path.basename(unique_output_path))

    doc.close()
    return image_file_names

def sanitize_filename(filename):
    # Unicode正規化
    filename = unicodedata.normalize('NFKC', filename)
    
    # 制御文字を除去
    filename = ''.join(ch for ch in filename if unicodedata.category(ch)[0] != 'C')
    
    # ファイル名に使用できない文字を置換
    invalid_chars = r'[<>:"/\\|?*]'
    filename = re.sub(invalid_chars, '_', filename)
    
    # 連続するアンダースコアを1つに置換
    filename = re.sub('_+', '_', filename)
    
    # 先頭と末尾の空白とアンダースコアを削除
    filename = filename.strip().strip('_')
    
    # ファイル名が空の場合、デフォルト名を使用
    if not filename:
        filename = "unnamed_file"
    
    # ファイル名の長さを制限(例:255文字)
    max_length = 255
    if len(filename) > max_length:
        filename = filename[:max_length]
    
    return filename

# create_output_folder 関数
def create_output_folder(pdf_path):
    # PDFファイル名を取得
    pdf_filename = os.path.basename(pdf_path)
    # 拡張子を除いたファイル名を取得
    pdf_name_without_ext = os.path.splitext(pdf_filename)[0]
    # 冒頭10文字を取得(10文字未満の場合はそのまま)
    folder_prefix = sanitize_filename(pdf_name_without_ext[:10])
    # 基本フォルダ名を作成
    base_folder_name = f"{folder_prefix}_images"
    # 出力フォルダのパスを作成
    output_folder = os.path.join(os.path.dirname(pdf_path), base_folder_name)
    
    # フォルダが既に存在する場合、連番を付けて新しいフォルダ名を生成
    counter = 1
    while os.path.exists(output_folder):
        new_folder_name = f"{base_folder_name}{counter}"
        output_folder = os.path.join(os.path.dirname(pdf_path), new_folder_name)
        counter += 1
    
    # フォルダを作成
    os.makedirs(output_folder)
    return output_folder

def extract_all_page_numbers(pdf_path):
    # PDFファイルを開く
    doc = fitz.open(pdf_path)
    
    # ページ番号を格納するリストを初期化
    page_numbers = []
    
    # 各ページからページ番号を抽出
    for page in doc:
        extracted_number = extract_page_number(page)
        page_numbers.append(extracted_number)
    
    # PDFファイルを閉じる
    doc.close()
    
    # 抽出したページ番号のリストを返す
    return page_numbers

def correct_page_numbers(page_numbers):
    if all(num is None for num in page_numbers):
        return list(range(1, len(page_numbers) + 1))

    # LLMにページ番号列を送信して修正
    prompt = f"""
    以下はPDFから抽出されたページ番号のリストです。連続していない値や異常値を修正し、
    正しいと思われるページ番号を返してください。Noneは抽出できなかったことを示します。

    抽出されたページ番号: {page_numbers}

    修正後のページ番号をカンマ区切りの数値で返してください。説明は不要です。
    例: 1,2,3,4,5,6,7,8,9,10
    """
    
    response = model.generate_content(prompt)
    response_text = response.text.strip()
    
    try:
        # カンマで分割し、各要素を整数に変換
        corrected_numbers = [int(num.strip()) for num in response_text.split(',')]
        
        # リストの長さが元のリストと同じであることを確認
        if len(corrected_numbers) != len(page_numbers):
            raise ValueError("修正後のリストの長さが元のリストと異なります")
        
        return corrected_numbers
    except Exception as e:
        print(f"ページ番号の修正中にエラーが発生しました: {e}")
        print("元のページ番号を使用します")
        return [num if num is not None else i+1 for i, num in enumerate(page_numbers)]

def main():
    pdf_path = select_pdf_file()
    if pdf_path:
        if MODE == "debug":
            # デバッグモード: ページ番号抽出と修正のテスト
            extracted_numbers = extract_all_page_numbers(pdf_path)
            print(f"抽出されたページ番号: {extracted_numbers}")
            
            corrected_numbers = correct_page_numbers(extracted_numbers)
            print(f"修正後のページ番号: {corrected_numbers}")
        else:
            # 通常の処理
            extracted_numbers = extract_all_page_numbers(pdf_path)
            corrected_numbers = correct_page_numbers(extracted_numbers)
            
            output_folder = create_output_folder(pdf_path)
            extract_images_from_pdf(pdf_path, output_folder, corrected_numbers)
            image_file_names = extract_table_images_from_pdf(pdf_path, output_folder, corrected_numbers)
            extract_tables_from_pdf(pdf_path, output_folder, image_file_names)
            print(f"画像と表の抽出が完了しました。保存先: {output_folder}")
    else:
        print("ファイルが選択されていません。")

if __name__ == "__main__":
    main()

生成AIの活用箇所

画像のキャプション生成関数 generate_caption(image_bytes):

画像の内容を表したファイル名を生み出すのに生成AIを使っています。

ページ番号の修正関数 correct_page_numbers(page_numbers):

ページ番号の修正に生成AIを使っています。PDFファイルの場合、ファイル内のページ数とページ内容に書かれているページ数が異なることが多いです。できる限りページ内容から抽出したページ番号を使うために、読み取ったページ番号を生成AIに送り、誤検出したページ番号について、前後の関係から正しいと推測されるページ番号を生成しています。

プログラムの主要関数リストと処理内容

  1. `select_pdf_file()`

    • 処理: GUIダイアログを使用してPDFファイルを選択する

    • 戻り値: 選択されたPDFファイルのパス

  2. `generate_caption(image_bytes)`

    • 処理: Google Gemini AIを使用して画像のキャプションを生成する

    • 引数: 画像のバイトデータ

    • 戻り値: 生成されたキャプション文字列

  3. `get_unique_filename(output_path)`

    • 処理: ファイル名の重複を避けるためにユニークなファイル名を生成する

    • 引数: 元のファイルパス

    • 戻り値: ユニークなファイルパス

  4. `extract_images_from_pdf(pdf_path, output_folder, page_numbers)`

    • 処理: PDFから画像を抽出し、キャプションを付けて保存する

    • 引数: PDFのパス、出力フォルダ、ページ番号リスト

  5. `extract_tables_from_pdf(pdf_path, output_folder, image_file_names)`

    • 処理: PDFから表を抽出し、CSVとして保存する

    • 引数: PDFのパス、出力フォルダ、画像ファイル名リスト

  6. `extract_page_number(page)`

    • 処理: PDFの各ページからページ番号を抽出する

    • 引数: PDFページオブジェクト

    • 戻り値: 抽出されたページ番号(文字列)

  7. `extract_table_images_from_pdf(pdf_path, output_folder, page_numbers)`

    • 処理: PDFから表の画像を抽出し、保存する

    • 引数: PDFのパス、出力フォルダ、ページ番号リスト

    • 戻り値: 保存された画像ファイル名のリスト

  8. `sanitize_filename(filename)`

    • 処理: ファイル名から不正な文字を除去し、安全なファイル名にする

    • 引数: 元のファイル名

    • 戻り値: サニタイズされたファイル名

  9. `create_output_folder(pdf_path)`

    • 処理: 抽出した画像と表を保存するための出力フォルダを作成する

    • 引数: PDFのパス

    • 戻り値: 作成された出力フォルダのパス

  10. `extract_all_page_numbers(pdf_path)`

    • 処理: PDFの全ページからページ番号を抽出する

    • 引数: PDFのパス

    • 戻り値: 抽出されたページ番号のリスト

  11. `correct_page_numbers(page_numbers)`

    • 処理: 抽出されたページ番号リストを AI を使用して修正する

    • 引数: 抽出されたページ番号リスト

    • 戻り値: 修正されたページ番号リスト

  12. `main()`

    • 処理: プログラムのメイン関数。全体の実行フローを制御する

    • 動作:

      1. PDFファイルの選択

      2. ページ番号の抽出と修正

      3. 画像と表の抽出

      4. 結果の保存

これらの関数が連携して動作することで、PDFからの図表抽出、キャプション生成、表のCSV変換などの一連の処理を自動化しています。

まとめ

このプログラムの生成はClaude 3.5 Sonnetを利用しました。デバッグや改良は、生成されたプログラムをCursor上に移し、そこでClaude 3.5 Sonnetと対話的に行いました。

はまった点は、CV2というPythonライブラリでファイル保存すると文字化けが発生することでした。この原因がわからず、半日程度時間が解けました。最後は人間側(つまり、私)の洞察で、本当の原因はCV2だろと指摘し、それを回避するプログラムをClaudeが書いてくれて完成しています。

AI時代になっても、人間によるスーパーバイジングは必要なようです。徐々にその必要性は下がっていくでしょうけれど。

この記事が参加している募集

よろしければサポートお願いします! いただいたサポートはクリエイターとしての活動費に使わせていただきます! (