haotian/duan8
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source 
* add a simple warm_up in yolov5_trt.py

* remove if-else condition in infer()

* be careful with bgr and rgb image
This commit is contained in:
wwqgtxx 2021-04-28 17:37:38 +08:00 committed by GitHub
parent e8653a776d
commit b58a7987bc
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 70 additions and 30 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

100
yolov5/yolov5_trt.py
View File

@ -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()