Update app.py

app.py

@@ -154,20 +154,61 @@ def format_description(description, breed):
     return formatted_description
 
 
+# async 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 = []
+#     boxes = []
+#     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()
+#             boxes.append((xyxy, confidence))
+    
+#     if not boxes:
+#         dogs.append((image, 1.0, [0, 0, image.width, image.height]))
+#     else:
+#         nms_boxes = non_max_suppression(boxes, iou_threshold)
+        
+#         for box, confidence in nms_boxes:
+#             x1, y1, x2, y2 = box
+#             w, h = x2 - x1, y2 - y1
+#             x1 = max(0, x1 - w * 0.05)
+#             y1 = max(0, y1 - h * 0.05)
+#             x2 = min(image.width, x2 + w * 0.05)
+#             y2 = min(image.height, y2 + h * 0.05)
+#             cropped_image = image.crop((x1, y1, x2, y2))
+#             dogs.append((cropped_image, confidence, [x1, y1, x2, y2]))
+    
+#     return dogs
+
+
 async 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)
+        topk_probs, topk_indices = torch.topk(probabilities, k=5)  # Increase to top 5
         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):
+
+async def detect_multiple_dogs(image, conf_threshold=0.25, iou_threshold=0.5):
     results = model_yolo(image, conf=conf_threshold, iou=iou_threshold)[0]
     dogs = []
     boxes = []
@@ -185,10 +226,10 @@ async def detect_multiple_dogs(image, conf_threshold=0.25, iou_threshold=0.4):
         for box, confidence in nms_boxes:
             x1, y1, x2, y2 = box
             w, h = x2 - x1, y2 - y1
-            x1 = max(0, x1 - w * 0.05)
-            y1 = max(0, y1 - h * 0.05)
-            x2 = min(image.width, x2 + w * 0.05)
-            y2 = min(image.height, y2 + h * 0.05)
+            x1 = max(0, x1 - w * 0.1)
+            y1 = max(0, y1 - h * 0.1)
+            x2 = min(image.width, x2 + w * 0.1)
+            y2 = min(image.height, y2 + h * 0.1)
             cropped_image = image.crop((x1, y1, x2, y2))
             dogs.append((cropped_image, confidence, [x1, y1, x2, y2]))