Selaa lähdekoodia

add filter_boxes_outside and filter_boxes_completely_outside

quarrying 3 vuotta sitten
vanhempi
commit
2256b0fa76
1 muutettua tiedostoa jossa 26 lisäystä ja 2 poistoa
  1. 26 2
      khandy/boxes/boxes_filter.py

+ 26 - 2
khandy/boxes/boxes_filter.py

@@ -2,7 +2,7 @@ import numpy as np
 
 
 def filter_small_boxes(boxes, min_width, min_height):
-    """Remove all boxes with side smaller than min size. 
+    """Filters all boxes with side smaller than min size. 
 
     Args:
         boxes: a numpy array with shape [N, 4] holding N boxes.
@@ -24,4 +24,28 @@ def filter_small_boxes(boxes, min_width, min_height):
     keep = (widths >= min_width)
     keep &= (heights >= min_height)
     return np.nonzero(keep)[0]
-    
+    
+
+def filter_boxes_outside(boxes, reference_box):
+    """Filters bounding boxes that fall outside reference box.
+    
+    References:
+        `prune_outside_window` in TensorFlow object detection API.
+    """
+    x_min, y_min, x_max, y_max = reference_box[:4]
+    keep = ((boxes[:, 0] >= x_min) & (boxes[:, 1] >= y_min) &
+            (boxes[:, 2] <= x_max) & (boxes[:, 3] <= y_max))
+    return np.nonzero(keep)[0]
+
+
+def filter_boxes_completely_outside(boxes, reference_box):
+    """Filters bounding boxes that fall completely outside of reference box.
+    
+    References:
+        `prune_completely_outside_window` in TensorFlow object detection API.
+    """
+    x_min, y_min, x_max, y_max = reference_box[:4]
+    keep = ((boxes[:, 0] < x_max) & (boxes[:, 1] < y_max) &
+            (boxes[:, 2] > x_min) & (boxes[:, 3] > y_min))
+    return np.nonzero(keep)[0]
+