Update app.py

app.py

@@ -19,7 +19,7 @@ logging.basicConfig(level=logging.DEBUG)
 logger = logging.getLogger(__name__)
 
 
-model_yolo = YOLO('yolov8s.pt')  
+model_yolo = YOLO('yolov8l.pt')  
 
 
 dog_breeds = ["Afghan_Hound", "African_Hunting_Dog", "Airedale", "American_Staffordshire_Terrier", 
@@ -166,95 +166,58 @@ async def predict_single_dog(image):
     return top1_prob, topk_breeds, topk_probs_percent
     
 
-# async def detect_multiple_dogs(image, conf_threshold=0.35, iou_threshold=0.55):
-#     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
-
-
-# def non_max_suppression(boxes, iou_threshold):
-#     keep = []
-#     boxes = sorted(boxes, key=lambda x: x[1], reverse=True)
-#     while boxes:
-#         current = boxes.pop(0)
-#         keep.append(current)
-#         boxes = [box for box in boxes if calculate_iou(current[0], box[0]) < iou_threshold]
-#     return keep
-
-# def calculate_iou(box1, box2):
-#     x1 = max(box1[0], box2[0])
-#     y1 = max(box1[1], box2[1])
-#     x2 = min(box1[2], box2[2])
-#     y2 = min(box1[3], box2[3])
-    
-#     intersection = max(0, x2 - x1) * max(0, y2 - y1)
-#     area1 = (box1[2] - box1[0]) * (box1[3] - box1[1])
-#     area2 = (box2[2] - box2[0]) * (box2[3] - box2[1])
-    
-#     iou = intersection / float(area1 + area2 - intersection)
-#     return iou
-
-
-def weighted_nms(boxes, scores, iou_threshold=0.5, score_threshold=0.05):
-    keep = nms(boxes, scores, iou_threshold)
-    keep = keep[scores[keep] > score_threshold]
-    return boxes[keep], scores[keep]
-
-async def detect_multiple_dogs(image, conf_threshold=0.3, iou_threshold=0.45):
+async def detect_multiple_dogs(image, conf_threshold=0.35, iou_threshold=0.55):
     results = model_yolo(image, conf=conf_threshold, iou=iou_threshold)[0]
     dogs = []
     boxes = []
-    scores = []
-    
     for box in results.boxes:
         if box.cls == 16:  # COCO dataset class for dog is 16
-            xyxy = box.xyxy[0]
-            confidence = box.conf
-            boxes.append(xyxy)
-            scores.append(confidence)
+            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:
-        boxes = torch.stack(boxes)
-        scores = torch.stack(scores).squeeze(-1)
-        nms_boxes, nms_scores = weighted_nms(boxes, scores, iou_threshold=iou_threshold, score_threshold=0.1)
+        nms_boxes = non_max_suppression(boxes, iou_threshold)
         
-        for box, score in zip(nms_boxes, nms_scores):
-            x1, y1, x2, y2 = box.tolist()
+        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, score.item(), [x1, y1, x2, y2]))
+            dogs.append((cropped_image, confidence, [x1, y1, x2, y2]))
     
     return dogs
 
 
+def non_max_suppression(boxes, iou_threshold):
+    keep = []
+    boxes = sorted(boxes, key=lambda x: x[1], reverse=True)
+    while boxes:
+        current = boxes.pop(0)
+        keep.append(current)
+        boxes = [box for box in boxes if calculate_iou(current[0], box[0]) < iou_threshold]
+    return keep
+
+def calculate_iou(box1, box2):
+    x1 = max(box1[0], box2[0])
+    y1 = max(box1[1], box2[1])
+    x2 = min(box1[2], box2[2])
+    y2 = min(box1[3], box2[3])
+    
+    intersection = max(0, x2 - x1) * max(0, y2 - y1)
+    area1 = (box1[2] - box1[0]) * (box1[3] - box1[1])
+    area2 = (box2[2] - box2[0]) * (box2[3] - box2[1])
+    
+    iou = intersection / float(area1 + area2 - intersection)
+    return iou
+
+
+
 async def process_single_dog(image):
     top1_prob, topk_breeds, topk_probs_percent = await predict_single_dog(image)
     if top1_prob < 0.15: