Create app.py

app.py

@@ -0,0 +1,179 @@
+# Gradio版 DINOv2セグメンテーション
+
+# 必要なライブラリのインストール
+!pip install torch torchvision
+!pip install git+https://github.com/facebookresearch/dinov2.git
+!pip install scikit-learn matplotlib Pillow numpy gradio
+
+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
+
+# デバイスの設定
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+print(f"使用デバイス: {device}")
+
+# DINOv2モデルのロード
+print("モデルを読み込んでいます...")
+model = torch.hub.load('facebookresearch/dinov2', 'dinov2_vits14')
+model = model.to(device)
+model.eval()
+print("モデルの読み込み完了！")
+
+# 画像の前処理
+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用のセグメンテーション関数
+    
+    Args:
+        image: PIL Image
+        n_segments: セグメント数
+        show_overlay: オーバーレイ表示するかどうか
+    
+    Returns:
+        セグメンテーション結果の画像
+    """
+    if image is None:
+        return None
+    
+    # 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=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 / 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
+
+# Gradioインターフェース
+with gr.Blocks(title="DINOv2 Image Segmentation") as demo:
+    gr.Markdown(
+        """
+        # 🎨 DINOv2 画像セグメンテーション
+        
+        DINOv2を使った画像のセマンティックセグメンテーションツールです。
+        画像をアップロードして、セグメント数を調整してみてください。
+        """
+    )
+    
+    with gr.Row():
+        with gr.Column():
+            input_image = gr.Image(
+                label="入力画像", 
+                type="pil",
+                height=400
+            )
+            
+            n_segments = gr.Slider(
+                minimum=2,
+                maximum=20,
+                value=8,
+                step=1,
+                label="セグメント数",
+                info="画像を分割する領域の数"
+            )
+            
+            show_overlay = gr.Checkbox(
+                label="オリジナル画像とオーバーレイ表示",
+                value=True
+            )
+            
+            segment_btn = gr.Button("セグメンテーション実行", variant="primary")
+        
+        with gr.Column():
+            output_image = gr.Image(
+                label="セグメンテーション結果",
+                height=400
+            )
+    
+    # 例を追加
+    gr.Examples(
+        examples=[
+            ["https://images.unsplash.com/photo-1506905925346-21bda4d32df4", 8, True],
+            ["https://images.unsplash.com/photo-1441974231531-c6227db76b6e", 6, True],
+            ["https://images.unsplash.com/photo-1472214103451-9374bd1c798e", 10, False],
+        ],
+        inputs=[input_image, n_segments, show_overlay],
+        outputs=output_image,
+        fn=segment_image_gradio,
+        cache_examples=False,
+    )
+    
+    gr.Markdown(
+        """
+        ### 使い方
+        1. 画像をアップロードまたはドラッ���&ドロップ
+        2. セグメント数を調整（2〜20）
+        3. オーバーレイ表示のオン/オフを選択
+        4. 「セグメンテーション実行」ボタンをクリック
+        
+        ### ヒント
+        - セグメント数が少ないほど大きな領域に分割されます
+        - オーバーレイ表示で元の画像と結果を比較できます
+        - 処理には数秒かかる場合があります（GPU使用時は高速）
+        """
+    )
+    
+    # ボタンクリック時のイベント
+    segment_btn.click(
+        fn=segment_image_gradio,
+        inputs=[input_image, n_segments, show_overlay],
+        outputs=output_image
+    )
+    
+    # スライダー変更時のリアルタイム更新（オプション）
+    # n_segments.change(
+    #     fn=segment_image_gradio,
+    #     inputs=[input_image, n_segments, show_overlay],
+    #     outputs=output_image
+    # )
+
+# アプリの起動
+demo.launch(share=True, debug=True)