Support batch_size when loading the neural network

environment.yml

@@ -23,7 +23,7 @@ dependencies:
   - click-plugins
   - darknet
   - entrypoints
-  - fsspec
+  - fsspec <=0.7.5
   - numpy
   - pillow
 

setup.py

@@ -14,7 +14,7 @@
     "click>=7.0",
     "click-plugins",
     "entrypoints",
-    "fsspec",
+    "fsspec <=0.7.5",
     "numpy",
     "pillow",
     # fmt: on

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, str nms_type, float nms_threshold):
+    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]
 
@@ -26,15 +46,16 @@ cdef class Network:
     cdef dn.network* _c_network
 
     @staticmethod
-    def open(config_url, weights_url):
+    def open(config_url, weights_url, batch_size=1):
         with fsspec_cache_open(config_url, mode="rt") as config:
             with fsspec_cache_open(weights_url, mode="rb") as weights:
-                return Network(config.name, weights.name)
-
+                return Network(config.name, weights.name, batch_size)
 
-    def __cinit__(self, config_file, weights_file):
-        clear = 1
-        self._c_network = dn.load_network(config_file.encode(), weights_file.encode(), clear)
+    def __cinit__(self, str config_file, str weights_file, int batch_size, bint clear=True):
+        self._c_network = dn.load_network_custom(config_file.encode(),
+                                                 weights_file.encode(),
+                                                 clear,
+                                                 batch_size)
         if self._c_network is NULL:
             raise RuntimeError("Failed to create the DarkNet Network...")
 
@@ -43,10 +64,26 @@ cdef class Network:
             dn.free_network(self._c_network[0])
             free(self._c_network)
 
+    @property
+    def batch_size(self):
+        return dn.network_batch_size(self._c_network)
+
     @property
     def shape(self):
         return dn.network_width(self._c_network), dn.network_height(self._c_network)
 
+    @property
+    def width(self):
+        return dn.network_width(self._c_network)
+
+    @property
+    def height(self):
+        return dn.network_height(self._c_network)
+
+    @property
+    def depth(self):
+        return dn.network_depth(self._c_network)
+
     def input_size(self):
         return dn.network_input_size(self._c_network)
 
@@ -81,38 +118,79 @@ 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)
+        rv = convert_detections_to_tuples(detections, num_dets, nms_type, nms_threshold)
+        dn.free_detections(detections, num_dets)
 
-        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']}")
+        return rv
 
+    def detect_batch(self,
+                     np.ndarray[dtype=np.float32_t, ndim=1, mode="c"] frames,
+                     frame_size=None,
+                     float threshold=.5,
+                     float hierarchical_threshold=.5,
+                     int relative=0,
+                     int letterbox=1,
+                     str nms_type="sort",
+                     float nms_threshold=.45
+                     ):
+        pred_width, pred_height = self.shape if frame_size is None else frame_size
+
+        cdef dn.image imr
+        # This looks awkward, but the batch predict *does not* use c, w, h.
+        imr.c = 0
+        imr.w = 0
+        imr.h = 0
+        imr.data = <float *> frames.data
+
+        if frames.size % self.input_size() != 0:
+            raise TypeError("The frames array is not divisible by network input size. "
+                            f"({frames.size} % {self.input_size()} != 0)")
+
+        num_frames = frames.size // self.input_size()
+        if num_frames > self.batch_size:
+            raise TypeError("There are more frames than the configured batch size. "
+                            f"({num_frames} > {self.batch_size})")
+
+        cdef dn.det_num_pair* batch_detections
+        batch_detections = dn.network_predict_batch(
+            self._c_network,
+            imr,
+            num_frames,
+            pred_width,
+            pred_height,
+            threshold,
+            hierarchical_threshold,
+            <int*>0,
+            relative,
+            letterbox
+        )
         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
+            convert_detections_to_tuples(batch_detections[b].dets, batch_detections[b].num, nms_type, nms_threshold)
+            for b in range(num_frames)
         ]
+        dn.free_batch_detections(batch_detections, num_frames)
+        return rv
+
 
-        dn.free_detections(detections, num_dets)
-        return sorted(rv, key=lambda x: x[1], reverse=True)

src/libdarknet/__init__.pxd

@@ -5,13 +5,19 @@ cdef extern from "darknet.h":
     /*
      * darknet.h forgot to extern some useful network functions
      */
+     static int network_depth(network* net) {
+         return net->c;
+     }
+     static int network_batch_size(network* net) {
+         return net->batch;
+     }
      static int network_input_size(network* net) {
-        return net->layers[0].inputs;
+         return net->layers[0].inputs;
      }
      static int network_output_size(network* net) {
-        int i;
-        for(i = net->n-1; i > 0; --i) if(net->layers[i].type != COST) break;
-        return net->layers[i].outputs;
+         int i;
+         for(i = net->n-1; i > 0; --i) if(net->layers[i].type != COST) break;
+         return net->layers[i].outputs;
      }
     """
 
@@ -51,6 +57,13 @@ cdef extern from "darknet.h":
 
     void free_detections(detection* detections, int len)
 
+    ctypedef struct det_num_pair:
+        int num;
+        detection* dets;
+
+    void free_batch_detections(det_num_pair* det_num_pairs, int len)
+
+
     void do_nms_sort(detection* detections, int len, int num_classes, float thresh)
     void do_nms_obj(detection* detections, int len, int num_classes, float thresh)
 
@@ -59,14 +72,19 @@ cdef extern from "darknet.h":
         pass
 
     network* load_network(char* cfg_filename, char* weights_filename, int clear)
+    network* load_network_custom(char* cfg_filename, char* weights_filename, int clear, int batch_size)
     void free_network(network self)
 
+    int network_batch_size(network *self);
     int network_width(network *self);
     int network_height(network *self);
+    int network_depth(network *self);
     int network_input_size(network* self);
     int network_output_size(network* self);
     float* network_predict(network, float* input)
     float* network_predict_image(network*, image)
+
     detection* get_network_boxes(network* self, int width, int height, float thresh, float hier_thresh, int* map, int relative, int* out_len, int letter)
+    det_num_pair* network_predict_batch(network* self, image, int batch_size, int width, int height, float thresh, float hier_thresh, int* map, int relative, int letter)