add a simple warm_up in yolov5_trt.py (wang-xinyu#513)

* add a simple warm_up in yolov5_trt.py

* remove if-else condition in infer()

* be careful with bgr and rgb image

Author: wwqgtxx

yolov5/yolov5_trt.py

@@ -123,7 +123,7 @@ def __init__(self, engine_file_path):
         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 @@ def infer(self, image_path_batch):
         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 @@ def infer(self, image_path_batch):
         # 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 @@ def infer(self, image_path_batch):
                         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 get_raw_image(self, image_path_batch):
+        """
+        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, input_image_path):
+    def preprocess_image(self, raw_bgr_image):
         """
-        description: Read an image from image path, convert it to RGB,
+        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 @@ def preprocess_image(self, input_image_path):
             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 @@ def post_process(self, output, origin_h, origin_w):
         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 @@ def run(self):
     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)
-
-    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()
+    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 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()