Update app.py

app.py

@@ -426,17 +426,15 @@ async def detect_multiple_dogs(image, conf_threshold=0.2, iou_threshold=0.5):
 
 async def predict(image):
     if image is None:
-        return "Please upload an image to start.", None, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
+        return "Please upload an image to start.", None, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), None
 
     try:
         if isinstance(image, np.ndarray):
             image = Image.fromarray(image)
 
-        # 嘗試檢測多隻狗，進一步降低閾值以提高檢測率
-        dogs = await detect_multiple_dogs(image, conf_threshold=0.05)  # 降低閾值以檢測更多狗
+        dogs = await detect_multiple_dogs(image, conf_threshold=0.05)
         
         if len(dogs) <= 1:
-            # 單狗情境
             return await process_single_dog(image)
 
         # 多狗情境
@@ -469,28 +467,51 @@ async def predict(image):
         final_explanation = "\n\n".join(explanations)
         if buttons:
             final_explanation += "\n\nClick on a button to view more information about the breed."
+            initial_state = {
+                "explanation": final_explanation,
+                "buttons": buttons,
+                "show_back": True
+            }
             return (final_explanation, annotated_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))  # 顯示 back 按鈕
+                    gr.update(visible=True),
+                    initial_state)
         else:
-            return final_explanation, annotated_image, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
+            initial_state = {
+                "explanation": final_explanation,
+                "buttons": [],
+                "show_back": False
+            }
+            return final_explanation, annotated_image, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), initial_state
 
     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)
+        error_msg = f"An error occurred: {str(e)}"
+        print(error_msg)  # 添加日誌輸出
+        return error_msg, None, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), None
 
 async def process_single_dog(image):
     top1_prob, topk_breeds, topk_probs_percent = await predict_single_dog(image)
     if top1_prob < 0.2:
-        return "The image is unclear or the breed is not in the dataset. Please upload a clearer image of a dog.", None, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
+        initial_state = {
+            "explanation": "The image is unclear or the breed is not in the dataset. Please upload a clearer image of a dog.",
+            "buttons": [],
+            "show_back": False
+        }
+        return initial_state["explanation"], None, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), initial_state
 
     breed = topk_breeds[0]
     description = get_dog_description(breed)
 
     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=False)
+        initial_state = {
+            "explanation": formatted_description,
+            "buttons": [],
+            "show_back": False
+        }
+        return formatted_description, image, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), initial_state
     else:
         explanation = (
             f"The model couldn't confidently identify the breed. Here are the top 3 possible breeds:\n\n"
@@ -499,22 +520,31 @@ async def process_single_dog(image):
             f"3. **{topk_breeds[2]}** ({topk_probs_percent[2]} confidence)\n\n"
             "Click on a button to view more information about the breed."
         )
-        return (explanation, image, 
-                gr.update(visible=True, value=f"More about {topk_breeds[0]}"),
-                gr.update(visible=True, value=f"More about {topk_breeds[1]}"),
-                gr.update(visible=True, value=f"More about {topk_breeds[2]}"),
-                gr.update(visible=True))  # 顯示 back 按鈕
-
-def show_details(choice, previous_output):
+        buttons = [
+            gr.update(visible=True, value=f"More about {topk_breeds[0]}"),
+            gr.update(visible=True, value=f"More about {topk_breeds[1]}"),
+            gr.update(visible=True, value=f"More about {topk_breeds[2]}")
+        ]
+        initial_state = {
+            "explanation": explanation,
+            "buttons": buttons,
+            "show_back": True
+        }
+        return explanation, image, buttons[0], buttons[1], buttons[2], gr.update(visible=True), initial_state
+
+def show_details(choice, previous_output, initial_state):
     if not choice:
-        return previous_output, gr.update(visible=True)
+        return previous_output, gr.update(visible=True), initial_state
 
     try:
         breed = choice.split("More about ")[-1]
         description = get_dog_description(breed)
-        return format_description(description, breed), gr.update(visible=True)
+        formatted_description = format_description(description, breed)
+        return formatted_description, gr.update(visible=True), initial_state
     except Exception as e:
-        return f"An error occurred while showing details: {e}", gr.update(visible=True)
+        error_msg = f"An error occurred while showing details: {e}"
+        print(error_msg)  # 添加日誌輸出
+        return error_msg, gr.update(visible=True), initial_state
 
 # 介面部分
 with gr.Blocks() as iface:
@@ -545,12 +575,16 @@ with gr.Blocks() as iface:
     for btn in [btn1, btn2, btn3]:
         btn.click(
             show_details,
-            inputs=[btn, output],
-            outputs=[output, back_button]
+            inputs=[btn, output, initial_state],
+            outputs=[output, back_button, initial_state]
         )
 
     back_button.click(
-        lambda x: x,
+        lambda state: (state["explanation"], 
+                       state["buttons"][0] if len(state["buttons"]) > 0 else gr.update(visible=False),
+                       state["buttons"][1] if len(state["buttons"]) > 1 else gr.update(visible=False),
+                       state["buttons"][2] if len(state["buttons"]) > 2 else gr.update(visible=False),
+                       gr.update(visible=state["show_back"])),
         inputs=[initial_state],
         outputs=[output, btn1, btn2, btn3, back_button]
     )
@@ -562,5 +596,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()