merged

python/darknet.py

@@ -23,6 +23,15 @@ class BOX(Structure):
                 ("w", c_float),
                 ("h", c_float)]
 
+class DETECTION(Structure):
+    _fields_ = [("bbox", BOX),
+                ("classes", c_int),
+                ("prob", POINTER(c_float)),
+                ("mask", POINTER(c_float)),
+                ("objectness", c_float),
+                ("sort_class", c_int)]
+
+
 class IMAGE(Structure):
     _fields_ = [("w", c_int),
                 ("h", c_int),
@@ -53,9 +62,16 @@ make_image = lib.make_image
 make_image.argtypes = [c_int, c_int, c_int]
 make_image.restype = IMAGE
 
-make_boxes = lib.make_boxes
-make_boxes.argtypes = [c_void_p]
-make_boxes.restype = POINTER(BOX)
+get_network_boxes = lib.get_network_boxes
+get_network_boxes.argtypes = [c_void_p, c_int, c_int, c_float, c_float, POINTER(c_int), c_int]
+get_network_boxes.restype = POINTER(DETECTION)
+
+make_network_boxes = lib.make_network_boxes
+make_network_boxes.argtypes = [c_void_p]
+make_network_boxes.restype = POINTER(DETECTION)
+
+free_detections = lib.free_detections
+free_detections.argtypes = [POINTER(DETECTION), c_int]
 
 free_ptrs = lib.free_ptrs
 free_ptrs.argtypes = [POINTER(c_void_p), c_int]
@@ -64,12 +80,8 @@ num_boxes = lib.num_boxes
 num_boxes.argtypes = [c_void_p]
 num_boxes.restype = c_int
 
-make_probs = lib.make_probs
-make_probs.argtypes = [c_void_p]
-make_probs.restype = POINTER(POINTER(c_float))
-
-detect = lib.network_predict
-detect.argtypes = [c_void_p, IMAGE, c_float, c_float, c_float, POINTER(BOX), POINTER(POINTER(c_float))]
+network_predict = lib.network_predict
+network_predict.argtypes = [c_void_p, POINTER(c_float)]
 
 reset_rnn = lib.reset_rnn
 reset_rnn.argtypes = [c_void_p]
@@ -78,6 +90,12 @@ load_net = lib.load_network
 load_net.argtypes = [c_char_p, c_char_p, c_int]
 load_net.restype = c_void_p
 
+do_nms_obj = lib.do_nms_obj
+do_nms_obj.argtypes = [POINTER(DETECTION), c_int, c_int, c_float]
+
+do_nms_sort = lib.do_nms_sort
+do_nms_sort.argtypes = [POINTER(DETECTION), c_int, c_int, c_float]
+
 free_image = lib.free_image
 free_image.argtypes = [IMAGE]
 
@@ -100,21 +118,6 @@ predict_image = lib.network_predict_image
 predict_image.argtypes = [c_void_p, IMAGE]
 predict_image.restype = POINTER(c_float)
 
-network_detect = lib.network_detect
-network_detect.argtypes = [c_void_p, IMAGE, c_float, c_float, c_float, POINTER(BOX), POINTER(POINTER(c_float))]
-
-import numpy
-def array_to_image(arr):
-    arr = arr.copy()
-    arr = arr.transpose(2,0,1)
-    c = arr.shape[0]
-    h = arr.shape[1]
-    w = arr.shape[2]
-    arr = (arr.astype(numpy.float32)/255.0).flatten()
-    data = c_array(c_float, arr)
-    im = IMAGE(w,h,c,data)
-    return im
-
 def classify(net, meta, im):
     out = predict_image(net, im)
     res = []
@@ -124,24 +127,21 @@ def classify(net, meta, im):
     return res
 
 def detect(net, meta, image, thresh=.5, hier_thresh=.5, nms=.45):
-    if type(image) == numpy.ndarray:
-        im = array_to_image(image)
-    else:
-        im = load_image(image, 0, 0)
-    boxes = make_boxes(net)
-    probs = make_probs(net)
+    im = load_image(image, 0, 0)
     num =   num_boxes(net)
-    network_detect(net, im, thresh, hier_thresh, nms, boxes, probs)
+    predict_image(net, im)
+    dets = get_network_boxes(net, im.w, im.h, thresh, hier_thresh, None, 0)
+    if (nms): do_nms_obj(dets, num, meta.classes, nms);
+
     res = []
     for j in range(num):
         for i in range(meta.classes):
-            if probs[j][i] > 0:
-                res.append((meta.names[i], probs[j][i], (boxes[j].x, boxes[j].y, boxes[j].w, boxes[j].h)))
+            if dets[j].prob[i] > 0:
+                b = dets[j].bbox
+                res.append((meta.names[i], dets[j].prob[i], (b.x, b.y, b.w, b.h)))
     res = sorted(res, key=lambda x: -x[1])
-
-    if type(image) != numpy.ndarray:
-        free_image(im)
-    free_ptrs(cast(probs, POINTER(c_void_p)), num)
+    free_image(im)
+    free_detections(dets, num)
     return res
 
 if __name__ == "__main__":
@@ -153,6 +153,4 @@ if __name__ == "__main__":
     net = load_net("cfg/tiny-yolo.cfg", "tiny-yolo.weights", 0)
     meta = load_meta("cfg/coco.data")
     r = detect(net, meta, "data/dog.jpg")
-    print(r)
-
-
+    print r