import torch
from PIL import Image
import requests
from transformers import CLIPProcessor, CLIPModel
import numpy as np

def download_image(url):
    """
    URLから画像をダウンロードしPIL Imageとして返す
    """
    return Image.open(requests.get(url, stream=True).raw)

def setup_clip():
    """
    CLIPモデルとプロセッサをセットアップする
    """
    # モデルとプロセッサをダウンロード
    model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
    processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
    return model, processor

def analyze_image(image_path, text_queries, model, processor):
    """
    画像とテキストクエリの類似度を計算する

    Args:
        image_path: 画像ファイルのパスまたはURL
        text_queries: 分析用テキストのリスト
        model: CLIPモデル
        processor: CLIPプロセッサ
    """
    # 画像の読み込み
    if image_path.startswith('http'):
        image = download_image(image_path)
    else:
        image = Image.open(image_path)

    # 画像とテキストの前処理
    inputs = processor(
        text=text_queries,
        images=image,
        return_tensors="pt",
        padding=True
    )

    # モデルによる推論
    outputs = model(**inputs)

    # テキストと画像の類似度スコアを計算
    logits_per_image = outputs.logits_per_image
    probs = logits_per_image.softmax(dim=1).detach().numpy()

    # 結果の表示
    for text, prob in zip(text_queries, probs[0]):
        print(f"「{text}」である確率: {prob:.2%}")

def main():
    # CLIPのセットアップ
    model, processor = setup_clip()

    # 分析する画像ファイルのパス
    image_path = "1.png"

    # 分析用のテキストクエリ
    text_queries = [
        "猫の写真",
        "犬の写真",
        "風景写真",
        "料理の写真"
    ]

    # 画像分析の実行
    analyze_image(image_path, text_queries, model, processor)

if __name__ == "__main__":
    main()