Spaces:
Sleeping
Sleeping
File size: 7,624 Bytes
ad8cacf |
1 2 3 4 5 6 7 8 9 10 11 12 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 44 45 46 47 48 49 50 51 52 53 54 55 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 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 |
import os
import base64
from litellm import completion
class ImageCaptioner:
"""
Geminiを使用して画像のキャプションを生成するクラス
"""
def __init__(self, api_key=None, model=None):
"""
初期化関数
Args:
api_key (str, optional): XAI API Key. 未指定の場合は.envのXAI_API_KEYを使用
model (str, optional): 使用するモデル名. 未指定の場合は.envのXAI_MODELまたはデフォルト値を使用
"""
# .envからAPIキーを取得（引数で指定がない場合）
self.api_key = api_key or os.getenv('XAI_API_KEY')
# .envからモデル名を取得（引数で指定がない場合）
# 環境変数にない場合はデフォルト値を使用
self.model = model or os.getenv('XAI_MODEL', 'xai/grok-2-vision-1212')
def set_api_key(self, api_key):
"""APIキーを設定する"""
self.api_key = api_key
def is_configured(self):
"""APIキーが設定されているかを確認する"""
return self.api_key is not None and self.api_key.strip() != ""
def generate_caption(self, image_path, prompt=None):
"""
画像のキャプションを生成する
Args:
image_path (str): 画像ファイルのパス
prompt (str, optional): キャプション生成のためのプロンプト
Returns:
str: 生成されたキャプション
"""
if not self.is_configured():
return "APIキーが設定されていません"
if not os.path.exists(image_path):
return f"画像ファイルが見つかりません: {image_path}"
# デフォルトプロンプト
if prompt is None:
prompt = "この画像に写っている動物を簡潔に説明してください。動物の種類と特徴を含めてください。"
try:
# APIキーを設定
api_key = self.api_key or os.getenv('XAI_API_KEY')
if not api_key:
return "APIキーが設定されていません"
os.environ['XAI_API_KEY'] = api_key
# 画像をBase64エンコード
with open(image_path, "rb") as f:
image_data = f.read()
base64_image = base64.b64encode(image_data).decode('utf-8')
# LitellmでXAI APIを呼び出す
response = completion(
model=self.model,
messages=[
{
"role": "user",
"content": [
{"type": "text", "text": prompt},
{
"type": "image_url",
"image_url": {
"url": f"data:image/jpeg;base64,{base64_image}"
}
}
]
}
]
)
# レスポンスからテキストを抽出
if response and hasattr(response, 'choices') and len(response.choices) > 0:
return response.choices[0].message.content
else:
return "キャプションを生成できませんでした"
except Exception as e:
return f"エラーが発生しました: {str(e)}"
def generate_filename(self, image_path, prefix="animal", prompt=None):
"""
画像のキャプションを生成してファイル名として使用する
Args:
image_path (str): 画像ファイルのパス
prefix (str): ファイル名のプレフィックス
prompt (str, optional): キャプション生成のためのプロンプト
Returns:
tuple: (新しいファイル名, 生成されたキャプション)
"""
# キャプション生成
caption = self.generate_caption(image_path, prompt)
# キャプションからファイル名を生成
# 不正な文字を除去し、長さを制限
invalid_chars = r'<>:"/\|?*'
filename = caption.lower()
# 不正な文字を除去
for char in invalid_chars:
filename = filename.replace(char, '')
# 空白をアンダースコアに置換
filename = filename.replace(' ', '_')
# 長さを制限
if len(filename) > 50:
filename = filename[:50]
# 拡張子を取得
_, ext = os.path.splitext(image_path)
# 最終的なファイル名を生成
final_filename = f"{prefix}_{filename}{ext}"
return final_filename, caption
def batch_generate_filenames(self, image_dir, output_dir=None, prefix="animal", prompt=None):
"""
ディレクトリ内の全画像のキャプションを生成してファイル名を変更する
Args:
image_dir (str): 画像ファイルのディレクトリ
output_dir (str, optional): 出力ディレクトリ
prefix (str): ファイル名のプレフィックス
prompt (str, optional): キャプション生成のためのプロンプト
Returns:
dict: {元のファイルパス: (新しいファイルパス, キャプション)}
"""
if not self.is_configured():
print("APIキーが設定されていません")
return {}
# 出力ディレクトリが指定されていない場合は入力ディレクトリを使用
if output_dir is None:
output_dir = image_dir
# 出力ディレクトリが存在しない場合は作成
if not os.path.exists(output_dir):
os.makedirs(output_dir)
results = {}
# ディレクトリ内の画像を処理
for filename in os.listdir(image_dir):
if not filename.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp')):
continue
input_path = os.path.join(image_dir, filename)
# キャプションとファイル名を生成
new_filename, caption = self.generate_filename(input_path, prefix, prompt)
output_path = os.path.join(output_dir, new_filename)
# ファイルをコピーまたは移動
if output_dir != image_dir:
# 別ディレクトリの場合はコピー
with open(input_path, 'rb') as src_file:
with open(output_path, 'wb') as dst_file:
dst_file.write(src_file.read())
else:
# 同じディレクトリの場合は名前変更
os.rename(input_path, output_path)
results[input_path] = (output_path, caption)
print(f"処理完了: {filename} → {new_filename}")
print(f"キャプション: {caption}")
return results
|