agnostic nms

Author: Koldim2001

patched_yolo_infer/nodes/CombineDetections.py

@@ -9,20 +9,23 @@ class CombineDetections:
 
     Args:
         element_crops (MakeCropsDetectThem): Object containing crop information.
-        nms_threshold (float): IoU/IoS threshold for non-maximum suppression.
-        match_metric (str): Matching metric, either 'IOU' or 'IOS'.
+        nms_threshold (float): IoU/IoS threshold for non-maximum suppression.  Dafault is 0.3.
+        match_metric (str): Matching metric, either 'IOU' or 'IOS'. Dafault is IoS.
+        class_agnostic_nms (bool) Determines the NMS mode in object detection. When set to True, NMS 
+            operates across all classes, ignoring class distinctions and suppressing less confident 
+            bounding boxes globally. Otherwise, NMS is applied separately for each class. Default is True.
         intelligent_sorter (bool): Enable sorting by area and rounded confidence parameter. 
-            If False, sorting will be done only by confidence (usual nms). (Dafault True)
+            If False, sorting will be done only by confidence (usual nms). Dafault is True.
         sorter_bins (int): Number of bins to use for intelligent_sorter. A smaller number of bins makes
             the NMS more reliant on object sizes rather than confidence scores. Defaults to 10.
 
     Attributes:
-        conf_treshold (float): Confidence threshold of yolov8.
-        class_names (dict): Dictionary containing class names pf yolov8 model.
+        class_names (dict): Dictionary containing class names of yolo model.
         crops (list): List to store the CropElement objects.
         image (np.ndarray): Source image in BGR.
         nms_threshold (float): IOU/IOS threshold for non-maximum suppression.
         match_metric (str): Matching metric (IOU/IOS).
+        class_agnostic_nms (bool) Determines the NMS mode in object detection.
         intelligent_sorter (bool): Flag indicating whether sorting by area and confidence parameter is enabled.
         sorter_bins (int): Number of bins to use for intelligent_sorter. 
         detected_conf_list_full (list): List of detected confidences.
@@ -46,9 +49,9 @@ def __init__(
         nms_threshold=0.3,
         match_metric='IOS',
         intelligent_sorter=True,
-        sorter_bins=10
+        sorter_bins=10,
+        class_agnostic_nms=True
     ) -> None:
-        self.conf_treshold = element_crops.conf
         self.class_names = element_crops.class_names_dict 
         self.crops = element_crops.crops  # List to store the CropElement objects
         if element_crops.resize_initial_size:
@@ -60,6 +63,7 @@ def __init__(
         self.match_metric = match_metric 
         self.intelligent_sorter = intelligent_sorter # enable sorting by area and confidence parameter
         self.sorter_bins = sorter_bins
+        self.class_agnostic_nms = class_agnostic_nms
 
         # Combinate detections of all patches
         (
@@ -74,27 +78,27 @@ def __init__(
             self.class_names[value] for value in self.detected_cls_id_list_full
         ]  # make str list
 
-        # Invoke the NMS for segmentation masks method for filtering predictions
-        if len(self.detected_masks_list_full) > 0:
-
-            self.filtered_indices = self.nms(
+        # Invoke the NMS:
+        if self.class_agnostic_nms:
+            self.filtered_indices = self.agnostic_nms(
                 self.detected_conf_list_full,
                 self.detected_xyxy_list_full,
                 self.match_metric,
                 self.nms_threshold,
                 self.detected_masks_list_full,
                 intelligent_sorter=self.intelligent_sorter
-            )  # for instance segmentation
+            )  
         else:
-            # Invoke the NMS method for filtering prediction
             self.filtered_indices = self.nms(
                 self.detected_conf_list_full,
                 self.detected_xyxy_list_full,
-                self.match_metric, 
+                self.detected_cls_id_list_full,
+                self.match_metric,
                 self.nms_threshold,
+                self.detected_masks_list_full,
                 intelligent_sorter=self.intelligent_sorter
-            )  # for detection
-
+            ) 
+            
         # Apply filtering (nms output indeces) to the prediction lists
         self.filtered_confidences = [self.detected_conf_list_full[i] for i in self.filtered_indices]
         self.filtered_boxes = [self.detected_xyxy_list_full[i] for i in self.filtered_indices]
@@ -208,7 +212,7 @@ def intersect_over_smaller(mask, masks_list):
             ios_scores.append(ios)
         return torch.tensor(ios_scores)
 
-    def nms(
+    def agnostic_nms(
         self,
         confidences: list,
         boxes: list,
@@ -218,15 +222,15 @@ def nms(
         intelligent_sorter=False,
     ):
         """
-        Apply non-maximum suppression to avoid detecting too many
+        Apply class-agnostic non-maximum suppression to avoid detecting too many
         overlapping bounding boxes for a given object.
 
         Args:
             confidences (list): List of confidence scores.
             boxes (list): List of bounding boxes.
             match_metric (str): Matching metric, either 'IOU' or 'IOS'.
             nms_threshold (float): The threshold for match metric.
-            masks (list, optional): List of masks. Defaults to None.
+            masks (list): List of masks. 
             intelligent_sorter (bool, optional): intelligent sorter 
 
         Returns:
@@ -320,7 +324,7 @@ def nms(
 
             # If masks are provided and IoU based on bounding boxes is greater than 0,
             # calculate IoU for masks and keep the ones with IoU < nms_threshold
-            if masks is not None and torch.any(match_metric_value > 0):
+            if len(masks) > 0 and torch.any(match_metric_value > 0):
 
                 mask_mask = match_metric_value > 0 
 

patched_yolo_infer/nodes/MakeCropsDetectThem.py

@@ -19,7 +19,7 @@ class MakeCropsDetectThem:
         iou (float): IoU threshold for non-maximum suppression YOLOv8 of single crop.
         classes_list (List[int] or None): List of classes to filter detections. If None, 
                                           all classes are considered. Defaults to None.
-        segment (bool): Whether to perform segmentation (YOLOv8-seg).
+        segment (bool): Whether to perform segmentation (YOLO-seg).
         shape_x (int): Size of the crop in the x-coordinate.
         shape_y (int): Size of the crop in the y-coordinate.
         overlap_x (int): Percentage of overlap along the x-axis.
@@ -42,7 +42,7 @@ class MakeCropsDetectThem:
         iou (float): IoU threshold for non-maximum suppression.
         classes_list (List[int] or None): List of classes to filter detections. If None, 
                                           all classes are considered. Defaults to None.
-        segment (bool): Whether to perform segmentation (YOLOv8-seg).
+        segment (bool): Whether to perform segmentation (YOLO-seg).
         shape_x (int): Size of the crop in the x-coordinate.
         shape_y (int): Size of the crop in the y-coordinate.
         overlap_x (int): Percentage of overlap along the x-axis.