Update app.py

app.py

@@ -249,6 +249,7 @@ def get_akc_breeds_link():
 # if __name__ == "__main__":
 #     iface.launch()
 
+
 def format_description(description, breed):
     if isinstance(description, dict):
         # 確保每一個描述項目換行顯示
@@ -267,6 +268,38 @@ def format_description(description, breed):
 
     return formatted_description
 
+async def predict_single_dog(image):
+    return await asyncio.to_thread(_predict_single_dog, image)
+
+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
+
+async def detect_multiple_dogs(image, conf_threshold=0.3):
+    # 調整 YOLO 模型的置信度閾值
+    return await asyncio.to_thread(_detect_multiple_dogs, image, conf_threshold)
+
+def _detect_multiple_dogs(image, conf_threshold):
+    results = model_yolo(image, conf=conf_threshold)
+    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()
+                if confidence >= conf_threshold:  # 只保留置信度高於閾值的框
+                    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:
         return "Please upload an image to start.", None, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)