first

app.py

@@ -1,6 +1,9 @@
 # app.py
 from collections import Counter
+from functools import lru_cache
 from typing import Tuple, Dict, Any, List
+import os
+import yaml
 
 import gradio as gr
 import numpy as np
@@ -36,6 +39,31 @@ except Exception:
     pass
 
 
+@lru_cache(maxsize=1)
+def load_class_names() -> List[str]:
+    """
+    設定ファイルからクラスリストを読み込む
+    """
+    config_path = "configs/deimv2_floorplan.yaml"
+    try:
+        with open(config_path, 'r', encoding='utf-8') as f:
+            config = yaml.safe_load(f)
+            # Modelセクションからclass_namesを取得
+            if 'Model' in config and 'class_names' in config['Model']:
+                return config['Model']['class_names']
+            else:
+                # フォールバック: デフォルトのクラスリスト
+                return ["kanki", "kanki_shikaku", "kanki_regisuta", "window1", "window2", 
+                        "door1", "door2", "bathtub1", "konro1", "sink1", "toilet1", 
+                        "kasaikeihou1", "kasaikeihou2", "houi1", "houi2", "houi3"]
+    except Exception as e:
+        # エラー時はデフォルトのクラスリストを返す
+        print(f"Warning: Failed to load class names from config: {e}")
+        return ["kanki", "kanki_shikaku", "kanki_regisuta", "window1", "window2", 
+                "door1", "door2", "bathtub1", "konro1", "sink1", "toilet1", 
+                "kasaikeihou1", "kasaikeihou2", "houi1", "houi2", "houi3"]
+
+
 def pil_to_np(img: Image.Image) -> np.ndarray:
     return np.array(img.convert("RGB"))
 
@@ -104,6 +132,7 @@ def summarize_detections(detections: List[Detection]) -> List[List[Any]]:
 def inference_pipeline(
     image: Image.Image,
     score_thresh: float = 0.8,
+    selected_classes: List[str] = None,
 ) -> Tuple[Image.Image, List[List[Any]]]:
     """Gradio から呼ばれるメイン処理"""
     if image is None:
@@ -125,6 +154,15 @@ def inference_pipeline(
         # DEIMv2 推論
         detections = run_inference(img_np, score_thresh=score_thresh)
 
+        # クラスフィルタリング: 選択されたクラスのみを残す
+        if selected_classes is not None and len(selected_classes) > 0:
+            # 選択されたクラスリストに含まれる検出結果のみをフィルタリング
+            filtered_detections = [
+                det for det in detections 
+                if det[4] in selected_classes  # det[4]はlabel_name
+            ]
+            detections = filtered_detections
+
         # 描画
         vis_pil = draw_detections(img_pil.copy(), detections)
 
@@ -141,9 +179,10 @@ def inference_pipeline(
 def gpu_inference(
     image: Image.Image,
     score_thresh: float = 0.8,  # UIのデフォルト値と統一
+    selected_classes: List[str] = None,
 ):
     """Spaces ZeroGPU が検出できるようにデコレータ付きの推論関数を用意"""
-    return inference_pipeline(image, score_thresh)
+    return inference_pipeline(image, score_thresh, selected_classes)
 
 
 # =========================
@@ -160,6 +199,9 @@ with gr.Blocks(title="DEIMv2 Floorplan Symbol Detection") as demo:
 """
     )
 
+    # クラスリストを読み込む
+    class_names = load_class_names()
+    
     with gr.Row():
         # 左: 入力
         with gr.Column(scale=1):
@@ -168,13 +210,23 @@ with gr.Blocks(title="DEIMv2 Floorplan Symbol Detection") as demo:
                 type="pil",
                 image_mode="RGB",
             )
-            score_thresh = gr.Slider(
-                minimum=0.0,
-                maximum=1.0,
-                value=0.8,
-                step=0.05,
-                label="スコア閾値",
-            )
+            
+            # 詳細設定タブ（デフォルトは閉じた状態）
+            with gr.Accordion("詳細設定", open=False):
+                score_thresh = gr.Slider(
+                    minimum=0.0,
+                    maximum=1.0,
+                    value=0.8,
+                    step=0.05,
+                    label="スコア閾値",
+                )
+                selected_classes = gr.CheckboxGroup(
+                    choices=class_names,
+                    value=class_names,  # デフォルトで全クラスを選択
+                    label="検出するクラス",
+                    info="選択したクラスの検出結果のみが表示されます",
+                )
+            
             run_button = gr.Button("検出を実行", variant="primary")
 
         # 中央: 出力画像
@@ -195,7 +247,7 @@ with gr.Blocks(title="DEIMv2 Floorplan Symbol Detection") as demo:
     # ボタンの動作
     run_button.click(
         fn=gpu_inference,
-        inputs=[input_image, score_thresh],
+        inputs=[input_image, score_thresh, selected_classes],
         outputs=[output_image, summary_dataframe],
     )