Update app.py

app.py

@@ -625,30 +625,21 @@ def preprocess_image(image):
     
 @device_handler
 async def predict_single_dog(image):
-    """
-    Predicts the dog breed using only the classifier.
-    Args:
-        image: PIL Image or numpy array
-    Returns:
-        tuple: (top1_prob, topk_breeds, relative_probs)
-    """
     image_tensor = preprocess_image(image)
     
     with torch.no_grad():
-        # Get model outputs (只使用logits，不需要features)
-        logits = model(image_tensor)[0]  # 如果model仍返回tuple，取第一個元素
+        outputs = model(image_tensor)
+        logits = outputs[0] if isinstance(outputs, tuple) else outputs
         probs = F.softmax(logits, dim=1)
         
-        # Classifier prediction
+        # 其餘代碼保持不變
         top5_prob, top5_idx = torch.topk(probs, k=5)
         breeds = [dog_breeds[idx.item()] for idx in top5_idx[0]]
         probabilities = [prob.item() for prob in top5_prob[0]]
         
-        # Calculate relative probabilities
-        sum_probs = sum(probabilities[:3])  # 只取前三個來計算相對概率
+        sum_probs = sum(probabilities[:3])
         relative_probs = [f"{(prob/sum_probs * 100):.2f}%" for prob in probabilities[:3]]
         
-        # Debug output
         print("\nClassifier Predictions:")
         for breed, prob in zip(breeds[:5], probabilities[:5]):
             print(f"{breed}: {prob:.4f}")