Update app.py

app.py

@@ -251,17 +251,17 @@ def get_akc_breeds_link():
 #     iface.launch()
 
 def format_description(description, breed):
-    # 分別將不同的屬性分開來顯示，保持結果的可讀性
     if isinstance(description, dict):
-        formatted_description = "\n".join([f"**{key}**: {value}" for key, value in description.items()])
+        formatted_description = "\n\n".join([f"**{key}**: {value}" for key, value in description.items()])
     else:
         formatted_description = description
 
     formatted_description = f"""
 **Breed**: {breed}
+
 {formatted_description}
 
-**Want to learn more about dog breeds?**  
+**Want to learn more about dog breeds?**
 [Visit the AKC dog breeds page]({get_akc_breeds_link()}) and search for {breed} to find detailed information.
 
 *Disclaimer: The external link provided leads to the American Kennel Club (AKC) dog breeds page. 
@@ -271,9 +271,7 @@ Please refer to the AKC's terms of use and privacy policy.*
 """
     return formatted_description
 
-
 def predict_single_dog(image):
-    # 直接使用模型進行預測，無需通過 YOLO
     image_tensor = preprocess_image(image)
     with torch.no_grad():
         output = model(image_tensor)
@@ -285,21 +283,18 @@ def predict_single_dog(image):
         topk_probs_percent = [f"{prob.item() * 100:.2f}%" for prob in topk_probs[0]]
     return top1_prob, topk_breeds, topk_probs_percent
 
-
 def detect_multiple_dogs(image):
-    # 使用 YOLO 檢測多隻狗
     results = model_yolo(image)
     dogs = []
     for result in results:
         for box in result.boxes:
-            if box.cls == 16:  # COCO 資料集中狗的類別是 16
+            if box.cls == 16:  # COCO dataset class for dog is 16
                 xyxy = box.xyxy[0].tolist()
                 confidence = box.conf.item()
                 cropped_image = image.crop((xyxy[0], xyxy[1], xyxy[2], xyxy[3]))
                 dogs.append((cropped_image, confidence, xyxy))
     return dogs
 
-
 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)
@@ -308,28 +303,34 @@ def predict(image):
         if isinstance(image, np.ndarray):
             image = Image.fromarray(image)
 
-        # 首先檢查圖片中是否有多隻狗
+        # Always start with single dog prediction
+        top1_prob, topk_breeds, topk_probs_percent = predict_single_dog(image)
+        
+        # Check if we need to use YOLO for multiple dogs
         dogs = detect_multiple_dogs(image)
-
-        if len(dogs) == 0:
-            # 沒有狗或 YOLO 未檢測到狗，使用單狗直接分類
-            top1_prob, topk_breeds, topk_probs_percent = 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)
+        
+        if len(dogs) <= 1:  # Single dog or no dog detected
             breed = topk_breeds[0]
             description = get_dog_description(breed)
             formatted_description = format_description(description, breed)
-            return formatted_description, image, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
+            
+            if top1_prob >= 0.5:
+                return formatted_description, image, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
+            elif 0.2 <= top1_prob < 0.5:
+                explanation = f"""
+Detected with moderate confidence. Here are the top 3 possible breeds:
 
-        if len(dogs) == 1:
-            # 檢測到一隻狗時，直接分類不使用 YOLO 來節省時間
-            top1_prob, topk_breeds, topk_probs_percent = predict_single_dog(image)
-            breed = topk_breeds[0]
-            description = get_dog_description(breed)
-            formatted_description = format_description(description, breed)
-            return formatted_description, image, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
+1. **{topk_breeds[0]}** ({topk_probs_percent[0]})
+2. **{topk_breeds[1]}** ({topk_probs_percent[1]})
+3. **{topk_breeds[2]}** ({topk_probs_percent[2]})
 
-        # 若有多隻狗，則使用 YOLO 的檢測結果來處理
+Click on a button below to view more information about each 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]}")
+            else:
+                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)
+        
+        # Multiple dogs detected, process each dog
         explanations = []
         visible_buttons = []
         annotated_image = image.copy()
@@ -348,6 +349,7 @@ def predict(image):
             elif 0.2 <= top1_prob < 0.5:
                 explanation = f"""
 Dog {i+1}: Detected with moderate confidence. Here are the top 3 possible breeds:
+
 1. **{topk_breeds[0]}** ({topk_probs_percent[0]})
 2. **{topk_breeds[1]}** ({topk_probs_percent[1]})
 3. **{topk_breeds[2]}** ({topk_probs_percent[2]})
@@ -358,51 +360,27 @@ Dog {i+1}: Detected with moderate confidence. Here are the top 3 possible breeds
                 explanations.append(f"Dog {i+1}: The image is unclear or the breed is not in the dataset.")
 
         final_explanation = "\n\n".join(explanations)
-        return final_explanation, annotated_image, gr.update(visible=len(visible_buttons) >= 1, value=visible_buttons[0] if visible_buttons else ""), gr.update(visible=len(visible_buttons) >= 2, value=visible_buttons[1] if len(visible_buttons) >= 2 else ""), gr.update(visible=len(visible_buttons) >= 3, value=visible_buttons[2] if len(visible_buttons) >= 3 else "")
+        return final_explanation, annotated_image, gr.update(visible=len(visible_buttons) >= 1, value=visible_buttons[0] if visible_buttons else ""), gr.update(visible=True, value=visible_buttons[1] if len(visible_buttons) >= 2 else ""), gr.update(visible=True, value=visible_buttons[2] if len(visible_buttons) >= 3 else "")
 
     except Exception as e:
         return f"An error occurred: {e}", None, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
 
-
 def show_details(breed):
     breed_name = breed.split("More about ")[-1]
     description = get_dog_description(breed_name)
     return format_description(description, breed_name)
 
-
+# Gradio interface setup
 with gr.Blocks(css="""
-    .container {
-        max-width: 900px;
-        margin: 0 auto;
-        padding: 20px;
-        background-color: rgba(255, 255, 255, 0.9);
-        border-radius: 15px;
-        box-shadow: 0 0 20px rgba(0, 0, 0, 0.1);
-    }
-    .gr-form { display: flex; flex-direction: column; align-items: center; }
-    .gr-box { width: 100%; max-width: 500px; }
-    .output-markdown, .output-image {
-        margin-top: 20px;
-        padding: 15px;
-        background-color: #f5f5f5;
-        border-radius: 10px;
-    }
-    .examples {
-        display: flex;
-        justify-content: center;
-        flex-wrap: wrap;
-        gap: 10px;
-        margin-top: 20px;
-    }
-    .examples img {
-        width: 100px;
-        height: 100px;
-        object-fit: cover;
-    }
+    .container { max-width: 900px; margin: auto; padding: 20px; }
+    .gr-box { border-radius: 15px; }
+    .output-markdown { margin-top: 20px; padding: 15px; background-color: #f5f5f5; border-radius: 10px; }
+    .examples { display: flex; justify-content: center; flex-wrap: wrap; gap: 10px; margin-top: 20px; }
+    .examples img { width: 100px; height: 100px; object-fit: cover; }
     """) as iface:
 
-    gr.HTML("<h1 style='font-family:Roboto; font-weight:bold; color:#2C3E50; text-align:center;'>🐶 Dog Breed Classifier 🔍</h1>")
-    gr.HTML("<p style='font-family:Open Sans; color:#34495E; 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>")
+    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")
@@ -426,7 +404,7 @@ with gr.Blocks(css="""
         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%20Breed%20Classifier">Dog Breed Classifier</a>')
+    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%20Breed_Classifier">Dog Breed Classifier</a>')
 
 if __name__ == "__main__":
     iface.launch()