| 8 | |
| 9 | |
| 10 | class Detector: |
| 11 | def __init__( |
| 12 | self, dll_path, trt_path, window_width=640, window_height=640, conf_thresh=0.25, iou_thresh=0.45, |
| 13 | num_class=80): |
| 14 | self.yolo = CDLL(dll_path) |
| 15 | self.max_bbox = 50 |
| 16 | |
| 17 | self.yolo.Detect.argtypes = [c_void_p, c_int, c_int, POINTER(c_ubyte), |
| 18 | npct.ndpointer(dtype=np.float32, ndim=2, shape=(self.max_bbox, 6), |
| 19 | flags="C_CONTIGUOUS")] |
| 20 | |
| 21 | self.yolo.Init.argtypes = [c_char_p, c_int, c_int, c_float, c_float, c_int] |
| 22 | self.yolo.Init.restype = c_void_p |
| 23 | |
| 24 | self.c_point = self.yolo.Init(trt_path.encode('utf-8'), window_width, window_height, conf_thresh, iou_thresh, |
| 25 | num_class) |
| 26 | |
| 27 | def predict(self, img): |
| 28 | rows, cols = img.shape[0], img.shape[1] |
| 29 | res_arr = np.zeros((self.max_bbox, 6), dtype=np.float32) |
| 30 | self.yolo.Detect(self.c_point, c_int(rows), c_int(cols), img.ctypes.data_as(POINTER(c_ubyte)), res_arr) |
| 31 | self.bbox_array = res_arr[~(res_arr == 0).all(1)] |
| 32 | return self.bbox_array |
| 33 | |
| 34 | |
| 35 | class_names = [ |
no outgoing calls
no test coverage detected