Update app.py

app.py

@@ -13,6 +13,10 @@ from urllib.parse import quote
 from ultralytics import YOLO
 import asyncio
 import traceback
+import logging
+
+logging.basicConfig(level=logging.DEBUG)
+logger = logging.getLogger(__name__)
 
 
 # 下載YOLOv8預訓練模型
@@ -152,29 +156,67 @@ def format_description(description, breed):
 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
+# 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.25, iou_threshold=0.4):
+#     results = model_yolo(image, conf=conf_threshold, iou=iou_threshold)[0]
+#     dogs = []
+#     for box in results.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 detect_multiple_dogs(image, conf_threshold=0.25, iou_threshold=0.4):
-    results = model_yolo(image, conf=conf_threshold, iou=iou_threshold)[0]
-    dogs = []
-    for box in results.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
+    try:
+        results = model_yolo(image, conf=conf_threshold, iou=iou_threshold)[0]
+        dogs = []
+        for box in results.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 no dogs are detected, use the whole image
+        if not dogs:
+            logger.info("No dogs detected, using the whole image.")
+            dogs = [(image, 1.0, [0, 0, image.width, image.height])]
+        
+        return dogs
+    except Exception as e:
+        logger.error(f"Error in detect_multiple_dogs: {str(e)}")
+        return [(image, 1.0, [0, 0, image.width, image.height])]
+
+async def predict_single_dog(image):
+    try:
+        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
+    except Exception as e:
+        logger.error(f"Error in predict_single_dog: {str(e)}")
+        return 0, ["Unknown"], ["0%"]
 
 
 async def process_single_dog(image):
@@ -434,8 +476,8 @@ async def predict(image):
             return final_explanation, annotated_image, [], gr.update(visible=False), initial_state
 
     except Exception as e:
-        error_msg = f"An error occurred: {str(e)}"
-        print(error_msg)  # Add log output
+        error_msg = f"An error occurred: {str(e)}\n\nTraceback:\n{traceback.format_exc()}"
+        logger.error(error_msg)
         return error_msg, None, [], gr.update(visible=False), None
 
 def show_details(choice, previous_output, initial_state):