Update app.py

app.py

@@ -181,42 +181,26 @@ async def predict_single_dog(image):
 #     return dogs
 
 
-async def detect_multiple_dogs(image, conf_threshold=0.25, iou_threshold=0.4):
-    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%"]
+    image_tensor = preprocess(image).unsqueeze(0)
+    with torch.no_grad():
+        output = model(image_tensor)
+        probabilities = torch.nn.functional.softmax(output[0], dim=0)
+        top3_prob, top3_catid = torch.topk(probabilities, 3)
+        top3_breeds = [dog_breeds[idx.item()] for idx in top3_catid]
+        top3_probs = [f"{prob.item()*100:.2f}%" for prob in top3_prob]
+    return top3_prob[0].item(), top3_breeds, top3_probs
+
+async def detect_multiple_dogs(image, conf_threshold=0.3, iou_threshold=0.45):
+    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))
+    return dogs
 
 
 async def process_single_dog(image):
@@ -418,7 +402,7 @@ async def process_single_dog(image):
 
 async def predict(image):
     if image is None:
-        return "Please upload an image to start.", None, gr.update(visible=False), None
+        return "Please upload an image to start.", None, [], gr.update(visible=False), None
 
     try:
         if isinstance(image, np.ndarray):
@@ -426,6 +410,9 @@ async def predict(image):
 
         dogs = await detect_multiple_dogs(image)
         
+        if len(dogs) == 0:
+            dogs = [(image, 1.0, [0, 0, image.width, image.height])]
+        
         color_list = ['#FF0000', '#00FF00', '#0000FF', '#FFFF00', '#00FFFF', '#FF00FF', '#800080', '#FFA500']
         explanations = []
         buttons = []
@@ -476,8 +463,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)}\n\nTraceback:\n{traceback.format_exc()}"
-        logger.error(error_msg)
+        error_msg = f"An error occurred: {str(e)}"
+        print(error_msg)
         return error_msg, None, gr.update(visible=False), None
         
 
@@ -508,8 +495,8 @@ def go_back(state):
         gr.update(visible=False),
         state
     )
+    
 
-# 修改 Gradio 界面結構
 with gr.Blocks() as iface:
     gr.HTML("<h1 style='text-align: center;'>🐶 Dog Breed Classifier 🔍</h1>")
     gr.HTML("<p style='text-align: center;'>Upload a picture of a dog, and the model will predict its breed, provide detailed information, and include an extra information link!</p>")
@@ -543,7 +530,7 @@ with gr.Blocks() as iface:
         inputs=[initial_state],
         outputs=[output, output_image, breed_buttons, back_button, initial_state]
     )
-
+    
     gr.Examples(
         examples=['Border_Collie.jpg', 'Golden_Retriever.jpeg', 'Saint_Bernard.jpeg', 'French_Bulldog.jpeg', 'Samoyed.jpg'],
         inputs=input_image