diff --git a/yolov5/yolov5_trt.py b/yolov5/yolov5_trt.py index 9986c69..af51f30 100644 --- a/yolov5/yolov5_trt.py +++ b/yolov5/yolov5_trt.py @@ -123,7 +123,7 @@ class YoLov5TRT(object): self.bindings = bindings self.batch_size = engine.max_batch_size - def infer(self, image_path_batch): + def infer(self, raw_image_generator): threading.Thread.__init__(self) # Make self the active context, pushing it on top of the context stack. self.ctx.push() @@ -141,8 +141,8 @@ class YoLov5TRT(object): batch_origin_h = [] batch_origin_w = [] batch_input_image = np.empty(shape=[self.batch_size, 3, self.input_h, self.input_w]) - for i, img_path in enumerate(image_path_batch): - input_image, image_raw, origin_h, origin_w = self.preprocess_image(img_path) + for i, image_raw in enumerate(raw_image_generator): + input_image, image_raw, origin_h, origin_w = self.preprocess_image(image_raw) batch_image_raw.append(image_raw) batch_origin_h.append(origin_h) batch_origin_w.append(origin_w) @@ -166,7 +166,7 @@ class YoLov5TRT(object): # Here we use the first row of output in that batch_size = 1 output = host_outputs[0] # Do postprocess - for i, img_path in enumerate(image_path_batch): + for i in range(self.batch_size): result_boxes, result_scores, result_classid = self.post_process( output[i * 6001: (i + 1) * 6001], batch_origin_h[i], batch_origin_w[i] ) @@ -180,19 +180,29 @@ class YoLov5TRT(object): categories[int(result_classid[j])], result_scores[j] ), ) - parent, filename = os.path.split(img_path) - save_name = os.path.join('output', filename) - # Save image - cv2.imwrite(save_name, batch_image_raw[i]) - print('input->{}, time->{:.2f}ms, saving into output/'.format(image_path_batch, (end - start) * 1000)) + return batch_image_raw, end - start def destroy(self): # Remove any context from the top of the context stack, deactivating it. self.ctx.pop() - - def preprocess_image(self, input_image_path): + + def get_raw_image(self, image_path_batch): """ - description: Read an image from image path, convert it to RGB, + description: Read an image from image path + """ + for img_path in image_path_batch: + yield cv2.imread(img_path) + + def get_raw_image_zeros(self, image_path_batch=None): + """ + description: Ready data for warmup + """ + for _ in range(self.batch_size): + yield np.zeros([self.input_h, self.input_w, 3], dtype=np.uint8) + + def preprocess_image(self, raw_bgr_image): + """ + description: Convert BGR image to RGB, resize and pad it to target size, normalize to [0,1], transform to NCHW format. param: @@ -203,7 +213,7 @@ class YoLov5TRT(object): h: original height w: original width """ - image_raw = cv2.imread(input_image_path) + image_raw = raw_bgr_image h, w, c = image_raw.shape image = cv2.cvtColor(image_raw, cv2.COLOR_BGR2RGB) # Calculate widht and height and paddings @@ -305,22 +315,45 @@ class YoLov5TRT(object): return result_boxes, result_scores, result_classid -class myThread(threading.Thread): - def __init__(self, func, args): +class inferThread(threading.Thread): + def __init__(self, yolov5_wrapper, image_path_batch): threading.Thread.__init__(self) - self.func = func - self.args = args + self.yolov5_wrapper = yolov5_wrapper + self.image_path_batch = image_path_batch def run(self): - self.func(*self.args) + batch_image_raw, use_time = self.yolov5_wrapper.infer(self.yolov5_wrapper.get_raw_image(self.image_path_batch)) + for i, img_path in enumerate(self.image_path_batch): + parent, filename = os.path.split(img_path) + save_name = os.path.join('output', filename) + # Save image + cv2.imwrite(save_name, batch_image_raw[i]) + print('input->{}, time->{:.2f}ms, saving into output/'.format(self.image_path_batch, use_time * 1000)) + + +class warmUpThread(threading.Thread): + def __init__(self, yolov5_wrapper): + threading.Thread.__init__(self) + self.yolov5_wrapper = yolov5_wrapper + + def run(self): + batch_image_raw, use_time = self.yolov5_wrapper.infer(self.yolov5_wrapper.get_raw_image_zeros()) + print('warm_up->{}, time->{:.2f}ms'.format(batch_image_raw[0].shape, use_time * 1000)) + if __name__ == "__main__": # load custom plugins PLUGIN_LIBRARY = "build/libmyplugins.so" - ctypes.CDLL(PLUGIN_LIBRARY) engine_file_path = "build/yolov5s.engine" + if len(sys.argv) > 1: + engine_file_path = sys.argv[1] + if len(sys.argv) > 2: + PLUGIN_LIBRARY = sys.argv[2] + + ctypes.CDLL(PLUGIN_LIBRARY) + # load coco labels categories = ["person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat", "traffic light", @@ -338,15 +371,22 @@ if __name__ == "__main__": os.makedirs('output/') # a YoLov5TRT instance yolov5_wrapper = YoLov5TRT(engine_file_path) - print('batch size is', yolov5_wrapper.batch_size) - image_dir = "samples/" - image_path_batches = get_img_path_batches(yolov5_wrapper.batch_size, image_dir) + try: + print('batch size is', yolov5_wrapper.batch_size) + + image_dir = "samples/" + image_path_batches = get_img_path_batches(yolov5_wrapper.batch_size, image_dir) - for batch in image_path_batches: - # create a new thread to do inference - thread1 = myThread(yolov5_wrapper.infer, [batch]) - thread1.start() - thread1.join() - - # destroy the instance - yolov5_wrapper.destroy() + for i in range(10): + # create a new thread to do warm_up + thread1 = warmUpThread(yolov5_wrapper) + thread1.start() + thread1.join() + for batch in image_path_batches: + # create a new thread to do inference + thread1 = inferThread(yolov5_wrapper, batch) + thread1.start() + thread1.join() + finally: + # destroy the instance + yolov5_wrapper.destroy()