From c19ab92172275164bc5e78030062d8b737251e86 Mon Sep 17 00:00:00 2001
From: chutchatut <33798875+chutchatut@users.noreply.github.com>
Date: Thu, 20 Mar 2025 12:21:15 +0000
Subject: [PATCH] add intersection and non max suppression SEGS filter (#940)

---
 __init__.py                  |   4 ++
 modules/impact/segs_nodes.py | 105 +++++++++++++++++++++++++++++++++++
 2 files changed, 109 insertions(+)

diff --git a/__init__.py b/__init__.py
index 97460b9..a3cf345 100644
--- a/__init__.py
+++ b/__init__.py
@@ -247,6 +247,8 @@ NODE_CLASS_MAPPINGS = {
     "ImpactSEGSLabelFilter": SEGSLabelFilter,
     "ImpactSEGSRangeFilter": SEGSRangeFilter,
     "ImpactSEGSOrderedFilter": SEGSOrderedFilter,
+    "ImpactSEGSIntersectionFilter": SEGSIntersectionFilter,
+    "ImpactSEGSNMSFilter": SEGSNMSFilter,
 
     "ImpactCompare": ImpactCompare,
     "ImpactConditionalBranch": ImpactConditionalBranch,
@@ -362,6 +364,8 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "ImpactSEGSLabelFilter": "SEGS Filter (label)",
     "ImpactSEGSRangeFilter": "SEGS Filter (range)",
     "ImpactSEGSOrderedFilter": "SEGS Filter (ordered)",
+    "ImpactSEGSIntersectionFilter": "SEGS Filter (intersection)",
+    "ImpactSEGSNMSFilter": "SEGS Filter (non max suppression)",
     "ImpactSEGSConcat": "SEGS Concat",
     "ImpactSEGSToMaskList": "SEGS to Mask List",
     "ImpactSEGSToMaskBatch": "SEGS to Mask Batch",
diff --git a/modules/impact/segs_nodes.py b/modules/impact/segs_nodes.py
index f0d892a..f3b8fe7 100644
--- a/modules/impact/segs_nodes.py
+++ b/modules/impact/segs_nodes.py
@@ -629,6 +629,111 @@ class SEGSRangeFilter:
         return (segs[0], new_segs), (segs[0], remained_segs),
 
 
+class SEGSIntersectionFilter:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": {
+                        "segs1": ("SEGS", ),
+                        "segs2": ("SEGS", ),
+                        "ioa_threshold": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}),
+                     },
+                }
+
+    RETURN_TYPES = ("SEGS",)
+    RETURN_NAMES = ("filtered_SEGS",)
+    FUNCTION = "doit"
+
+    CATEGORY = "ImpactPack/Util"
+
+    def compute_ioa(self, mask1, mask2):
+        """Compute Intersection over Area (IoA) between two boxes."""
+        inter_mask = utils.bitwise_and_masks(mask1, mask2)
+        
+        inter_area = (inter_mask > 0).sum()
+        area1 = (mask1 > 0).sum()
+
+        return inter_area / area1 if area1 > 0 else 0
+
+    def doit(self, segs1, segs2, ioa_threshold):
+        """Remove segments from segs1 if their IoA with any segment in segs2 exceeds the threshold."""
+        # Extract bounding boxes for all segments in segs1 and segs2
+        keep = []
+
+        # Iterate over all segments in segs1
+        for idx1, seg1 in enumerate(segs1[1]):
+            keep_segment = True  # Assume the segment should be kept
+            mask1 = core.segs_to_combined_mask((segs1[0], [seg1]))
+
+            # Compare with every segment in segs2
+            for seg2 in segs2[1]:
+                mask2 = core.segs_to_combined_mask((segs2[0], [seg2]))
+                ioa = self.compute_ioa(mask1, mask2)  # IoA between segment 1 and segment 2
+                
+                if ioa > ioa_threshold:  # If IoA exceeds the threshold, mark the segment for removal
+                    keep_segment = False
+                    break  # If one overlap exceeds threshold, break early and mark for removal
+
+            # Keep the segment if it did not exceed the threshold with any other segment
+            if keep_segment:
+                keep.append(segs1[1][idx1])
+
+        return (segs1[0], keep),  # Return the updated SEGS
+
+
+class SEGSNMSFilter:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "segs": ("SEGS",),
+                "iou_threshold": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}),
+            },
+        }
+
+    RETURN_TYPES = ("SEGS",)
+    RETURN_NAMES = ("filtered_SEGS",)
+    FUNCTION = "doit"
+
+    CATEGORY = "ImpactPack/Util"
+
+    def compute_iou(self, mask1, mask2):
+        """Compute IoU between two bounding boxes (x1, y1, x2, y2)."""
+        inter_mask = utils.bitwise_and_masks(mask1, mask2)
+        union_mask = utils.add_masks(mask1, mask2)
+        
+        inter_area = (inter_mask > 0).sum()
+        union_area = (union_mask > 0).sum()
+
+        return inter_area / union_area if union_area > 0 else 0
+
+    def doit(self, segs, iou_threshold):
+        """Perform NMS to filter overlapping segments."""
+        confidences = np.ndarray.flatten(np.array([seg.confidence for seg in segs[1]]))
+
+        # Sort boxes by confidence (high to low)
+        sorted_indices = np.argsort(confidences)[::-1].tolist()
+        keep = []
+
+        while len(sorted_indices) > 0:
+            idx = sorted_indices[0]
+            mask1 = core.segs_to_combined_mask((segs[0], [segs[1][idx]]))
+            keep.append(idx)
+            sorted_indices = sorted_indices[1:]
+
+            # Filter indices only contain the indices where the bbox does not intersect
+            filtered_indices = []
+            for i in sorted_indices:
+                mask2 = core.segs_to_combined_mask((segs[0], [segs[1][i]]))
+                iou = self.compute_iou(mask1, mask2)
+                if iou < iou_threshold:
+                    filtered_indices.append(i)
+
+            sorted_indices = np.array(filtered_indices)
+
+        filtered_segs = [segs[1][i] for i in keep]
+        return (segs[0], filtered_segs),
+
+
 class SEGSToImageList:
     @classmethod
     def INPUT_TYPES(s):