Update app.py

app.py

@@ -176,6 +176,7 @@ 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:
@@ -183,7 +184,8 @@ async def process_single_dog(image):
             "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
+            "image": None,
+            "is_multi_dog": False
         }
         return initial_state["explanation"], None, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), initial_state
 
@@ -196,7 +198,8 @@ async def process_single_dog(image):
             "explanation": formatted_description,
             "buttons": [],
             "show_back": False,
-            "image": image
+            "image": image,
+            "is_multi_dog": False
         }
         return formatted_description, image, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), initial_state
     else:
@@ -216,7 +219,8 @@ async def process_single_dog(image):
             "explanation": explanation,
             "buttons": buttons,
             "show_back": True,
-            "image": image
+            "image": image,
+            "is_multi_dog": False
         }
         return explanation, image, buttons[0], buttons[1], buttons[2], gr.update(visible=True), initial_state
 
@@ -230,20 +234,20 @@ async def predict(image):
         
         if len(dogs) <= 1:
             return await process_single_dog(image)
-
+        
         color_list = ['#FF0000', '#00FF00', '#0000FF', '#FFFF00', '#00FFFF', '#FF00FF', '#800080', '#FFA500']
         explanations = []
         buttons = []
         annotated_image = image.copy()
         draw = ImageDraw.Draw(annotated_image)
         font = ImageFont.load_default()
-
+        
         for i, (cropped_image, _, box) in enumerate(dogs):
             top1_prob, topk_breeds, topk_probs_percent = await predict_single_dog(cropped_image)
             color = color_list[i % len(color_list)]
             draw.rectangle(box, outline=color, width=3)
             draw.text((box[0], box[1]), f"Dog {i+1}", fill=color, font=font)
-
+            
             breed = topk_breeds[0]
             if top1_prob >= 0.5:
                 description = get_dog_description(breed)
@@ -256,7 +260,7 @@ async def predict(image):
                 buttons.extend([gr.update(visible=True, value=f"Dog {i+1}: More about {breed}") for breed in topk_breeds[:3]])
             else:
                 explanations.append(f"Dog {i+1}: The image is unclear or the breed is not in the dataset.")
-
+        
         final_explanation = "\n\n".join(explanations)
         if buttons:
             final_explanation += "\n\nClick on a button to view more information about the breed."
@@ -264,7 +268,9 @@ async def predict(image):
                 "explanation": final_explanation,
                 "buttons": buttons,
                 "show_back": True,
-                "image": annotated_image
+                "image": annotated_image,
+                "is_multi_dog": True,
+                "dogs_info": explanations
             }
             return (final_explanation, annotated_image, 
                     buttons[0] if len(buttons) > 0 else gr.update(visible=False),
@@ -277,7 +283,9 @@ async def predict(image):
                 "explanation": final_explanation,
                 "buttons": [],
                 "show_back": False,
-                "image": annotated_image
+                "image": annotated_image,
+                "is_multi_dog": True,
+                "dogs_info": explanations
             }
             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:
@@ -286,6 +294,7 @@ async def predict(image):
         return error_msg, None, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), None
 
 
+
 # async def detect_multiple_dogs(image, conf_threshold=0.25, iou_threshold=0.4, merge_threshold=0.5):
 #     results = model_yolo(image, conf=conf_threshold, iou=iou_threshold)[0]
 #     dogs = []
@@ -446,24 +455,38 @@ 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
+        initial_state["current_description"] = formatted_description  # 保存當前顯示的描述
+        initial_state["show_back"] = True  # 確保 back 按鈕可見
         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)  # 添加日誌輸出
+        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
-    )
+    if state.get("is_multi_dog", False):
+        # 恢復到多狗情境的初始狀態
+        buttons = state.get("buttons", [])
+        return (
+            state["explanation"],
+            state["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=False),  # 隱藏 back 按鈕
+            state
+        )
+    else:
+        # 單狗情境，不需要特殊處理
+        return (
+            state["explanation"],
+            state["image"],
+            gr.update(visible=False),
+            gr.update(visible=False),
+            gr.update(visible=False),
+            gr.update(visible=False),
+            state
+        )
 
 with gr.Blocks() as iface:
     gr.HTML("<h1 style='text-align: center;'>🐶 Dog Breed Classifier 🔍</h1>")