Update app.py

app.py

@@ -249,19 +249,19 @@ def get_akc_breeds_link():
 # if __name__ == "__main__":
 #     iface.launch()
 
-def format_description(description, breed, is_multi_dog=False, dog_number=None):
+
+def format_description(description, breed):
+    # 分別將不同的屬性分開來顯示，保持結果的可讀性
     if isinstance(description, dict):
-        formatted_description = "\n\n".join([f"**{key}**: {value}" for key, value in description.items() if key != "Breed"])
+        formatted_description = "\n".join([f"**{key}**: {value}" for key, value in description.items()])
     else:
         formatted_description = description
 
-    header = f"**Dog {dog_number}: {breed}**\n\n" if is_multi_dog else f"**Breed: {breed}**\n\n"
-    
     formatted_description = f"""
-{header}
+**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. 
@@ -272,9 +272,14 @@ Please refer to the AKC's terms of use and privacy policy.*
     return formatted_description
 
 async def predict_single_dog(image):
+    # 使用 asyncio.to_thread 將同步操作轉換為異步
+    return await asyncio.to_thread(_predict_single_dog, image)
+
+def _predict_single_dog(image):
+    # 直接使用模型進行預測，無需通過 YOLO
     image_tensor = preprocess_image(image)
     with torch.no_grad():
-        output = model(image_tensor.to(device))
+        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)
@@ -283,34 +288,22 @@ async 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
 
-
 async def detect_multiple_dogs(image):
-    try:
-        img = image.copy()
-        img.thumbnail((640, 640))
-        
-        results = model_yolo(img, conf=0.1)  # 降低閾值以檢測更多狗
-        dogs = []
-        for result in results:
-            for box in result.boxes:
-                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))
-        
-        # 如果只檢測到一隻狗，嘗試檢測其他可能的狗
-        if len(dogs) == 1:
-            # 使用整張圖像進行品種預測
-            full_image_prob, full_image_breeds, _ = await predict_single_dog(image)
-            if full_image_prob >= 0.3 and full_image_breeds[0] != dogs[0][0]:
-                # 如果整張圖像的預測結果不同且置信度較高，添加為第二隻狗
-                dogs.append((image, full_image_prob, [0, 0, image.width, image.height]))
-        
-        return dogs
-    except Exception as e:
-        print(f"Error in detect_multiple_dogs: {e}")
-        return []
+    # 使用 asyncio.to_thread 將同步操作轉換為異步
+    return await asyncio.to_thread(_detect_multiple_dogs, image)
+
+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
+                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
 
 async def predict(image):
     if image is None:
@@ -320,96 +313,60 @@ async def predict(image):
         if isinstance(image, np.ndarray):
             image = Image.fromarray(image)
 
+        # 首先檢查圖片中是否有多隻狗
         dogs = await detect_multiple_dogs(image)
 
         if len(dogs) == 0:
-            # 如果沒有檢測到狗，嘗試對整張圖像進行預測
+            # 沒有狗或 YOLO 未檢測到狗，使用單狗直接分類
             top1_prob, topk_breeds, topk_probs_percent = await predict_single_dog(image)
-            if top1_prob >= 0.3:
-                return await process_single_dog_result(top1_prob, topk_breeds, topk_probs_percent, image, [0, 0, image.width, image.height])
-            else:
-                return "No dogs detected in the image. Please upload a clear image of a dog.", None, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
-        elif len(dogs) == 1:
-            cropped_image, _, box = dogs[0]
-            top1_prob, topk_breeds, topk_probs_percent = await predict_single_dog(cropped_image)
-            return await process_single_dog_result(top1_prob, topk_breeds, topk_probs_percent, image, box)
-        else:
-            return await process_multiple_dogs_result(dogs, image)
-
-    except Exception as e:
-        return f"An error occurred: {str(e)}", None, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
+            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)
+            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)
 
-async def process_multiple_dogs_result(dogs, image):
-    annotated_image = image.copy()
-    draw = ImageDraw.Draw(annotated_image)
-    font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", 20)
-    
-    explanations = []
-    buttons = []
-    
-    for i, (cropped_image, _, box) in enumerate(dogs, 1):
-        top1_prob, topk_breeds, topk_probs_percent = await predict_single_dog(cropped_image)
-        
-        optimized_box = optimize_box(box, image.size)
-        draw.rectangle(optimized_box, outline="red", width=3)
-        draw.text((optimized_box[0], optimized_box[1]), f"Dog {i}", fill="yellow", font=font, stroke_width=2, stroke_fill="black")
-        
-        if top1_prob >= 0.2:
+        if len(dogs) == 1:
+            # 檢測到一隻狗時，直接分類不使用 YOLO 來節省時間
+            top1_prob, topk_breeds, topk_probs_percent = await predict_single_dog(image)
             breed = topk_breeds[0]
             description = get_dog_description(breed)
-            explanation = f"Dog {i}: **{breed}**\n\n"
-            explanation += "\n".join([f"**{key}**: {value}" for key, value in description.items() if key != "Breed"])
-            explanation += f"\n\n**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."
-            explanations.append(explanation)
-            
-            if top1_prob < 0.5:
-                buttons.append(f"More about Dog {i}: {breed}")
-                buttons.append(f"More about Dog {i}: {topk_breeds[1]}")
-                buttons.append(f"More about Dog {i}: {topk_breeds[2]}")
-        else:
-            explanations.append(f"Dog {i}: The image is unclear or the breed is not in the dataset. Please upload a clearer image of this dog.")
+            formatted_description = format_description(description, breed)
+            return formatted_description, image, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
 
-    final_explanation = "\n\n---\n\n".join(explanations)
-    
-    if not explanations:
-        return "No dogs were confidently detected in the image. Please upload a clearer image of a dog.", annotated_image, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
-    
-    return final_explanation, annotated_image, gr.update(visible=bool(buttons), choices=buttons), gr.update(visible=False), gr.update(visible=False)
+        # 若有多隻狗，則使用 YOLO 的檢測結果來處理
+        explanations = []
+        visible_buttons = []
+        annotated_image = image.copy()
+        draw = ImageDraw.Draw(annotated_image)
 
-async def process_single_dog_result(top1_prob, topk_breeds, topk_probs_percent, image, box):
-    annotated_image = image.copy()
-    draw = ImageDraw.Draw(annotated_image)
-    optimized_box = optimize_box(box, image.size)
-    draw.rectangle(optimized_box, outline="red", width=3)
-    draw.text((optimized_box[0], optimized_box[1]), "Dog", fill="yellow", font=ImageFont.load_default())
+        for i, (cropped_image, _, box) in enumerate(dogs):
+            top1_prob, topk_breeds, topk_probs_percent = await predict_single_dog(cropped_image)
 
-    if top1_prob >= 0.5:
-        breed = topk_breeds[0]
-        description = get_dog_description(breed)
-        formatted_description = format_description(description, breed)
-        return formatted_description, annotated_image, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
-    elif top1_prob >= 0.2:
-        explanation = (
-            f"The model couldn't confidently identify the breed. Here are the top 3 possible breeds:\n\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\n"
-            "Click on a button to view more information about the breed."
-        )
-        breed_buttons = [f"More about {breed}" for breed in topk_breeds[:3]]
-        return explanation, annotated_image, gr.update(visible=True, choices=breed_buttons), gr.update(visible=False), gr.update(visible=False)
-    else:
-        return "The image is unclear or the breed is not in the dataset. Please upload a clearer image of a dog.", annotated_image, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
-
-def optimize_box(box, image_size):
-    x1, y1, x2, y2 = box
-    w, h = image_size
-    # 擴大邊界框以確保完整包含狗
-    x1 = max(0, x1 - 10)
-    y1 = max(0, y1 - 10)
-    x2 = min(w, x2 + 10)
-    y2 = min(h, y2 + 10)
-    return [x1, y1, x2, y2]
+            draw.rectangle(box, outline="red", width=3)
+            draw.text((box[0], box[1]), f"Dog {i+1}", fill="red")
+
+            if top1_prob >= 0.5:
+                breed = topk_breeds[0]
+                description = get_dog_description(breed)
+                explanations.append(f"Dog {i+1}:\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:
+1. **{topk_breeds[0]}** ({topk_probs_percent[0]})
+2. **{topk_breeds[1]}** ({topk_probs_percent[1]})
+3. **{topk_breeds[2]}** ({topk_probs_percent[2]})
+"""
+                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(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 "")
+
+    except Exception as e:
+        return f"An error occurred: {e}", None, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
 
 async def show_details(choice):
     if not choice: