Merge pull request #52 from Paper-Chart-Extraction-Project/add-nms-sorting-parameter

Add a Custom Sorting Function for NMS

Author: RyanDoesMath

ChartExtractor/extraction/checkboxes.py

@@ -133,6 +133,7 @@ def detect_checkboxes(
         detections=detections,
         threshold=0.8,
         overlap_comparator=intersection_over_minimum,
+        sorting_fn=lambda det: det.annotation.area * det.annotation.confidence,
     )
     return [det.annotation for det in detections]
 

ChartExtractor/extraction/extraction.py

@@ -331,7 +331,9 @@ def make_document_landmark_detections(
         MODEL_CONFIG["intraoperative_document_landmarks"]["vert_overlap_proportion"],
     )
     detections = non_maximum_suppression(
-        detections, overlap_comparator=intersection_over_minimum
+        detections,
+        overlap_comparator=intersection_over_minimum,
+        sorting_fn=lambda det: det.annotation.area * det.confidence,
     )
     del document_model
     return detections
@@ -479,13 +481,22 @@ def tile_predict(
     )
 
     sys_dets: List[Detection] = non_maximum_suppression(
-        sys_dets, 0.5, intersection_over_minimum
+        sys_dets,
+        0.5,
+        intersection_over_minimum,
+        lambda det: det.annotation.area * det.confidence,
     )
     dia_dets: List[Detection] = non_maximum_suppression(
-        dia_dets, 0.5, intersection_over_minimum
+        dia_dets,
+        0.5,
+        intersection_over_minimum,
+        lambda det: det.annotation.area * det.confidence,
     )
     hr_dets: List[Detection] = non_maximum_suppression(
-        hr_dets, 0.5, intersection_over_minimum
+        hr_dets,
+        0.5,
+        intersection_over_minimum,
+        lambda det: det.annotation.area * det.confidence,
     )
 
     dets: List[Detection] = sys_dets + dia_dets + hr_dets

ChartExtractor/extraction/extraction_utilities.py

@@ -146,6 +146,7 @@ def detect_numbers(
         detections=detections,
         threshold=0.5,
         overlap_comparator=intersection_over_minimum,
+        sorting_fn=lambda det: det.annotation.area * det.confidence,
     )
     return detections
 

ChartExtractor/utilities/annotations.py

@@ -223,6 +223,11 @@ def box(self) -> List[int]:
         """A list containing this `BoundingBox`'s [left, top, right, bottom]."""
         return [self.left, self.top, self.right, self.bottom]
 
+    @property
+    def area(self) -> float:
+        """The area of the box."""
+        return (self.right - self.left) * (self.bottom - self.top)
+
     def set_box(self, new_left: int, new_top: int, new_right: int, new_bottom: int):
         """Sets this BoundingBox's values for left, top, right, bottom.
 

ChartExtractor/utilities/detection_reassembly.py

@@ -95,6 +95,7 @@ def non_maximum_suppression(
     overlap_comparator: Callable[
         [Detection, Detection], float
     ] = intersection_over_union,
+    sorting_fn: Callable[[Detection], bool] = lambda d: d.confidence,
 ) -> List[Detection]:
     """Applies Non-Maximum Suppression (NMS) to a list of detections.
 
@@ -117,7 +118,7 @@ def non_maximum_suppression(
     Returns:
         A list of `Detection` objects containing the filtered detections after applying NMS.
     """
-    detections = sorted(detections, key=lambda d: d.confidence, reverse=True)
+    detections = sorted(detections, key=sorting_fn, reverse=True)
     ix = 0
     while ix < len(detections):
         jx = ix + 1