Update app.py

app.py

@@ -13,7 +13,7 @@ os.system('pip install ultralytics')
 from ultralytics import YOLO
 
 # 下載YOLOv5預訓練模型
-model_yolo = YOLO('yolov5s.pt')  # 使用 YOLOv5 預訓練模型
+model_yolo = YOLO('yolov8n.pt')  # 使用 YOLOv5 預訓練模型
 
 
 dog_breeds = ["Afghan_Hound", "African_Hunting_Dog", "Airedale", "American_Staffordshire_Terrier", 
@@ -166,6 +166,103 @@ def get_akc_breeds_link():
 #         return f"An error occurred: {e}", gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
 
 
+# def format_description(description, breed):
+#     if isinstance(description, dict):
+#         formatted_description = "\n\n".join([f"**{key}**: {value}" for key, value in description.items()])
+#     else:
+#         formatted_description = description
+
+#     akc_link = get_akc_breeds_link()
+#     formatted_description += f"\n\n**Want to learn more about dog breeds?** [Visit the AKC dog breeds page]({akc_link}) and search for {breed} to find detailed information."
+
+#     disclaimer = ("\n\n*Disclaimer: The external link provided leads to the American Kennel Club (AKC) dog breeds page. "
+#                   "You may need to search for the specific breed on that page. "
+#                   "I am not responsible for the content on external sites. "
+#                   "Please refer to the AKC's terms of use and privacy policy.*")
+#     formatted_description += disclaimer
+
+#     return formatted_description
+
+# def show_details(breed):
+#     breed_name = breed.split("More about ")[-1]
+#     description = get_dog_description(breed_name)
+#     return format_description(description, breed_name)
+
+# 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;
+#     }
+#     """) 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>")
+    
+#     with gr.Row():
+#         input_image = gr.Image(label="Upload a dog image", type="numpy")
+#         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)
+
+#     input_image.change(predict, inputs=input_image, outputs=[output, btn1, btn2, btn3])
+    
+#     btn1.click(show_details, inputs=btn1, outputs=output)
+#     btn2.click(show_details, inputs=btn2, outputs=output)
+#     btn3.click(show_details, inputs=btn3, outputs=output)
+
+#     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%20Breed%20Classifier">Dog Breed Classifier</a>')
+
+# # launch the program
+# if __name__ == "__main__":
+#     iface.launch()
+
+
+# 使用 YOLOv8 進行狗偵測
+def detect_dogs(image):
+    results = yolo_model.predict(image)
+    dogs = []
+
+    for result in results:
+        for box in result.boxes:
+            if box.cls == 16:  # COCO 資料集中的狗類別是16
+                xyxy = box.xyxy
+                confidence = box.conf
+                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:
@@ -177,22 +274,16 @@ def predict(image):
             image = Image.fromarray(image)
 
         # 使用 YOLO 偵測狗
-        results = model_yolo(image)
-        boxes = results[0].boxes  # 提取 YOLO 偵測結果中的邊界框
-
-        if len(boxes) == 0:
-            return "The image is too unclear or the dog breed is not in the dataset. Please upload a clearer image of the dog.", gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
-
-        # 檢查 YOLO 偵測到的邊界框���量
-        print(f"Detected {len(boxes)} dogs in the image.")
+        dogs = detect_dogs(image)
+        if len(dogs) == 0:
+            return "No dogs detected or the image is too unclear. Please upload a clearer image.", gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
 
+        # 開始處理每一隻狗
         explanations = []
         visible_buttons = []
+        annotated_image = image.copy()
 
-        # 處理每個偵測到的狗
-        for i, box in enumerate(boxes):
-            x1, y1, x2, y2 = map(int, box.xyxy[0])
-            cropped_image = image.crop((x1, y1, x2, y2))  # 裁剪出每隻狗的區域
+        for i, (cropped_image, confidence, box) in enumerate(dogs):
             image_tensor = preprocess_image(cropped_image)
 
             with torch.no_grad():
@@ -206,32 +297,36 @@ def predict(image):
                 topk_breeds = [dog_breeds[idx.item()] for idx in topk_indices[0]]
                 topk_probs_percent = [f"{prob.item() * 100:.2f}%" for prob in topk_probs[0]]
 
-                # 根據置信度決定返回結果
+                # 標註狗的邊界框
+                draw = ImageDraw.Draw(annotated_image)
+                draw.rectangle(box.tolist(), outline="red", width=3)
+                draw.text((box[0], box[1]), f"Dog {i+1}", fill="red")
+
+                # 信心度大於 50%，顯示詳細品種資訊
                 if top1_prob >= 0.5:
                     breed = topk_breeds[0]
                     description = get_dog_description(breed)
-                    explanations.append(f"Detected a dog: **{breed}** with {topk_probs_percent[0]} confidence.\n{format_description(description, breed)}")
+                    explanations.append(f"Dog {i+1}: **{breed}**\n{format_description(description, breed)}")
+                # 信心度 20%-49%，顯示 Top 3 品種
                 elif 0.2 <= top1_prob < 0.5:
                     explanation = (
-                        f"Detected a dog with moderate confidence. Here are the top 3 possible breeds:\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"
+                        f"Dog {i+1}: Detected with moderate confidence. Here are the top 3 possible breeds:\n"
+                        f"1. **{topk_breeds[0]}** ({topk_probs_percent[0]})\n"
+                        f"2. **{topk_breeds[1]}** ({topk_probs_percent[1]})\n"
+                        f"3. **{topk_breeds[2]}** ({topk_probs_percent[2]})\n"
                     )
                     explanations.append(explanation)
                     visible_buttons.extend([f"More about {topk_breeds[0]}", f"More about {topk_breeds[1]}", f"More about {topk_breeds[2]}"])
                 else:
-                    explanations.append("The image is too unclear or the breed is not in the dataset. Please upload a clearer image.")
+                    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, gr.update(visible=len(visible_buttons) >= 1), gr.update(visible=len(visible_buttons) >= 2), gr.update(visible=len(visible_buttons) >= 3)
+        return annotated_image, final_explanation, gr.update(visible=len(visible_buttons) >= 1), gr.update(visible=len(visible_buttons) >= 2), gr.update(visible=len(visible_buttons) >= 3)
 
     except Exception as e:
         return f"An error occurred: {e}", gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
 
 
-
 def format_description(description, breed):
     if isinstance(description, dict):
         formatted_description = "\n\n".join([f"**{key}**: {value}" for key, value in description.items()])
@@ -249,11 +344,6 @@ def format_description(description, breed):
 
     return formatted_description
 
-def show_details(breed):
-    breed_name = breed.split("More about ")[-1]
-    description = get_dog_description(breed_name)
-    return format_description(description, breed_name)
-
 with gr.Blocks(css="""
     .container {
         max-width: 900px;
@@ -290,6 +380,7 @@ with gr.Blocks(css="""
     
     with gr.Row():
         input_image = gr.Image(label="Upload a dog image", type="numpy")
+        output_image = gr.Image(label="Annotated Image")
         output = gr.Markdown(label="Prediction Results")
     
     with gr.Row():
@@ -297,7 +388,7 @@ with gr.Blocks(css="""
         btn2 = gr.Button("View More 2", visible=False)
         btn3 = gr.Button("View More 3", visible=False)
 
-    input_image.change(predict, inputs=input_image, outputs=[output, btn1, btn2, btn3])
+    input_image.change(predict, inputs=input_image, outputs=[output_image, output, btn1, btn2, btn3])
     
     btn1.click(show_details, inputs=btn1, outputs=output)
     btn2.click(show_details, inputs=btn2, outputs=output)
@@ -314,3 +405,4 @@ with gr.Blocks(css="""
 if __name__ == "__main__":
     iface.launch()
 
+