Update app.py

app.py

@@ -177,7 +177,6 @@ def get_akc_breeds_link():
 #     except Exception as e:
 #         return f"An error occurred: {e}"
 
-# Prediction function
 def predict(image):
     try:
         image_tensor = preprocess_image(image)
@@ -187,37 +186,43 @@ def predict(image):
                 logits = output[0]
             else:
                 logits = output
-            
+
+            # 取得預測的top k結果
             probabilities = F.softmax(logits, dim=1)
-            top_confidence, top_index = torch.max(probabilities, 1)
-            top_confidence = top_confidence.item()
-            top_breed = dog_breeds[top_index.item()]
+            topk_probs, topk_indices = torch.topk(probabilities, k=3)
 
-            # If confidence is higher than 60%, return the top prediction directly
-            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
+            # 檢查最高的預測機率
+            top1_prob = topk_probs[0][0].item()
 
-            # If confidence is lower than 60%, return top 3 results and explain why
-            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()
+            if top1_prob >= 0.5:
+                # 正確辨識時，返回該品種資訊
+                predicted = topk_indices[0][0]
+                breed = dog_breeds[predicted.item()]
+                description = get_dog_description(breed)
+                akc_link = get_akc_breeds_link()
 
-                top3_info = "\n\n".join([f"{i+1}. {breed} ({conf*100:.2f}% confidence)" 
-                                         for i, (breed, conf) in enumerate(zip(top3_breeds, top3_confidences))])
+                description_str = f"**Breed**: {breed}\n\n**Description**: {description}"
+                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."
 
-                # Return top 3 breeds and an explanation
-                message = (f"The model couldn't confidently identify the breed. Here are the top 3 possible breeds:\n\n"
-                           f"{top3_info}\n\n"
-                           "This can happen if the image quality is low or the breed is rare in the dataset. "
-                           "Please try uploading a clearer image or a different angle of the dog. "
-                           "For more accurate results, ensure the dog is the main subject of the photo.")
-                return message
+                return description_str
 
+            else:
+                # 不確定時，返回top 3的預測結果
+                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]]
+
+                # 用粗體返回品種和機率
+                topk_results = "\n".join([f"**{i+1}. {breed}** ({prob} confidence)" for i, (breed, prob) in enumerate(zip(topk_breeds, topk_probs_percent))])
+
+                # 提供說明
+                explanation = (
+                    f"The model couldn't confidently identify the breed. Here are the top 3 possible breeds:\n\n{topk_results}\n\n"
+                    "This can happen if the image quality is low or the breed is rare in the dataset. "
+                    "Please try uploading a clearer image or a different angle of the dog. "
+                    "For more accurate results, ensure the dog is the main subject of the photo."
+                )
+
+                return explanation
     except Exception as e:
         return f"An error occurred: {e}"