import torch
import torchvision.transforms as transforms
from PIL import Image
import matplotlib.pyplot as plt
import numpy as np
from sklearn.cluster import KMeans
import gradio as gr
import os

# デバイスの設定
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"使用デバイス: {device}")
print(f"PyTorch version: {torch.__version__}")

# モデルのキャッシュディレクトリを設定
os.environ['TORCH_HOME'] = '/tmp/torch_cache'

# DINOv2モデルのロード
print("モデルを読み込んでいます...")
try:
    model = torch.hub.load('facebookresearch/dinov2', 'dinov2_vits14', trust_repo=True)
    model = model.to(device)
    model.eval()
    print("モデルの読み込み完了！")
except Exception as e:
    print(f"モデル読み込みエラー: {e}")
    raise

# 画像の前処理
transform = transforms.Compose([
    transforms.Resize((518, 518)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], 
                       std=[0.229, 0.224, 0.225])
])

def segment_image_gradio(image, n_segments, show_overlay):
    """
    Gradio用のセグメンテーション関数
    """
    if image is None:
        return None
    
    try:
        # PIL Imageに変換
        if not isinstance(image, Image.Image):
            image = Image.fromarray(image).convert('RGB')
        
        original_size = image.size
        
        # 画像の前処理
        img_tensor = transform(image).unsqueeze(0).to(device)
        
        # 特徴量の抽出
        with torch.no_grad():
            features = model.forward_features(img_tensor)
            patch_features = features['x_norm_patchtokens']
        
        # 特徴量を整形
        patch_h = patch_w = 518 // 14
        patch_features = patch_features.reshape(patch_h, patch_w, -1)
        patch_features = patch_features.cpu().numpy()
        
        # K-meansクラスタリング
        features_flat = patch_features.reshape(-1, patch_features.shape[-1])
        kmeans = KMeans(n_clusters=int(n_segments), random_state=42, n_init=10)
        segments = kmeans.fit_predict(features_flat)
        segments = segments.reshape(patch_h, patch_w)
        
        # 元の画像サイズにリサイズ
        segments_resized = Image.fromarray(segments.astype(np.uint8))
        segments_resized = segments_resized.resize(original_size, Image.NEAREST)
        segments_final = np.array(segments_resized)
        
        # カラーマップを適用
        cmap = plt.cm.get_cmap('tab20')
        segments_colored = cmap(segments_final / int(n_segments))
        segments_colored = (segments_colored[:, :, :3] * 255).astype(np.uint8)
        
        # オーバーレイ表示
        if show_overlay:
            image_array = np.array(image)
            result = (image_array * 0.6 + segments_colored * 0.4).astype(np.uint8)
        else:
            result = segments_colored
        
        return result
    
    except Exception as e:
        print(f"エラーが発生しました: {str(e)}")
        return None

# Gradioインターフェース
css = """
.gradio-container {
    font-family: 'IBM Plex Sans', sans-serif;
}
"""

with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        # 🎨 DINOv2 画像セグメンテーション
        
        DINOv2（Meta AI）を使った画像のセマンティックセグメンテーションツールです。
        """
    )
    
    with gr.Row():
        with gr.Column(scale=1):
            input_image = gr.Image(
                label="📸 入力画像", 
                type="pil",
                height=400
            )
            
            with gr.Accordion("⚙️ 設定", open=True):
                n_segments = gr.Slider(
                    minimum=2,
                    maximum=20,
                    value=8,
                    step=1,
                    label="セグメント数",
                    info="画像を分割する領域の数"
                )
                
                show_overlay = gr.Checkbox(
                    label="オーバーレイ表示",
                    value=True,
                    info="元の画像と結果を重ねて表示"
                )
            
            segment_btn = gr.Button("🚀 セグメンテーション実行", variant="primary", size="lg")
        
        with gr.Column(scale=1):
            output_image = gr.Image(
                label="✨ セグメンテーション結果",
                height=400
            )
    
    gr.Markdown(
        """
        ### 📖 使い方
        1. **画像をアップロード** - ドラッグ&ドロップまたはクリックして選択
        2. **セグメント数を調整** - スライダーで2〜20の間で設定
        3. **オーバーレイ表示** - チェックボックスでオン/オフ
        4. **実行ボタンをクリック** - セグメンテーション開始
        
        ### 💡 ヒント
        - セグメント数が**少ない**ほど大きな領域に分割されます
        - セグメント数が**多い**ほど細かく分割されます
        - オーバーレイ表示で元の画像と結果を比較できます
        - 初回実行時はモデルのダウンロードで時間がかかる場合があります
        
        ### ℹ️ 技術情報
        - モデル: **DINOv2-ViT-S/14** (Meta AI)
        - セグメンテーション手法: **K-means clustering**
        - 特徴量: パッチレベルの視覚的特徴量
        """
    )
    
    # ボタンクリック時のイベント
    segment_btn.click(
        fn=segment_image_gradio,
        inputs=[input_image, n_segments, show_overlay],
        outputs=output_image
    )

# アプリの起動
if __name__ == "__main__":
    demo.queue()
    demo.launch()