Refactor NMS and detections to tuple array convertion.

Author: sodre

src/darknet/py/network.pyx

@@ -10,6 +10,26 @@ from .util import fsspec_cache_open
 
 np.import_array()
 
+cdef convert_detections_to_tuples(dn.detection* detections, int num_dets, float nms_threshold, str nms_type):
+    if nms_threshold > 0 and num_dets > 0:
+        if nms_type == "obj":
+            dn.do_nms_obj(detections, num_dets, detections[0].classes, nms_threshold)
+        elif nms_type == "sort":
+            dn.do_nms_sort(detections, num_dets, detections[0].classes, nms_threshold)
+        else:
+            raise ValueError(f"non-maximum-suppression type {nms_type} is not one of {['obj', 'sort']}")
+    rv = [
+        (j,
+         detections[i].prob[j],
+         (detections[i].bbox.x, detections[i].bbox.y, detections[i].bbox.w, detections[i].bbox.h)
+        )
+        for i in range(num_dets)
+        for j in range(detections[i].classes)
+        if detections[i].prob[j] > 0
+    ]
+    return sorted(rv, key=lambda x: x[1], reverse=True)
+
+
 cdef class Metadata:
     classes = []  # typing: List[AnyStr]
 
@@ -31,7 +51,6 @@ cdef class Network:
             with fsspec_cache_open(weights_url, mode="rb") as weights:
                 return Network(config.name, weights.name)
 
-
     def __cinit__(self, config_file, weights_file, clear = True, batch_size = 1):
         self._c_network = dn.load_network_custom(config_file.encode(),
                                                  weights_file.encode(),
@@ -83,38 +102,28 @@ cdef class Network:
         output_shape[0] = self.output_size()
         return np.PyArray_SimpleNewFromData(1, output_shape, np.NPY_FLOAT32, output)
 
-    def detect(self, frame_size=None,
-               float threshold=.5, float hierarchical_threshold=.5,
-               int relative=0, int letterbox=1,
-               str nms_type="sort", float nms_threshold=.45,
+    def detect(self,
+               frame_size=None,
+               float threshold=.5,
+               float hierarchical_threshold=.5,
+               int relative=0,
+               int letterbox=1,
+               str nms_type="sort",
+               float nms_threshold=.45,
                ):
-        frame_size = self.shape if frame_size is None else frame_size
+        pred_width, pred_height =  self.shape if frame_size is None else frame_size
+
         cdef int num_dets = 0
         cdef dn.detection* detections
-
         detections = dn.get_network_boxes(self._c_network,
-                                          frame_size[0], frame_size[1],
-                                          threshold, hierarchical_threshold,
+                                          pred_width,
+                                          pred_height,
+                                          threshold,
+                                          hierarchical_threshold,
                                           <int*>0,
                                           relative,
                                           &num_dets,
                                           letterbox)
-
-        if nms_threshold > 0 and num_dets:
-            if nms_type == "obj":
-                dn.do_nms_obj(detections, num_dets, detections[0].classes, nms_threshold)
-            elif nms_type == "sort":
-                dn.do_nms_sort(detections, num_dets, detections[0].classes, nms_threshold)
-            else:
-                raise ValueError(f"non-maximum-suppression type {nms_type} is not one of {['obj', 'sort']}")
-
-        rv = [
-            (j, detections[i].prob[j],
-             (detections[i].bbox.x, detections[i].bbox.y, detections[i].bbox.w, detections[i].bbox.h))
-            for i in range(num_dets)
-            for j in range(detections[i].classes)
-            if detections[i].prob[j] > 0
-        ]
-
+        rv = convert_detections_to_tuples(detections, num_dets, nms_type, nms_threshold)
         dn.free_detections(detections, num_dets)
         return sorted(rv, key=lambda x: x[1], reverse=True)