bedrock API

import boto3
import json
import base64
import pandas as pd
from pathlib import Path

def encode_image(image_path):
    """画像をbase64エンコードする"""
    with open(image_path, "rb") as image_file:
        return base64.b64encode(image_file.read()).decode('utf-8')

def process_image_with_bedrock(image_path):
    """
    Bedrock Claude 3.5 Sonnetで画像を処理し、選択項目を抽出する
    """
    # Bedrockクライアントの初期化
    bedrock = boto3.client(
        service_name='bedrock-runtime',
        region_name='us-east-1'  # 適切なリージョンに変更してください
    )

    # 画像のエンコード
    base64_image = encode_image(image_path)

    # プロンプトの作成
    prompt = """
    この画像は災害報告書のエクセルファイルです。
    画像内で丸が付けられている全ての選択項目を抽出し、
    カテゴリーと選択項目のペアとしてCSVフォーマットで出力してください。
    出力形式は以下の通りです：
    カテゴリー,選択項目

    回答は必ずCSVフォーマットのみを出力し、余分な説明は含めないでください。
    """

    # リクエストボディの作成
    body = json.dumps({
        "anthropic_version": "bedrock-2023-05-31",
        "max_tokens": 1000,
        "messages": [
            {
                "role": "user",
                "content": [
                    {
                        "type": "image",
                        "source": {
                            "type": "base64",
                            "media_type": "image/png",
                            "data": base64_image
                        }
                    },
                    {
                        "type": "text",
                        "text": prompt
                    }
                ]
            }
        ]
    })

    try:
        # Bedrockモデルの呼び出し
        response = bedrock.invoke_model(
            modelId='anthropic.claude-3-sonnet-20240229-v1:0',
            body=body
        )
        
        # レスポンスの解析
        response_body = json.loads(response['body'].read())
        return response_body['messages'][0]['content'][0]['text']
    
    except Exception as e:
        print(f"エラーが発生しました: {str(e)}")
        return None

def save_to_csv(csv_content, output_path):
    """
    CSVコンテンツをファイルとして保存
    """
    # CSVテキストを行に分割
    lines = csv_content.strip().split('\n')
    
    # ヘッダーと内容を分離
    header = lines[0].split(',')
    data = [line.split(',') for line in lines[1:]]
    
    # DataFrameの作成
    df = pd.DataFrame(data, columns=header)
    
    # CSVファイルとして保存
    df.to_csv(output_path, index=False, encoding='utf-8-sig')

def process_directory(input_dir, output_dir):
    """
    指定されたディレクトリ内の全ての画像を処理
    """
    # 出力ディレクトリの作成
    Path(output_dir).mkdir(parents=True, exist_ok=True)
    
    # 画像ファイルの処理
    for image_path in Path(input_dir).glob('*.png'):
        print(f"処理中: {image_path.name}")
        
        # 画像の処理
        csv_content = process_image_with_bedrock(str(image_path))
        
        if csv_content:
            # 出力ファイル名の生成
            output_path = Path(output_dir) / f"{image_path.stem}.csv"
            
            # CSVとして保存
            save_to_csv(csv_content, output_path)
            print(f"保存完了: {output_path}")
        else:
            print(f"エラー: {image_path.name} の処理に失敗しました")

def main():
    # 入力と出力のディレクトリを設定
    input_dir = "input_images"  # 画像ファイルのあるディレクトリ
    output_dir = "output_csv"   # CSV出力先ディレクトリ
    
    try:
        process_directory(input_dir, output_dir)
        print("全ての処理が完了しました")
    except Exception as e:
        print(f"エラーが発生しました: {str(e)}")

if __name__ == "__main__":
    main()