[ADD] Add a python demo.

darknetTR.py

@@ -0,0 +1,185 @@
+#!/usr/bin/env python
+# -*- encoding: utf-8 -*-
+"""
+@File    :   darknetTR.py.py
+@Contact :   JZ
+
+@Modify Time      @Author    @Version    @Desciption
+------------      -------    --------    -----------
+2020/6/12 14:40   JZ      1.0         None
+"""
+
+from ctypes import *
+import cv2
+import numpy as np
+import argparse
+import os
+from threading import Thread
+import time
+
+class IMAGE(Structure):
+    _fields_ = [("w", c_int),
+                ("h", c_int),
+                ("c", c_int),
+                ("data", POINTER(c_float))]
+
+class BOX(Structure):
+    _fields_ = [("x", c_float),
+                ("y", c_float),
+                ("w", c_float),
+                ("h", c_float)]
+
+class DETECTION(Structure):
+    _fields_ = [("cl", c_int),
+                ("bbox", BOX),
+                ("prob", c_float),
+                ("name", c_char*20),
+                ]
+
+lib = CDLL("./build/libdarknetTR.so", RTLD_GLOBAL)
+
+load_network = lib.load_network
+load_network.argtypes = [c_char_p, c_int, c_int]
+load_network.restype = c_void_p
+
+copy_image_from_bytes = lib.copy_image_from_bytes
+copy_image_from_bytes.argtypes = [IMAGE,c_char_p]
+
+make_image = lib.make_image
+make_image.argtypes = [c_int, c_int, c_int]
+make_image.restype = IMAGE
+
+do_inference = lib.do_inference
+do_inference.argtypes = [c_void_p, IMAGE]
+
+get_network_boxes = lib.get_network_boxes
+get_network_boxes.argtypes = [c_void_p, c_float, c_int, POINTER(c_int)]
+get_network_boxes.restype = POINTER(DETECTION)
+
+# cfg = 'yolo4_fp16.rt'
+# netMain = load_network(cfg.encode("ascii"), 80, 1)  # batch size = 1
+#
+#
+# darknet_image = make_image(512, 512, 3)
+# image = cv2.imread('/home/juzheng/dataset/mask/image/20190821004325_55.jpg')
+# frame_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+# image = cv2.resize(frame_rgb,
+#                    (512, 512),
+#                    interpolation=cv2.INTER_LINEAR)
+#
+# # frame_data = np.asarray(image, dtype=np.uint8)
+# # print(frame_data.shape)
+# frame_data = image.ctypes.data_as(c_char_p)
+# copy_image_from_bytes(darknet_image, frame_data)
+#
+# num = c_int(0)
+#
+# pnum = pointer(num)
+# do_inference(netMain, darknet_image)
+# dets = get_network_boxes(netMain, 0.5, 0, pnum)
+# print('end')
+# print(dets[0].cl, dets[0].prob)
+
+def detect_image(net, meta, darknet_image, thresh=.5):
+    num = c_int(0)
+
+    pnum = pointer(num)
+    do_inference(net, darknet_image)
+    dets = get_network_boxes(net, 0.5, 0, pnum)
+    res = []
+    for i in range(pnum[0]):
+        b = dets[i].bbox
+        res.append((dets[i].name.decode("ascii"), dets[i].prob, (b.x, b.y, b.w, b.h)))
+
+    return res
+
+
+def loop_detect(detect_m, video_path):
+    stream = cv2.VideoCapture(video_path)
+    start = time.time()
+    cnt = 0
+    while stream.isOpened():
+        ret, image = stream.read()
+        if ret is False:
+            break
+        frame_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+        image = cv2.resize(frame_rgb,
+                           (512, 512),
+                           interpolation=cv2.INTER_LINEAR)
+        detections = detect_m.detect(image, need_resize=False)
+        cnt += 1
+        for det in detections:
+            print(det)
+    end = time.time()
+    print("frame:{},time:{:.3f},FPS:{:.2f}".format(cnt, end-start, cnt/(end-start)))
+    stream.release()
+
+
+# class myThread(threading.Thread):
+#    def __init__(self, func, args):
+#       threading.Thread.__init__(self)
+#       self.func = func
+#       self.args = args
+#    def run(self):
+#       # print ("Starting " + self.args[0])
+#       self.func(*self.args)
+#       print ("Exiting " )
+
+
+class YOLO4RT(object):
+    def __init__(self,
+                 input_size=512,
+                 weight_file='./yolo4_fp16.rt',
+                 metaPath='Models/yolo4/coco.data',
+                 nms=0.2,
+                 conf_thres=0.3,
+                 device='cuda'):
+        self.input_size = input_size
+        self.metaMain =None
+        self.model = load_network(weight_file.encode("ascii"), 80, 1)
+        self.darknet_image = make_image(input_size, input_size, 3)
+        self.thresh = conf_thres
+        # self.resize_fn = ResizePadding(input_size, input_size)
+        # self.transf_fn = transforms.ToTensor()
+
+    def detect(self, image, need_resize=True, expand_bb=5):
+        try:
+            if need_resize:
+                frame_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+                image = cv2.resize(frame_rgb,
+                                   (self.input_size, self.input_size),
+                                   interpolation=cv2.INTER_LINEAR)
+            frame_data = image.ctypes.data_as(c_char_p)
+
+            copy_image_from_bytes(self.darknet_image, frame_data)
+
+            detections = detect_image(self.model, self.metaMain, self.darknet_image, thresh=self.thresh)
+
+            # cvDrawBoxes(detections, image)
+            # cv2.imshow("1", image)
+            # cv2.waitKey(1)
+            # detections = self.filter_results(detections, "person")
+            return detections
+        except Exception as e_s:
+            print(e_s)
+
+def parse_args():
+    parser = argparse.ArgumentParser(description='tkDNN detect')
+    parser.add_argument('weight', help='rt file path')
+    parser.add_argument('--video',  type=str, help='video path')
+    args = parser.parse_args()
+
+    return args
+
+
+
+if __name__ == '__main__':
+    args = parse_args()
+    detect_m = YOLO4RT(weight_file=args.weight)
+    t = Thread(target=loop_detect, args=(detect_m, args.video), daemon=True)
+
+    # thread1 = myThread(loop_detect, [detect_m])
+
+    # Start new Threads
+    t.start()
+    t.join()