feat: 식물종 선택 드롭다운 + 지원 작물 안내 배너 + plant_filter 진단 개선

app.py

@@ -321,38 +321,46 @@ TTA_TF = [
     T.Compose([T.Resize(256), T.CenterCrop(224), T.RandomHorizontalFlip(p=1.0), T.ToTensor(), _NORM]),
 ]
 
-def _predict_ensemble(pil_img):
-    """TTA + 앙상블 예측"""
+def _predict_ensemble(pil_img, plant_filter=None):
+    """TTA + 앙상블 예측. plant_filter가 있으면 해당 식물 클래스만 후보로 사용."""
     if not _models:
         return None, 0.0
-    
+
     device = next(_models[0][0].parameters()).device
-    
+
     # TTA 배치 준비
     tensors = [tf(pil_img) for tf in TTA_TF]
     batch = torch.stack(tensors).to(device)
-    
+
     logits_sum = torch.zeros(NUM_CLASSES, device=device)
-    
+
     with torch.no_grad():
         for model, weight, desc in _models:
             outs = model(batch)   # [TTA, 38]
             avg = outs.mean(0)    # [38]
             logits_sum += avg * weight
-    
+
     # 칼리브레이션
     for idx, factor in CALIBRATION.items():
         logits_sum[idx] = logits_sum[idx] * factor
-    
+
+    # 식물종 필터 — 선택한 작물 클래스만 활성화
+    if plant_filter:
+        mask = torch.full((NUM_CLASSES,), float('-inf'), device=device)
+        for idx, label in ID2LABEL.items():
+            if plant_filter.lower() in label.lower():
+                mask[idx] = 0.0
+        logits_sum = logits_sum + mask
+
     probs = torch.softmax(logits_sum, dim=0)
     top3_probs, top3_ids = probs.topk(3)
-    
+
     top3 = [(ID2LABEL[int(i)], float(p)) for i, p in zip(top3_ids, top3_probs)]
     return top3, float(probs.max())
 
 # ── Main diagnosis function ──────────────────────────────────────────────────
 @spaces.GPU(duration=30)
-def diagnose(image, language):
+def diagnose(image, language, plant_type="Auto-detect"):
     t0 = time.time()
     
     if image is None:
@@ -372,9 +380,24 @@ def diagnose(image, language):
     if not isinstance(image, Image.Image):
         image = Image.fromarray(image)
     pil_img = image.convert("RGB")
-    
+
+    # 식물종 필터 매핑
+    PLANT_KEYWORD_MAP = {
+        "Apple": "Apple", "Blueberry": "Blueberry", "Cherry": "Cherry",
+        "Corn": "Corn", "Maize": "Corn", "Grape": "Grape", "Orange": "Orange",
+        "Peach": "Peach", "Pepper": "Pepper", "Potato": "Potato",
+        "Raspberry": "Raspberry", "Soybean": "Soybean", "Squash": "Squash",
+        "Strawberry": "Strawberry", "Tomato": "Tomato",
+    }
+    selected_plant = None
+    if plant_type and "Auto-detect" not in plant_type:
+        for kw, plant in PLANT_KEYWORD_MAP.items():
+            if kw.lower() in plant_type.lower():
+                selected_plant = plant
+                break
+
     # 예측
-    top3, conf = _predict_ensemble(pil_img)
+    top3, conf = _predict_ensemble(pil_img, plant_filter=selected_plant)
     elapsed = time.time() - t0
     
     if top3 is None:
@@ -493,14 +516,50 @@ with gr.Blocks(
     </div>
     """)
     
+    # 지원 작물 목록
+    SUPPORTED_PLANTS = [
+        "🔍 Auto-detect (자동감지)",
+        "🍎 Apple (사과)",
+        "🫐 Blueberry (블루베리)",
+        "🍒 Cherry (체리)",
+        "🌽 Corn / Maize (옥수수)",
+        "🍇 Grape (포도)",
+        "🍊 Orange (오렌지)",
+        "🍑 Peach (복숭아)",
+        "🫑 Bell Pepper (피망/고추)",
+        "🥔 Potato (감자)",
+        "🫐 Raspberry (라즈베리)",
+        "🌱 Soybean (콩)",
+        "🎃 Squash (호박)",
+        "🍓 Strawberry (딸기)",
+        "🍅 Tomato (토마토)",
+    ]
+
+    gr.HTML("""
+    <div style="background:#fff8e1;border-left:4px solid #ffc107;padding:12px 16px;border-radius:8px;margin-bottom:12px;">
+        <b>⚠️ 지원 작물 안내 / Supported Crops</b><br>
+        이 앱은 아래 <b>14종 작물</b>의 병해를 진단합니다. 지원하지 않는 식물은 오진될 수 있습니다.<br>
+        <small>This app diagnoses diseases for <b>14 specific crops</b> only. Other plants may be misclassified.</small><br>
+        <br>
+        🍎 Apple &nbsp; 🫐 Blueberry &nbsp; 🍒 Cherry &nbsp; 🌽 Corn &nbsp; 🍇 Grape &nbsp; 🍊 Orange &nbsp; 🍑 Peach<br>
+        🫑 Pepper &nbsp; 🥔 Potato &nbsp; 🫐 Raspberry &nbsp; 🌱 Soybean &nbsp; 🎃 Squash &nbsp; 🍓 Strawberry &nbsp; 🍅 Tomato
+    </div>
+    """)
+
     with gr.Row():
         with gr.Column(scale=1):
+            plant_select = gr.Dropdown(
+                choices=SUPPORTED_PLANTS,
+                value="🔍 Auto-detect (자동감지)",
+                label="🌿 Select Crop / 작물 선택 (선택 시 정확도 향상)",
+                info="작물을 먼저 선택하면 더 정확하게 진단합니다",
+            )
             img_input = gr.Image(
                 type="pil",
                 label="📸 Plant Photo / 식물 사진",
                 elem_classes="upload-area",
                 sources=["upload", "webcam", "clipboard"],
-                height=320,
+                height=300,
             )
             lang_select = gr.Radio(
                 choices=["English", "한국어 (Korean)"],
@@ -512,12 +571,6 @@ with gr.Blocks(
                 variant="primary",
                 size="lg",
             )
-            
-            gr.Markdown("""
-**Supported crops**: Apple, Blueberry, Cherry, Corn, Grape, Orange, Peach, Pepper, Potato, Raspberry, Soybean, Squash, Strawberry, Tomato
-
-**38 disease/healthy classes** | PlantVillage accuracy: **99.91%**
-            """)
         
         with gr.Column(scale=1):
             label_output = gr.Label(
@@ -532,12 +585,12 @@ with gr.Blocks(
     
     btn.click(
         fn=diagnose,
-        inputs=[img_input, lang_select],
+        inputs=[img_input, lang_select, plant_select],
         outputs=[label_output, detail_output],
     )
     img_input.change(
         fn=diagnose,
-        inputs=[img_input, lang_select],
+        inputs=[img_input, lang_select, plant_select],
         outputs=[label_output, detail_output],
     )