Update app.py

app.py

@@ -224,207 +224,53 @@ async def process_single_dog(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 = []
-        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)
-                formatted_description = format_description(description, breed)
-                explanations.append(f"Dog {i+1}: {formatted_description}")
-            elif top1_prob >= 0.2:
-                dog_explanation = f"Dog {i+1}: Top 3 possible breeds:\n"
-                dog_explanation += "\n".join([f"{j+1}. **{breed}** ({prob} confidence)" for j, (breed, prob) in enumerate(zip(topk_breeds[:3], topk_probs_percent[:3]))])
-                explanations.append(dog_explanation)
-                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."
-            initial_state = {
-                "explanation": final_explanation,
-                "buttons": buttons,
-                "show_back": True,
-                "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),
-                    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),
-                    initial_state)
-        else:
-            initial_state = {
-                "explanation": final_explanation,
-                "buttons": [],
-                "show_back": False,
-                "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:
-        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 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 = []
-    
-#     image_area = image.width * image.height
-#     min_area_ratio = 0.005  # 最小檢測面積佔整個圖像的比例
-    
-#     for box in results.boxes:
-#         if box.cls == 16:  # COCO 數據集中狗的類別是 16
-#             xyxy = box.xyxy[0].tolist()
-#             area = (xyxy[2] - xyxy[0]) * (xyxy[3] - xyxy[1])
-#             if area / image_area >= min_area_ratio:
-#                 confidence = box.conf.item()
-#                 dogs.append((xyxy, confidence))
-    
-#     if dogs:
-#         boxes = torch.tensor([dog[0] for dog in dogs])
-#         scores = torch.tensor([dog[1] for dog in dogs])
-        
-#         # 應用 NMS
-#         keep = nms(boxes, scores, iou_threshold)
-        
-#         merged_dogs = []
-#         for i in keep:
-#             xyxy = boxes[i].tolist()
-#             confidence = scores[i].item()
-#             merged_dogs.append((xyxy, confidence))
-        
-#         # 後處理：分離過於接近的檢測框
-#         final_dogs = []
-#         while merged_dogs:
-#             base_dog = merged_dogs.pop(0)
-#             to_merge = [base_dog]
-            
-#             i = 0
-#             while i < len(merged_dogs):
-#                 iou = box_iou(torch.tensor([base_dog[0]]), torch.tensor([merged_dogs[i][0]]))[0][0].item()
-#                 if iou > merge_threshold:
-#                     to_merge.append(merged_dogs.pop(i))
-#                 else:
-#                     i += 1
-            
-#             if len(to_merge) == 1:
-#                 final_dogs.append(base_dog)
-#             else:
-#                 # 如果檢測到多個重疊框，嘗試分離它們
-#                 centers = torch.tensor([[((box[0] + box[2]) / 2, (box[1] + box[3]) / 2)] for box, _ in to_merge])
-#                 distances = torch.cdist(centers, centers)
-                
-#                 if torch.any(distances > 0):  # 確保不是完全重疊
-#                     max_distance = distances.max()
-#                     if max_distance > (base_dog[0][2] - base_dog[0][0]) * 0.5:  # 如果最大距離大於框寬度的一半
-#                         final_dogs.extend(to_merge)
-#                     else:
-#                         # 合併為一個框
-#                         merged_box = torch.tensor([box for box, _ in to_merge]).mean(dim=0)
-#                         merged_confidence = max(conf for _, conf in to_merge)
-#                         final_dogs.append((merged_box.tolist(), merged_confidence))
-#                 else:
-#                     # 完全重疊的情況，保留置信度最高的
-#                     best_dog = max(to_merge, key=lambda x: x[1])
-#                     final_dogs.append(best_dog)
-        
-#         # 擴展邊界框並創建剪裁的圖像
-#         expanded_dogs = []
-#         for xyxy, confidence in final_dogs:
-#             expanded_xyxy = [
-#                 max(0, xyxy[0] - 20),
-#                 max(0, xyxy[1] - 20),
-#                 min(image.width, xyxy[2] + 20),
-#                 min(image.height, xyxy[3] + 20)
-#             ]
-#             cropped_image = image.crop(expanded_xyxy)
-#             expanded_dogs.append((cropped_image, confidence, expanded_xyxy))
-        
-#         return expanded_dogs
-    
-#     # 如果沒有檢測到狗狗，返回整張圖片
-#     return [(image, 1.0, [0, 0, image.width, image.height])]
-
 # 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)
+#         dogs = await detect_multiple_dogs(image, conf_threshold=0.25, iou_threshold=0.4)
         
-#         # 如果沒有檢測到狗狗或只檢測到一隻，使用整張圖像進行分類
 #         if len(dogs) <= 1:
-#             top1_prob, topk_breeds, topk_probs_percent = await predict_single_dog(image)
-#             if top1_prob >= 0.5:
-#                 return await process_single_dog(image)
-#             else:
-#                 dogs = [(image, 1.0, [0, 0, image.width, image.height])]
-
-#         # 多狗情境處理保持不變
+#             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)
 #                 formatted_description = format_description(description, breed)
 #                 explanations.append(f"Dog {i+1}: {formatted_description}")
-#             else:
+#             elif top1_prob >= 0.2:
 #                 dog_explanation = f"Dog {i+1}: Top 3 possible breeds:\n"
 #                 dog_explanation += "\n".join([f"{j+1}. **{breed}** ({prob} confidence)" for j, (breed, prob) in enumerate(zip(topk_breeds[:3], topk_probs_percent[:3]))])
 #                 explanations.append(dog_explanation)
 #                 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."
 #             initial_state = {
 #                 "explanation": final_explanation,
 #                 "buttons": buttons,
-#                 "show_back": True
+#                 "show_back": True,
+#                 "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),
@@ -436,27 +282,172 @@ async def predict(image):
 #             initial_state = {
 #                 "explanation": final_explanation,
 #                 "buttons": [],
-#                 "show_back": False
+#                 "show_back": False,
+#                 "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:
 #         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
+
+
+# def show_details(choice, previous_output, initial_state):
+#     if not choice:
+#         return previous_output, gr.update(visible=True), initial_state
+
+#     try:
+#         breed = choice.split("More about ")[-1]
+#         description = get_dog_description(breed)
+#         formatted_description = format_description(description, breed)
+        
+#         # 保存當前描述和原始按鈕狀態
+#         initial_state["current_description"] = formatted_description
+#         initial_state["original_buttons"] = initial_state.get("buttons", [])
         
+#         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["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
+#     )
+
+# 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>")
+    
+#     with gr.Row():
+#         input_image = gr.Image(label="Upload a dog image", type="pil")
+#         output_image = gr.Image(label="Annotated Image")
+    
+#     output = gr.Markdown(label="Prediction Results")
+    
+#     with gr.Row():
+#         btn1 = gr.Button("View More 1", visible=False)
+#         btn2 = gr.Button("View More 2", visible=False)
+#         btn3 = gr.Button("View More 3", visible=False)
+    
+#     back_button = gr.Button("Back", visible=False)
+    
+#     initial_state = gr.State()
+    
+#     input_image.change(
+#         predict,
+#         inputs=input_image,
+#         outputs=[output, output_image, btn1, btn2, btn3, back_button, initial_state]
+#     )
+
+#     for btn in [btn1, btn2, btn3]:
+#         btn.click(
+#             show_details,
+#             inputs=[btn, output, initial_state],
+#             outputs=[output, back_button, initial_state]
+#         )
+
+#     back_button.click(
+#         go_back,
+#         inputs=[initial_state],
+#         outputs=[output, output_image, btn1, btn2, btn3, back_button, initial_state]
+#     )
+
+#     gr.Examples(
+#         examples=['Border_Collie.jpg', 'Golden_Retriever.jpeg', 'Saint_Bernard.jpeg', 'French_Bulldog.jpeg', 'Samoyed.jpg'],
+#         inputs=input_image
+#     )
+
+#     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()
+
+
+async def predict(image):
+    if image is None:
+        return "Please upload an image to start.", None, [], gr.update(visible=False), None
+
+    try:
+        if isinstance(image, np.ndarray):
+            image = Image.fromarray(image)
+
+        dogs = await detect_multiple_dogs(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)
 
+            if top1_prob >= 0.5:
+                breed = topk_breeds[0]
+                description = get_dog_description(breed)
+                formatted_description = format_description(description, breed)
+                explanations.append(f"Dog {i+1}: {formatted_description}")
+            elif top1_prob >= 0.2:
+                dog_explanation = f"Dog {i+1}: Top 3 possible breeds:\n"
+                dog_explanation += "\n".join([f"{j+1}. **{breed}** ({prob} confidence)" for j, (breed, prob) in enumerate(zip(topk_breeds[:3], topk_probs_percent[:3]))])
+                explanations.append(dog_explanation)
+                buttons.extend([gr.Button(f"Dog {i+1}: More about {breed}", visible=True) 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."
+            initial_state = {
+                "explanation": final_explanation,
+                "buttons": buttons,
+                "show_back": True,
+                "image": annotated_image,
+                "is_multi_dog": len(dogs) > 1,
+                "dogs_info": explanations
+            }
+            return final_explanation, annotated_image, buttons, gr.update(visible=True), initial_state
+        else:
+            initial_state = {
+                "explanation": final_explanation,
+                "buttons": [],
+                "show_back": False,
+                "image": annotated_image,
+                "is_multi_dog": len(dogs) > 1,
+                "dogs_info": explanations
+            }
+            return final_explanation, annotated_image, [], gr.update(visible=False), initial_state
+
+    except Exception as e:
+        error_msg = f"An error occurred: {str(e)}"
+        print(error_msg)  # Add log output
+        return error_msg, None, [], gr.update(visible=False), None
 
 def show_details(choice, previous_output, initial_state):
     if not choice:
         return previous_output, gr.update(visible=True), initial_state
 
     try:
-        breed = choice.split("More about ")[-1]
+        dog_num, breed = choice.split(": More about ")
         description = get_dog_description(breed)
         formatted_description = format_description(description, breed)
         
-        # 保存當前描述和原始按鈕狀態
+        # Save current description and original button state
         initial_state["current_description"] = formatted_description
         initial_state["original_buttons"] = initial_state.get("buttons", [])
         
@@ -471,10 +462,8 @@ def go_back(state):
     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 按鈕
+        buttons,
+        gr.update(visible=False),  # Hide back button
         state
     )
 
@@ -488,10 +477,9 @@ with gr.Blocks() as iface:
     
     output = gr.Markdown(label="Prediction Results")
     
-    with gr.Row():
-        btn1 = gr.Button("View More 1", visible=False)
-        btn2 = gr.Button("View More 2", visible=False)
-        btn3 = gr.Button("View More 3", visible=False)
+    button_group = gr.Group()
+    with button_group:
+        buttons = []
     
     back_button = gr.Button("Back", visible=False)
     
@@ -500,20 +488,25 @@ with gr.Blocks() as iface:
     input_image.change(
         predict,
         inputs=input_image,
-        outputs=[output, output_image, btn1, btn2, btn3, back_button, initial_state]
+        outputs=[output, output_image, button_group, back_button, initial_state]
     )
 
-    for btn in [btn1, btn2, btn3]:
-        btn.click(
-            show_details,
-            inputs=[btn, output, initial_state],
-            outputs=[output, back_button, initial_state]
-        )
+    def update_buttons(buttons):
+        button_group.clear()
+        for btn in buttons:
+            button = gr.Button(btn.value)
+            button.click(
+                show_details,
+                inputs=[button, output, initial_state],
+                outputs=[output, back_button, initial_state]
+            )
+            buttons.append(button)
+        return button_group
 
     back_button.click(
         go_back,
         inputs=[initial_state],
-        outputs=[output, output_image, btn1, btn2, btn3, back_button, initial_state]
+        outputs=[output, output_image, button_group, back_button, initial_state]
     )
 
     gr.Examples(
@@ -524,4 +517,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()
+    iface.launch()
+
+