two models

app.py

@@ -6,7 +6,7 @@ from sahi import AutoDetectionModel
 from PIL import Image
 import plotly.graph_objects as go
 import torch
-import spaces
+#import spaces
 
 import os
 import shutil
@@ -18,40 +18,143 @@ import subprocess
 
 device = "cuda:0" if torch.cuda.is_available() else "cpu"
 
+from torchvision.ops import box_iou
+
+
 
 class Detection:
 
+    # def __init__(self):
+    #     # Set the model path and confidence threshold
+    #     yolov8_model_path = "./model/train_model.pt"  # Update to your model path
+    #     #yolov8_model_path = "./model/best_100epochs_latest.pt"  # Update to your model path
+         
+
+    #     # Initialize the AutoDetectionModel
+    #     self.model = AutoDetectionModel.from_pretrained(
+    #         model_type='yolov8',
+    #         model_path=yolov8_model_path,
+    #         confidence_threshold=0.3,
+    #         device='cpu'  # Change to 'cuda:0' if you are using a GPU
+    #     )
+
+    # def detect_from_image(self, image):
+    #     # Perform sliced prediction with SAHI
+    #     results = get_sliced_prediction(
+    #         image=image,
+    #         detection_model=self.model,
+    #         slice_height=256,
+    #         slice_width=256,
+    #         overlap_height_ratio=0.5,
+    #         overlap_width_ratio=0.5,
+    #         postprocess_type='NMS',
+    #         postprocess_match_metric='IOU',
+    #         postprocess_match_threshold=0.1,
+    #         postprocess_class_agnostic=True,
+    #     )
+
+    #     # Retrieve COCO annotations
+    #     coco_annotations = results.to_coco_annotations()
+    #     return coco_annotations
+
+
+  
     def __init__(self):
-        # Set the model path and confidence threshold
-        yolov8_model_path = "./model/best_100epochs_latest.pt"  # Update to your model path
+        # Set the paths for the two YOLOv8 models
+        yolov8_model_path1 = "./model/train_model.pt"  # Update to your model path
+        yolov8_model_path2 = "./model/best_100epochs_latest.pt"  # Update to the second model path
 
-        # Initialize the AutoDetectionModel
-        self.model = AutoDetectionModel.from_pretrained(
+        # Initialize the AutoDetectionModels
+        self.model1 = AutoDetectionModel.from_pretrained(
             model_type='yolov8',
-            model_path=yolov8_model_path,
+            model_path=yolov8_model_path1,
             confidence_threshold=0.3,
-            device='cpu'  # Change to 'cuda:0' if you are using a GPU
+            device='cuda:0'  # Change to 'cpu' if not using GPU
+        )
+
+        self.model2 = AutoDetectionModel.from_pretrained(
+            model_type='yolov8',
+            model_path=yolov8_model_path2,
+            confidence_threshold=0.3,
+            device='cuda:0'
         )
 
     def detect_from_image(self, image):
-        # Perform sliced prediction with SAHI
-        results = get_sliced_prediction(
+        # Perform sliced prediction with both models
+        results1 = get_sliced_prediction(
             image=image,
-            detection_model=self.model,
+            detection_model=self.model1,
             slice_height=256,
             slice_width=256,
-            overlap_height_ratio=0.2,
-            overlap_width_ratio=0.2,
+            overlap_height_ratio=0.5,
+            overlap_width_ratio=0.5,
             postprocess_type='NMS',
             postprocess_match_metric='IOU',
             postprocess_match_threshold=0.1,
             postprocess_class_agnostic=True,
         )
 
-        # Retrieve COCO annotations
-        coco_annotations = results.to_coco_annotations()
-        return coco_annotations
+        results2 = get_sliced_prediction(
+            image=image,
+            detection_model=self.model2,
+            slice_height=256,
+            slice_width=256,
+            overlap_height_ratio=0.5,
+            overlap_width_ratio=0.5,
+            postprocess_type='NMS',
+            postprocess_match_metric='IOU',
+            postprocess_match_threshold=0.1,
+            postprocess_class_agnostic=True,
+        )
 
+        # Convert results to COCO annotations
+        annotations1 = results1.to_coco_annotations()
+        annotations2 = results2.to_coco_annotations()
+
+        # Combine results from both models
+        combined_annotations = self.combine_results(annotations1, annotations2)
+
+        return combined_annotations
+
+    def combine_results(self, annotations1, annotations2, iou_threshold=0.5):
+        """
+        Combine the results of two sets of annotations, keeping only the higher-confidence
+        prediction when IoU between two bounding boxes is above the threshold.
+
+        :param annotations1: COCO annotations from model 1
+        :param annotations2: COCO annotations from model 2
+        :param iou_threshold: IoU threshold to consider two boxes overlapping
+        :return: Combined annotations list
+        """
+        combined = annotations1.copy()
+
+        for ann2 in annotations2:
+            box2 = ann2['bbox']
+            conf2 = ann2['score']
+
+            keep = True
+            for ann1 in combined:
+                box1 = ann1['bbox']
+                conf1 = ann1['score']
+
+                # Compute IoU between boxes
+                box1_array = np.array([[box1[0], box1[1], box1[0] + box1[2], box1[1] + box1[3]]])
+                box2_array = np.array([[box2[0], box2[1], box2[0] + box2[2], box2[1] + box2[3]]])
+
+                iou = box_iou(torch.tensor(box1_array), torch.tensor(box2_array)).item()
+
+                # If IoU is high and model 2 confidence is lower, discard the annotation from model 2
+                if iou > iou_threshold and conf2 <= conf1:
+                    keep = False
+                    break
+
+            if keep:
+                combined.append(ann2)
+
+        return combined
+
+    #-----------------------------------------------------------------------------------------------------------------------
+    
     def draw_annotations(self, image, annotations):
         """Draw bounding boxes on the image based on COCO annotations using OpenCV."""
         # Define colors for each category in BGR (OpenCV uses BGR format)
@@ -217,7 +320,7 @@ def upload_image(image):
     """Process the uploaded image (if needed) and display it.""" 
     return image
 
-@spaces.GPU
+#@spaces.GPU
 def apply_detection(image):
     """Run object detection on the uploaded image and return the annotated image.""" 
     # Convert image from PIL to NumPy array

model/best.pt

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:67424dbaf2d9c3f07f356a59c37187ef1a7b9f59ebabf77c5cb7f9cb9507f107
-size 38138560