Update app.py

app.py

@@ -300,7 +300,7 @@ def _predict_single_dog(image):
 #                     dogs.append((cropped_image, confidence, xyxy))
 #     return dogs
 
-async def detect_multiple_dogs(image, conf_threshold=0.3, iou_threshold=0.5):
+async def detect_multiple_dogs(image, conf_threshold=0.2, iou_threshold=0.5):
     results = model_yolo(image, conf=conf_threshold, iou=iou_threshold)[0]
     dogs = []
     for box in results.boxes:
@@ -432,10 +432,10 @@ async def predict(image):
         if isinstance(image, np.ndarray):
             image = Image.fromarray(image)
 
-        # 嘗試檢測多隻狗，使用較高的閾值以避免錯誤檢測
-        dogs = await detect_multiple_dogs(image, conf_threshold=0.3, iou_threshold=0.5)
+        # 嘗試檢測多隻狗，使用較低的閾值以提高檢測率
+        dogs = await detect_multiple_dogs(image, conf_threshold=0.2, iou_threshold=0.5)
         
-        # 如果只檢測到一隻狗，使用單狗處理邏輯
+        # 如果只檢測到一隻或沒有檢測到狗，使用單狗處理邏輯
         if len(dogs) <= 1:
             return await process_single_dog(image)
 
@@ -473,9 +473,9 @@ async def predict(image):
                     buttons[0] if len(buttons) > 0 else gr.update(visible=False),
                     buttons[1] if len(buttons) > 1 else gr.update(visible=False),
                     buttons[2] if len(buttons) > 2 else gr.update(visible=False),
-                    gr.update(visible=True))
+                    gr.update(visible=False))
         else:
-            return final_explanation, annotated_image, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=True)
+            return final_explanation, annotated_image, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
 
     except Exception as e:
         return f"An error occurred: {str(e)}", None, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
@@ -490,7 +490,7 @@ async def process_single_dog(image):
 
     if top1_prob >= 0.5:
         formatted_description = format_description(description, breed)
-        return formatted_description, image, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=True)
+        return formatted_description, image, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
     else:
         explanation = (
             f"The model couldn't confidently identify the breed. Here are the top 3 possible breeds:\n\n"
@@ -534,10 +534,12 @@ with gr.Blocks() as iface:
     
     back_button = gr.Button("Back", visible=False)
     
+    initial_output = gr.State()
+    
     input_image.change(
         predict,
         inputs=input_image,
-        outputs=[output, output_image, btn1, btn2, btn3, back_button]
+        outputs=[output, output_image, btn1, btn2, btn3, back_button, initial_output]
     )
 
     for btn in [btn1, btn2, btn3]:
@@ -548,8 +550,9 @@ with gr.Blocks() as iface:
         )
 
     back_button.click(
-        lambda: (gr.update(value=""), gr.update(visible=True), gr.update(visible=True), gr.update(visible=True), gr.update(visible=False)),
-        outputs=[output, btn1, btn2, btn3, back_button]
+        lambda initial: (initial, gr.update(visible=False)),
+        inputs=[initial_output],
+        outputs=[output, back_button]
     )
 
     gr.Examples(
@@ -559,5 +562,6 @@ with gr.Blocks() as iface:
 
     gr.HTML('For more details on this project and other work, feel free to visit my GitHub <a href="https://github.com/Eric-Chung-0511/Learning-Record/tree/main/Data%20Science%20Projects/Dog_Breed_Classifier">Dog Breed Classifier</a>')
 
+
 if __name__ == "__main__":
     iface.launch()