Update app.py

app.py

@@ -251,33 +251,45 @@ 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()])
+    else:
+        formatted_description = description
+
+    akc_link = get_akc_breeds_link()
+    formatted_description += f"\n\n**Want to learn more about dog breeds?**\n[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 predict_single_dog(image):
+    image_tensor = preprocess_image(image)
+    with torch.no_grad():
+        output = model(image_tensor)
+        logits = output[0] if isinstance(output, tuple) else output
+        probabilities = F.softmax(logits, dim=1)
+        topk_probs, topk_indices = torch.topk(probabilities, k=3)
+        top1_prob = topk_probs[0][0].item()
+        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]]
+    return top1_prob, topk_breeds, topk_probs_percent
+
 def detect_dogs(image):
-    # 使用 YOLO 模型進行偵測
     results = model_yolo(image)
-    
-    # 打印 YOLO 偵測結果
-    print(f"YOLO detection results: {results}")
-    
     dogs = []
     for result in results:
-        # 打印每個結果
-        print(f"Result: {result}")
         for box in result.boxes:
-            # 打印邊界框的類別和座標
-            print(f"Detected class: {box.cls}, Confidence: {box.conf}, Box coordinates: {box.xyxy}")
-            
-            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()
-                
-                # 確認圖片裁切過程正確
-                print(f"Cropping image at coordinates: {xyxy}")
-                
                 cropped_image = image.crop((xyxy[0], xyxy[1], xyxy[2], xyxy[3]))
                 dogs.append((cropped_image, confidence, xyxy))
-    
-    # 最後打印偵測到的狗的數量
-    print(f"Number of dogs detected: {len(dogs)}")
     return dogs
 
 
@@ -301,6 +313,17 @@ def predict(image):
         if isinstance(image, np.ndarray):
             image = Image.fromarray(image)
 
+        # First, try single dog prediction
+        top1_prob, topk_breeds, topk_probs_percent = predict_single_dog(image)
+        
+        if top1_prob >= 0.5:
+            # If confident enough, use single dog prediction
+            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 not confident, use YOLO for multiple dog detection
         dogs = detect_dogs(image)
         if len(dogs) == 0:
             return "No dogs detected or the image is unclear. Please upload a clearer image of a dog.", None, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
@@ -311,7 +334,7 @@ def predict(image):
         draw = ImageDraw.Draw(annotated_image)
 
         for i, (cropped_image, _, box) in enumerate(dogs):
-            top1_prob, topk_breeds, topk_probs_percent = predict_breed(cropped_image)
+            top1_prob, topk_breeds, topk_probs_percent = predict_single_dog(cropped_image)
 
             draw.rectangle(box, outline="red", width=3)
             draw.text((box[0], box[1]), f"Dog {i+1}", fill="red")
@@ -319,7 +342,7 @@ def predict(image):
             if top1_prob >= 0.5:
                 breed = topk_breeds[0]
                 description = get_dog_description(breed)
-                explanations.append(f"Dog {i+1}: **{breed}**\n{format_description(description, breed)}")
+                explanations.append(f"Dog {i+1}:\n**Breed**: {breed}\n{format_description(description, breed)}")
             elif 0.2 <= top1_prob < 0.5:
                 explanation = (
                     f"Dog {i+1}: Detected with moderate confidence. Here are the top 3 possible breeds:\n"
@@ -338,22 +361,6 @@ def predict(image):
     except Exception as e:
         return f"An error occurred: {e}", None, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
 
-def format_description(description, breed):
-    if isinstance(description, dict):
-        formatted_description = "\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]