Update app.py

app.py

@@ -272,6 +272,34 @@ Please refer to the AKC's terms of use and privacy policy.*
 """
     return formatted_description
 
+def predict_single_dog(image):
+    # 直接使用模型進行預測，無需通過 YOLO
+    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_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
+
+
 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)