Update app.py

app.py

@@ -271,11 +271,10 @@ Please refer to the AKC's terms of use and privacy policy.*
 """
     return formatted_description
 
-# 預測單隻狗的品種
 async def predict_single_dog(image):
     image_tensor = preprocess_image(image)
     with torch.no_grad():
-        output = model(image_tensor)
+        output = model(image_tensor.to(device))
         logits = output[0] if isinstance(output, tuple) else output
         probabilities = F.softmax(logits, dim=1)
         topk_probs, topk_indices = torch.topk(probabilities, k=3)
@@ -309,7 +308,7 @@ async def predict(image):
         if isinstance(image, np.ndarray):
             image = Image.fromarray(image)
 
-        # First, try to predict without YOLO
+        # Always start with single dog prediction
         top1_prob, topk_breeds, topk_probs_percent = await predict_single_dog(image)
 
         if top1_prob >= 0.5:
@@ -329,50 +328,9 @@ async def predict(image):
             breed_buttons = [f"More about {breed}" for breed in topk_breeds[:3]]
             return explanation, image, gr.update(visible=True, choices=breed_buttons), gr.update(visible=False), gr.update(visible=False)
         
-        elif top1_prob < 0.2:
-            # Only use YOLO if the confidence is very low
-            dogs = await detect_multiple_dogs(image)
-
-            if len(dogs) > 1:
-                # Multiple dogs detected, process each one
-                explanations = []
-                buttons = []
-                annotated_image = image.copy()
-                draw = ImageDraw.Draw(annotated_image)
-                font = ImageFont.load_default()
-
-                for i, (cropped_image, _, box) in enumerate(dogs, 1):
-                    top1_prob, topk_breeds, topk_probs_percent = await predict_single_dog(cropped_image)
-
-                    draw.rectangle(box, outline="red", width=3)
-                    draw.text((box[0], box[1]), f"Dog {i}", fill="yellow", font=font)
-
-                    if top1_prob >= 0.5:
-                        breed = topk_breeds[0]
-                        description = get_dog_description(breed)
-                        explanations.append(format_description(description, breed, is_multi_dog=True, dog_number=i))
-                    else:
-                        explanation = f"""
-    Dog {i}: 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)
-                        for breed in topk_breeds:
-                            buttons.append(f"More about Dog {i}: {breed}")
-
-                final_explanation = "\n\n---\n\n".join(explanations)
-                
-                if buttons:
-                    return final_explanation, annotated_image, gr.update(visible=True, choices=buttons), gr.update(visible=False), gr.update(visible=False)
-                else:
-                    return final_explanation, annotated_image, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
-
-            else:
-                # No dogs detected or only one dog with low confidence
-                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)
-            
+        else:  # 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)
+
     except Exception as e:
         return f"An error occurred: {e}", None, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)