Gemini Vision で壁紙アプリの自動カテゴリ分類を実装した話

import google.generativeai as genai
from PIL import Image
import json
 
genai.configure(api_key="YOUR_GEMINI_API_KEY")
 
# カテゴリ定義（アプリのカテゴリ一覧）
CATEGORIES = [
    "nature", "city", "abstract", "animals", "space",
    "ocean", "mountain", "forest", "flower", "architecture",
    "art", "minimal", "dark", "colorful", "japan",
]
 
def classify_wallpaper(image_path: str) -> dict:
    """
    壁紙画像をGemini Visionで分類する
    戻り値: {"category": "nature", "confidence": "high", "reason": "..."}
    """
    model = genai.GenerativeModel("gemini-2.0-flash")
    image = Image.open(image_path)
 
    prompt = f"""
あなたは壁紙アプリのカテゴリ分類システムです。
この画像を以下のカテゴリから1つだけ選んでください:
{', '.join(CATEGORIES)}
 
以下のJSON形式で回答してください:
{{
  "category": "選んだカテゴリ（英語）",
  "confidence": "high / medium / low",
  "reason": "分類の根拠（日本語で1文）"
}}
 
必ずJSON形式のみで回答し、他の説明は不要です。
"""
 
    response = model.generate_content([prompt, image])
 
    # JSON抽出（モデルが```jsonブロックで囲んで返すことがある）
    text = response.text.strip()
    if "```json" in text:
        text = text.split("```json")[1].split("```")[0].strip()
    elif "```" in text:
        text = text.split("```")[1].split("```")[0].strip()
 
    return json.loads(text)
 
# 使用例
result = classify_wallpaper("wallpaper_001.jpg")
print(result)
# 出力例: {"category": "nature", "confidence": "high", "reason": "緑豊かな森の朝霧を捉えた写真"}

import re
 
def extract_json_safely(text: str) -> dict | None:
    """レスポンステキストからJSONを安全に抽出する"""
    # コードブロックを除去
    text = re.sub(r'```(?:json)?\n?', '', text).strip()
 
    # JSON部分だけを正規表現で抽出
    json_match = re.search(r'\{[^{}]*\}', text, re.DOTALL)
    if not json_match:
        return None
 
    try:
        return json.loads(json_match.group())
    except json.JSONDecodeError:
        return None

CATEGORY_DESCRIPTIONS = {
    "nature": "森・草原・田舎の風景。都市の建物が入っていたら architecture か city を選ぶ",
    "ocean": "海・海岸・波。水族館の水中写真も含む。川や湖は nature",
    "abstract": "現実の物体が認識できない抽象的なパターン・グラデーション・デジタルアート",
    "art": "絵画・イラスト・アニメ風・水彩画。写真でも強く加工されているものも含む",
    "japan": "日本の神社仏閣・桜・浮世絵・着物。都市の場合でも日本的な要素が主題であれば japan",
}
 
def build_prompt_with_descriptions(categories: list, descriptions: dict) -> str:
    cat_lines = []
    for cat in categories:
        desc = descriptions.get(cat, "")
        if desc:
            cat_lines.append(f"- {cat}: {desc}")
        else:
            cat_lines.append(f"- {cat}")
    return "\n".join(cat_lines)

import time
 
def classify_batch(image_paths: list, rpm_limit: int = 30) -> list:
    """
    レート制限を守りながらバッチ処理する
    rpm_limit: 1分あたりの最大リクエスト数
    """
    results = []
    delay = 60 / rpm_limit  # リクエスト間隔（秒）
 
    for i, path in enumerate(image_paths):
        try:
            result = classify_wallpaper(path)
            results.append({"path": path, "status": "ok", **result})
        except Exception as e:
            # 429エラーは少し待ってリトライ
            if "429" in str(e):
                time.sleep(30)
                result = classify_wallpaper(path)
                results.append({"path": path, "status": "retry", **result})
            else:
                results.append({"path": path, "status": "error", "error": str(e)})
 
        if (i + 1) % 10 == 0:
            print(f"進行状況: {i + 1}/{len(image_paths)}")
 
        if i < len(image_paths) - 1:
            time.sleep(delay)
 
    return results

from google import genai
from pydantic import BaseModel
from enum import Enum
 
class Confidence(str, Enum):
    high = "high"
    medium = "medium"
    low = "low"
 
class Classification(BaseModel):
    category: str
    confidence: Confidence
    reason: str
 
client = genai.Client(api_key="YOUR_GEMINI_API_KEY")
 
def classify_structured(image_path: str) -> Classification:
    img = client.files.upload(file=image_path)
    resp = client.models.generate_content(
        model="gemini-2.5-flash",
        contents=["この壁紙を1カテゴリに分類してください。", img],
        config={
            "response_mime_type": "application/json",
            "response_schema": Classification,
        },
    )
    return Classification.model_validate_json(resp.text)