Update app.py

app.py

@@ -143,6 +143,40 @@ def preprocess_image(image):
 def get_akc_breeds_link():
     return "https://www.akc.org/dog-breeds/"
 
+# def predict(image):
+#     try:
+#         image_tensor = preprocess_image(image)
+#         with torch.no_grad():
+#             output = model(image_tensor)
+#             if isinstance(output, tuple):
+#                 logits = output[0]
+#             else:
+#                 logits = output
+#             _, predicted = torch.max(logits, 1) # predicted is the max value's index on dim=1
+#         breed = dog_breeds[predicted.item()]
+        
+#         description = get_dog_description(breed)
+#         akc_link = get_akc_breeds_link()
+        
+#         if isinstance(description, dict):
+#             description_str = "\n\n".join([f"**{key}**: {value}" for key, value in description.items()])
+#         else:
+#             description_str = description
+        
+#         # Add AKC link as an option
+#         description_str += 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."
+        
+#         # Add disclaimer
+#         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.*")
+#         description_str += disclaimer
+        
+#         return description_str
+#     except Exception as e:
+#         return f"An error occurred: {e}"
+
 def predict(image):
     try:
         image_tensor = preprocess_image(image)
@@ -152,30 +186,52 @@ def predict(image):
                 logits = output[0]
             else:
                 logits = output
-            _, predicted = torch.max(logits, 1) # predicted is the max value's index on dim=1
-        breed = dog_breeds[predicted.item()]
+            
+            # 計算預測的概率分佈
+            probabilities = F.softmax(logits, dim=1)
+            
+            # 取得最高的預測分數以及對應的品種
+            top_confidence, top_index = torch.max(probabilities, 1)
+            top_confidence = top_confidence.item()  # 轉成 Python 數值
+            top_breed = dog_breeds[top_index.item()]
+
+            # 如果最高預測分數大於等於 60%，直接返回該品種的資訊
+            if top_confidence >= 0.60:
+                description = get_dog_description(top_breed)
+                akc_link = get_akc_breeds_link()
+                description_str = f"**Breed**: {top_breed}\n\n**Description**: {description}\n"
+                description_str += f"\n\n**Want to learn more about dog breeds?** [Visit the AKC dog breeds page]({akc_link}) and search for {top_breed}."
+                return description_str
+            
+            # 如果預測分數小於 60%，返回 Top-3 預測並讓用戶選擇
+            else:
+                top3_confidences, top3_indices = torch.topk(probabilities, 3, dim=1)
+                top3_breeds = [dog_breeds[idx] for idx in top3_indices.squeeze().tolist()]
+                top3_confidences = top3_confidences.squeeze().tolist()
+
+                return {
+                    "top3_breeds": top3_breeds,
+                    "top3_confidences": [f"{conf * 100:.2f}%" for conf in top3_confidences],
+                    "selected_breed": None,
+                    "message": "The confidence score is low. Please select the correct breed from the options or select 'None of the above' if none are correct."
+                }
+
+    except Exception as e:
+        return f"An error occurred: {e}"
+
+# 處理用戶選擇的結果
+def handle_user_selection(top3_breeds, selected_breed):
+    if selected_breed in top3_breeds:
+        breed_index = top3_breeds.index(selected_breed)
+        description = get_dog_description(selected_breed)
         
-        description = get_dog_description(breed)
         akc_link = get_akc_breeds_link()
-        
-        if isinstance(description, dict):
-            description_str = "\n\n".join([f"**{key}**: {value}" for key, value in description.items()])
-        else:
-            description_str = description
-        
-        # Add AKC link as an option
-        description_str += 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."
-        
-        # Add disclaimer
-        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.*")
-        description_str += disclaimer
-        
+        description_str = f"**Breed**: {selected_breed}\n\n**Description**: {description}\n"
+        description_str += f"\n\n**Want to learn more about dog breeds?** [Visit the AKC dog breeds page]({akc_link}) and search for {selected_breed}."
+
         return description_str
-    except Exception as e:
-        return f"An error occurred: {e}"
+    else:
+        return "Sorry, the breed could not be identified. Please try uploading a clearer image or another breed."
 
 
 iface = gr.Interface(