Jelajahi Sumber

add filter_small_boxes

quarrying 3 tahun lalu
induk
melakukan
1c9c7855f0
2 mengubah file dengan 28 tambahan dan 0 penghapusan
  1. 1 0
      khandy/boxes/__init__.py
  2. 27 0
      khandy/boxes/boxes_filter.py

+ 1 - 0
khandy/boxes/__init__.py

@@ -1,5 +1,6 @@
 from .boxes_clip import *
 from .boxes_overlap import *
+from .boxes_filter import *
 
 from .boxes_transform_flip import *
 from .boxes_transform_rotate import *

+ 27 - 0
khandy/boxes/boxes_filter.py

@@ -0,0 +1,27 @@
+import numpy as np
+
+
+def filter_small_boxes(boxes, min_width, min_height):
+    """Remove all boxes with side smaller than min size. 
+
+    Args:
+        boxes: a numpy array with shape [N, 4] holding N boxes.
+        min_width (float): minimum width
+        min_height (float): minimum height
+
+    Returns:
+        keep: indices of the boxes that have width larger than
+            min_width and height larger than min_height.
+
+    References:
+        `_filter_boxes` in py-faster-rcnn
+        `prune_small_boxes` in TensorFlow object detection API.
+        `structures.Boxes.nonempty` in detectron2
+        `ops.boxes.remove_small_boxes` in torchvision
+    """
+    widths = boxes[:, 2] - boxes[:, 0]
+    heights = boxes[:, 3] - boxes[:, 1] 
+    keep = (widths >= min_width)
+    keep &= (heights >= min_height)
+    return np.nonzero(keep)[0]
+