Update app.py

app.py

@@ -288,76 +288,9 @@ 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, choices=[]), 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, detection_confidence, 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)
-
-#             combined_confidence = detection_confidence * top1_prob
-
-#             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 combined_confidence >= 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([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. Please upload a clearer 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,
-#                 "image": annotated_image,
-#                 "is_multi_dog": len(dogs) > 1,
-#                 "dogs_info": explanations
-#             }
-#             return final_explanation, annotated_image, gr.update(visible=True, choices=buttons), 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, choices=[]), initial_state
-
-#     except Exception as e:
-#         error_msg = f"An error occurred: {str(e)}\n\nTraceback:\n{traceback.format_exc()}"
-#         print(error_msg)
-#         return error_msg, None, gr.update(visible=False, choices=[]), None
-
-
 async def predict(image):
     if image is None:
-        return "Please upload an image to start.", None, gr.update(visible=False), None
+        return "Please upload an image to start.", None, gr.update(visible=False, choices=[]), None
 
     try:
         if isinstance(image, np.ndarray):
@@ -365,30 +298,61 @@ async def predict(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, detection_confidence, box) in enumerate(dogs):
             top1_prob, topk_breeds, topk_probs_percent = await predict_single_dog(cropped_image)
-            draw.rectangle(box, outline="#FF0000", width=3)  # 標記框選
+            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)
+
             combined_confidence = detection_confidence * top1_prob
 
-            if combined_confidence >= 0.5:
+            if top1_prob >= 0.5:
                 breed = topk_breeds[0]
                 description = get_dog_description(breed)
-                explanations.append(f"Dog {i+1}: {description}")
+                formatted_description = format_description(description, breed)
+                explanations.append(f"Dog {i+1}: {formatted_description}")
+            elif combined_confidence >= 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([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.")
-        
-        # 返回最終解釋與圖片
-        return "\n".join(explanations), annotated_image
+                explanations.append(f"Dog {i+1}: The image is unclear or the breed is not in the dataset. Please upload a clearer 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,
+                "image": annotated_image,
+                "is_multi_dog": len(dogs) > 1,
+                "dogs_info": explanations
+            }
+            return final_explanation, annotated_image, gr.update(visible=True, choices=buttons), 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, choices=[]), initial_state
 
     except Exception as e:
-        error_msg = f"Error: {str(e)}"
+        error_msg = f"An error occurred: {str(e)}\n\nTraceback:\n{traceback.format_exc()}"
         print(error_msg)
-        return error_msg, None
+        return error_msg, None, gr.update(visible=False, choices=[]), None
 
 
 def show_details(choice, previous_output, initial_state):