Update app.py

app.py

@@ -176,21 +176,61 @@ async def detect_multiple_dogs(image, conf_threshold=0.25, iou_threshold=0.4):
             dogs.append((cropped_image, confidence, xyxy))
     return dogs
 
+async def process_single_dog(image):
+    top1_prob, topk_breeds, topk_probs_percent = await predict_single_dog(image)
+    if top1_prob < 0.2:
+        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,
+            "image": None
+        }
+        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)
+        initial_state = {
+            "explanation": formatted_description,
+            "buttons": [],
+            "show_back": False,
+            "image": image
+        }
+        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"
+            f"1. **{topk_breeds[0]}** ({topk_probs_percent[0]} confidence)\n"
+            f"2. **{topk_breeds[1]}** ({topk_probs_percent[1]} confidence)\n"
+            f"3. **{topk_breeds[2]}** ({topk_probs_percent[2]} confidence)\n\n"
+            "Click on a button to view more information about the breed."
+        )
+        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,
+            "image": image
+        }
+        return explanation, image, buttons[0], buttons[1], buttons[2], gr.update(visible=True), initial_state
 
 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), None
-
     try:
         if isinstance(image, np.ndarray):
             image = Image.fromarray(image)
-
         dogs = await detect_multiple_dogs(image, conf_threshold=0.25, iou_threshold=0.4)
         
         if len(dogs) <= 1:
             return await process_single_dog(image)
 
-        # 多狗情境
         color_list = ['#FF0000', '#00FF00', '#0000FF', '#FFFF00', '#00FFFF', '#FF00FF', '#800080', '#FFA500']
         explanations = []
         buttons = []
@@ -223,7 +263,8 @@ async def predict(image):
             initial_state = {
                 "explanation": final_explanation,
                 "buttons": buttons,
-                "show_back": True
+                "show_back": True,
+                "image": annotated_image
             }
             return (final_explanation, annotated_image, 
                     buttons[0] if len(buttons) > 0 else gr.update(visible=False),
@@ -235,10 +276,10 @@ async def predict(image):
             initial_state = {
                 "explanation": final_explanation,
                 "buttons": [],
-                "show_back": False
+                "show_back": False,
+                "image": annotated_image
             }
             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:
         error_msg = f"An error occurred: {str(e)}"
         print(error_msg)  # 添加日誌輸出
@@ -397,47 +438,6 @@ async def predict(image):
         
 
 
-async def process_single_dog(image):
-    top1_prob, topk_breeds, topk_probs_percent = await predict_single_dog(image)
-    if top1_prob < 0.2:
-        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)
-        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"
-            f"1. **{topk_breeds[0]}** ({topk_probs_percent[0]} confidence)\n"
-            f"2. **{topk_breeds[1]}** ({topk_probs_percent[1]} confidence)\n"
-            f"3. **{topk_breeds[2]}** ({topk_probs_percent[2]} confidence)\n\n"
-            "Click on a button to view more information about the breed."
-        )
-        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), initial_state
@@ -446,13 +446,25 @@ def show_details(choice, previous_output, initial_state):
         breed = choice.split("More about ")[-1]
         description = get_dog_description(breed)
         formatted_description = format_description(description, breed)
+        initial_state["explanation"] = formatted_description
         return formatted_description, gr.update(visible=True), initial_state
     except Exception as e:
         error_msg = f"An error occurred while showing details: {e}"
         print(error_msg)  # 添加日誌輸出
         return error_msg, gr.update(visible=True), initial_state
 
-# 介面部分
+def go_back(state):
+    buttons = state.get("buttons", [])
+    return (
+        state["explanation"],
+        state.get("image", None),
+        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=state["show_back"]),
+        state
+    )
+
 with gr.Blocks() as iface:
     gr.HTML("<h1 style='text-align: center;'>🐶 Dog Breed Classifier 🔍</h1>")
     gr.HTML("<p style='text-align: center;'>Upload a picture of a dog, and the model will predict its breed, provide detailed information, and include an extra information link!</p>")
@@ -486,12 +498,9 @@ with gr.Blocks() as iface:
         )
 
     back_button.click(
-        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),
-                       gr.update(visible=state["show_back"])),
+        go_back,
         inputs=[initial_state],
-        outputs=[output, btn1, btn2, btn3, back_button]
+        outputs=[output, output_image, btn1, btn2, btn3, back_button, initial_state]
     )
 
     gr.Examples(
@@ -501,5 +510,5 @@ 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()
+if __name__ == "__main__":
+    iface.launch()