Feature/issue 160 core nms remove ensemble boxes dependency (#161)

* Remove ensemble_boxes dependency

* Remove package from setup

Author: mmeendez8

pyodi/apps/crops/crops_merge.py

@@ -19,7 +19,6 @@ def crops_merge(
     output_file: str,
     predictions_file: Optional[str] = None,
     apply_nms: bool = True,
-    nms_mode: str = "nms",
     score_thr: float = 0.0,
     iou_thr: float = 0.5,
 ) -> str:
@@ -33,7 +32,6 @@ def crops_merge(
             If not None, the annotations of predictions_file will be merged instead of ground_truth_file's.
         apply_nms: Whether to apply Non Maximum Supression to the merged predictions of
             each image. Defaults to True.
-        nms_mode: Non Maximum Supression mode. Defaults to "nms".
         score_thr: Predictions bellow `score_thr` will be filtered. Only used if
             `apply_nms`. Defaults to 0.0.
         iou_thr: None of the filtered predictions will have an iou above `iou_thr` to
@@ -76,12 +74,10 @@ def crops_merge(
         crop["original_image_shape"] = stem_to_original_shape[stem]
 
     if apply_nms:
-        annotations = nms_predictions(
-            annotations, score_thr=score_thr, nms_mode=nms_mode, iou_thr=iou_thr
-        )
+        annotations = nms_predictions(annotations, score_thr=score_thr, iou_thr=iou_thr)
         output_file = str(
             Path(output_file).parent
-            / f"{Path(output_file).stem}_{nms_mode}_{score_thr}_{iou_thr}.json"
+            / f"{Path(output_file).stem}_{score_thr}_{iou_thr}.json"
         )
 
     if predictions_file is not None:

pyodi/core/nms.py

@@ -1,24 +1,61 @@
 from collections import defaultdict
 from typing import Any, Dict, List, Tuple
 
-import ensemble_boxes
 import numpy as np
 from loguru import logger
+from numba import jit
 
 from pyodi.core.boxes import coco_to_corners, corners_to_coco, denormalize, normalize
 
 
+@jit(nopython=True)
+def nms(dets: np.ndarray, scores: np.ndarray, iou_thr: float) -> np.ndarray:
+    """Non Maximum supression algorithm from https://github.com/ZFTurbo/Weighted-Boxes-Fusion/blob/master/ensemble_boxes/ensemble_boxes_nms.py.
+
+    Args:
+        dets: Array of predictions in corner format.
+        scores: Array of scores for each prediction.
+        iou_thr: None of the filtered predictions will have an iou above `iou_thr`
+            to any other.
+
+    Returns:
+        List of filtered predictions in COCO format.
+
+    """
+    x1 = dets[:, 0]
+    y1 = dets[:, 1]
+    x2 = dets[:, 2]
+    y2 = dets[:, 3]
+
+    areas = (x2 - x1) * (y2 - y1)
+    order = scores.argsort()[::-1]
+
+    keep = []
+    while order.size > 0:
+        i = order[0]
+        keep.append(i)
+        xx1 = np.maximum(x1[i], x1[order[1:]])
+        yy1 = np.maximum(y1[i], y1[order[1:]])
+        xx2 = np.minimum(x2[i], x2[order[1:]])
+        yy2 = np.minimum(y2[i], y2[order[1:]])
+
+        w = np.maximum(0.0, xx2 - xx1)
+        h = np.maximum(0.0, yy2 - yy1)
+        inter = w * h
+        ovr = inter / (areas[i] + areas[order[1:]] - inter)
+        inds = np.where(ovr <= iou_thr)[0]
+        order = order[inds + 1]
+
+    return keep
+
+
 def nms_predictions(
-    predictions: List[Dict[Any, Any]],
-    nms_mode: str = "nms",
-    score_thr: float = 0.0,
-    iou_thr: float = 0.5,
+    predictions: List[Dict[Any, Any]], score_thr: float = 0.0, iou_thr: float = 0.5,
 ) -> List[Dict[Any, Any]]:
     """Apply Non Maximum supression to all the images in a COCO `predictions` list.
 
     Args:
         predictions: List of predictions in COCO format.
-        nms_mode: Non Maximum Supression mode. Defaults to "nms".
         score_thr: Predictions below `score_thr` will be filtered. Defaults to 0.0.
         iou_thr: None of the filtered predictions will have an iou above `iou_thr`
             to any other. Defaults to 0.5.
@@ -27,7 +64,6 @@ def nms_predictions(
         List of filtered predictions in COCO format.
 
     """
-    nms = getattr(ensemble_boxes, nms_mode)
     new_predictions = []
     image_id_to_all_boxes: Dict[str, List[List[float]]] = defaultdict(list)
     image_id_to_shape: Dict[str, Tuple[int, int]] = dict()
@@ -51,9 +87,10 @@ def nms_predictions(
         boxes = normalize(coco_to_corners(boxes), image_width, image_height)
 
         logger.info(f"Before nms: {boxes.shape}")
-        boxes, scores, categories = nms(
-            [boxes], [scores], [categories], iou_thr=iou_thr
-        )
+        keep = nms(boxes, scores, iou_thr=iou_thr)
+        # Filter out predictions
+        boxes, categories, scores = boxes[keep], categories[keep], scores[keep]
+
         logger.info(f"After nms: {boxes.shape}")
 
         logger.info(f"Before score threshold: {boxes.shape}")

pyproject.toml

@@ -22,4 +22,4 @@ exclude = '''
 
 [tool.isort]
 profile = "black"
-known_third_party = ["PIL", "ensemble_boxes", "loguru", "matplotlib", "numba", "numpy", "pandas", "plotly", "pycocotools", "pytest", "setuptools", "sklearn", "typer"]
+known_third_party = ["PIL", "loguru", "matplotlib", "numba", "numpy", "pandas", "plotly", "pycocotools", "pytest", "setuptools", "sklearn", "typer"]

setup.cfg

@@ -12,7 +12,6 @@ packages = find_namespace:
 python_requires = >=3.7
 install_requires =
   numpy==1.20.1
-  ensemble-boxes==1.0.4
   loguru==0.5.3
   matplotlib==3.3.4
   numba==0.52.0