Tavi's Travelog

import pathlib
import textwrap
import os
import glob

import google.generativeai as genai

from PIL import Image


def main(args: argparse.Namespace):

    # 環境変数に GOOGLE_API_KEY が設定されていれば不要
    # GOOGLE_API_KEY=os.getenv('GOOGLE_API_KEY')
    #genai.configure(api_key="GOOGLE_API_KEY")
    genai.configure(api_key="xxxxxxxxxxxxxxx")


    gemini_pro = genai.GenerativeModel("gemini-1.5-pro")


    # 複数ファイル一気に読み取りのためにinputディレクトリ内にジャンル別で分けたディレクトリを指定
    genre = args.genre


    # プロンプトは試行錯誤
    #prompt = 'これらの画像から、"例文の日本語訳"と"例文"の項目だけをOCRで読み取って、表にして'
    #prompt = 'これらの画像から、"例文の日本語訳"と"例文"の項目だけを抜き出して、カンマ区切りにして'
    prompt = 'これらの画像から、"例文の日本語訳"と"例文"の項目だけを抜き出して、カンマ区切りの表にして'
    #prompt = 'これらの画像から、下記項目の文字をOCRで読み取って、表にして'\
        #         '・例文の日本語訳'\
        #         '・例文'


    # 画像ファイルのPATH取得
    files = sorted(glob.glob(f'./input/{genre}/*.jpeg'))

    # 画像取り出し
    images = []
    for f in files:
      img = Image.open(f)
      images.append(img)

    # 出力ディレクトリがなかったら作成
    make_dir_path = f'./output/{genre}'
    if os.path.isdir(make_dir_path):
      pass
    else:
      os.makedirs(make_dir_path)

    # GeminiにAPIを投げる
    response = gemini_pro.generate_content([prompt, *images])

    # 結果をファイルに追記
    with open(f'./output/{genre}/output.txt', 'a') as of:
      of.write(response.text)

if __name__ == '__main__':
    
    parser = argparse.ArgumentParser()
    parser.add_argument('genre', type=str, help='What genre')
    args = parser.parse_args()

    main(args)

import pathlib
import textwrap
import os
import glob

import time

from PIL import Image

from openai import OpenAI
import base64
import argparse



def encode_image(image_path):
  with open(image_path, 'rb') as image_file:
    return base64.b64encode(image_file.read()).decode('utf-8')



def main(args: argparse.Namespace):

    # 環境変数に OPEN_API_KEY が設定されていれば不要
    # OPENAI_API_KEY=os.getenv('OPENAI_API_KEY')
    #client = OpenAI(api_key=os.environ.get("OPENAI_API_KEY"))
    client = OpenAI(api_key="xxxxxxxxxxxxxxx")
    
    # 複数ファイル一気に読み取りのためにinputディレクトリ内にジャンル別で分けたディレクトリを指定
    genre = args.genre

    # 画像ファイルのPATH取得
    files = sorted(glob.glob(f'./input/{genre}/*.jpeg'))
    
    # 出力ディレクトリがなかったら作成
    make_dir_path = f'./output/{genre}'
    if os.path.isdir(make_dir_path):
      pass
    else:
      os.makedirs(make_dir_path)


    
    for f in files:
        base64_image = encode_image(f)


        user_content = [{
            "type": "text",
            "text": 'この画像から"例文"と"例文の日本語訳"の項目を抜き出して、"例文の日本語訳"のほうだけスペースを取り除いてから、タブ区切りの表にして'
            }]
        user_content.append({
            "type": "image_url",
            "image_url": {
                "url": f"data:image/jpeg;base64,{base64_image}",
                }
            })

        system_block = {
            "role": "system",
            "content": "You are an Optical Character Recognition (OCR) machine. You will extract all the characters from the image file provided by the user, and you will only privide the extracted text in your response. As an OCR machine, You can only respond with the extracted text.The response should be output in CSV format."
            }

        user_block = {
            "role": "user",
            "content": user_content
            }
        

        messages = [
            system_block,
            user_block
        ]
        
        #print(messages)

        # OpenAI APIを呼び出してレスポンスを取得
        res = client.chat.completions.create(
            #model="gpt-4o-2024-08-06",  # 使用するモデル
            model="gpt-4o",  # 使用するモデル
            messages=messages,
            temperature=0.0,
        )

        # API実行状況と使用トークン数をチェックするためprint
        print(f)
        print(res.usage)
        
        # 結果をファイルに追記
        with open(f'./output/{genre}/output.txt', 'a') as of:
          of.write(f)
          of.write("\n")
          of.write(res.choices[0].message.content)
          of.write("\n")
          of.write("\n")

        # 連続送信しないように待つ(けど、あまり意味ないかも
        time.sleep(3)

if __name__ == '__main__':
    
    parser = argparse.ArgumentParser()
    parser.add_argument('genre', type=str, help='What genre')
    args = parser.parse_args()

    main(args)