From cc725017a46c4ae059214780dcd8ec33d479e993 Mon Sep 17 00:00:00 2001 From: lindsayshuo Date: Wed, 15 May 2024 10:13:54 +0800 Subject: [PATCH] repaire python infer code (#1526) * Add the generation of multi-class pose engines * Change grids in forwardGpu to one-dimensional arrays * Update README.md * repaire python infer code * Update yolov8_det_trt.py * Update yolov8_pose_trt.py * Update yolov8_seg_trt.py * Optimize parameters * Optimize parameters * Optimize parameters --------- Co-authored-by: lindsayshuo --- yolov8/yolov8_det_trt.py | 28 +++++++++------- yolov8/yolov8_pose_trt.py | 17 +++++----- yolov8/yolov8_seg_trt.py | 67 ++++++++++++++++++++++----------------- 3 files changed, 64 insertions(+), 48 deletions(-) diff --git a/yolov8/yolov8_det_trt.py b/yolov8/yolov8_det_trt.py index f1dd1ae..6413639 100644 --- a/yolov8/yolov8_det_trt.py +++ b/yolov8/yolov8_det_trt.py @@ -10,12 +10,15 @@ import threading import time import cv2 import numpy as np -import pycuda.autoinit +import pycuda.autoinit # noqa: F401 import pycuda.driver as cuda import tensorrt as trt CONF_THRESH = 0.5 IOU_THRESHOLD = 0.4 +POSE_NUM = 17 * 3 +DET_NUM = 6 +SEG_NUM = 32 def get_img_path_batches(batch_size, img_dir): @@ -36,7 +39,7 @@ def plot_one_box(x, img, color=None, label=None, line_thickness=None): """ description: Plots one bounding box on image img, this function comes from YoLov8 project. - param: + param: x: a box likes [x1,y1,x2,y2] img: a opencv image object color: color to draw rectangle, such as (0,255,0) @@ -121,6 +124,7 @@ class YoLov8TRT(object): self.cuda_outputs = cuda_outputs self.bindings = bindings self.batch_size = engine.max_batch_size + self.det_output_length = host_outputs[0].shape[0] def infer(self, raw_image_generator): threading.Thread.__init__(self) @@ -129,7 +133,6 @@ class YoLov8TRT(object): # Restore stream = self.stream context = self.context - engine = self.engine host_inputs = self.host_inputs cuda_inputs = self.cuda_inputs host_outputs = self.host_outputs @@ -167,7 +170,8 @@ class YoLov8TRT(object): # Do postprocess for i in range(self.batch_size): result_boxes, result_scores, result_classid = self.post_process( - output[i * 38001: (i + 1) * 38001], batch_origin_h[i], batch_origin_w[i] + output[i * self.det_output_length: (i + 1) * self.det_output_length], batch_origin_h[i], + batch_origin_w[i] ) # Draw rectangles and labels on the original image for j in range(len(result_boxes)): @@ -279,7 +283,7 @@ class YoLov8TRT(object): """ description: postprocess the prediction param: - output: A numpy likes [num_boxes,cx,cy,w,h,conf,cls_id, cx,cy,w,h,conf,cls_id, ...] + output: A numpy likes [num_boxes,cx,cy,w,h,conf,cls_id, cx,cy,w,h,conf,cls_id, ...] origin_h: height of original image origin_w: width of original image return: @@ -287,10 +291,12 @@ class YoLov8TRT(object): result_scores: finally scores, a numpy, each element is the score correspoing to box result_classid: finally classid, a numpy, each element is the classid correspoing to box """ + num_values_per_detection = DET_NUM + SEG_NUM + POSE_NUM # Get the num of boxes detected num = int(output[0]) # Reshape to a two dimentional ndarray - pred = np.reshape(output[1:], (-1, 38))[:num, :] + # pred = np.reshape(output[1:], (-1, 38))[:num, :] + pred = np.reshape(output[1:], (-1, num_values_per_detection))[:num, :] # Do nms boxes = self.non_max_suppression(pred, origin_h, origin_w, conf_thres=CONF_THRESH, nms_thres=IOU_THRESHOLD) result_boxes = boxes[:, :4] if len(boxes) else np.array([]) @@ -303,7 +309,7 @@ class YoLov8TRT(object): description: compute the IoU of two bounding boxes param: box1: A box coordinate (can be (x1, y1, x2, y2) or (x, y, w, h)) - box2: A box coordinate (can be (x1, y1, x2, y2) or (x, y, w, h)) + box2: A box coordinate (can be (x1, y1, x2, y2) or (x, y, w, h)) x1y1x2y2: select the coordinate format return: iou: computed iou @@ -325,8 +331,8 @@ class YoLov8TRT(object): inter_rect_x2 = np.minimum(b1_x2, b2_x2) inter_rect_y2 = np.minimum(b1_y2, b2_y2) # Intersection area - inter_area = np.clip(inter_rect_x2 - inter_rect_x1 + 1, 0, None) * \ - np.clip(inter_rect_y2 - inter_rect_y1 + 1, 0, None) + inter_area = (np.clip(inter_rect_x2 - inter_rect_x1 + 1, 0, None) + * np.clip(inter_rect_y2 - inter_rect_y1 + 1, 0, None)) # Union Area b1_area = (b1_x2 - b1_x1 + 1) * (b1_y2 - b1_y1 + 1) b2_area = (b2_x2 - b2_x1 + 1) * (b2_y2 - b2_y1 + 1) @@ -403,7 +409,7 @@ class warmUpThread(threading.Thread): if __name__ == "__main__": # load custom plugin and engine PLUGIN_LIBRARY = "build/libmyplugins.so" - engine_file_path = "yolov8n.engine" + engine_file_path = "yolov8s.engine" if len(sys.argv) > 1: engine_file_path = sys.argv[1] @@ -437,7 +443,7 @@ if __name__ == "__main__": try: print('batch size is', yolov8_wrapper.batch_size) - image_dir = "samples/" + image_dir = "images/" image_path_batches = get_img_path_batches(yolov8_wrapper.batch_size, image_dir) for i in range(10): diff --git a/yolov8/yolov8_pose_trt.py b/yolov8/yolov8_pose_trt.py index f56a61f..80f11e7 100644 --- a/yolov8/yolov8_pose_trt.py +++ b/yolov8/yolov8_pose_trt.py @@ -14,10 +14,11 @@ import pycuda.autoinit # noqa: F401 import pycuda.driver as cuda import tensorrt as trt - CONF_THRESH = 0.5 IOU_THRESHOLD = 0.4 - +POSE_NUM = 17 * 3 +DET_NUM = 6 +SEG_NUM = 32 keypoint_pairs = [ (0, 1), (0, 2), (0, 5), (0, 6), (1, 2), (1, 3), (2, 4), (5, 6), (5, 7), (5, 11), @@ -128,7 +129,7 @@ class YoLov8TRT(object): self.cuda_outputs = cuda_outputs self.bindings = bindings self.batch_size = engine.max_batch_size - self.det_output_size = 89001 + self.det_output_size = host_outputs[0].shape[0] def infer(self, raw_image_generator): threading.Thread.__init__(self) @@ -329,7 +330,7 @@ class YoLov8TRT(object): each element represents keypoints for a box, shaped as (#keypoints, 3) """ # Number of values per detection: 38 base values + 17 keypoints * 3 values each - num_values_per_detection = 38 + 17 * 3 + num_values_per_detection = DET_NUM + SEG_NUM + POSE_NUM # Get the number of boxes detected num = int(output[0]) # Reshape to a two-dimensional ndarray with the full detection shape @@ -344,7 +345,7 @@ class YoLov8TRT(object): result_boxes = boxes[:, :4] if len(boxes) else np.array([]) result_scores = boxes[:, 4] if len(boxes) else np.array([]) result_classid = boxes[:, 5] if len(boxes) else np.array([]) - result_keypoints = boxes[:, -51:] if len(boxes) else np.array([]) + result_keypoints = boxes[:, -POSE_NUM:] if len(boxes) else np.array([]) # Return the post-processed results including keypoints return result_boxes, result_scores, result_classid, result_keypoints @@ -404,11 +405,11 @@ class YoLov8TRT(object): # Trandform bbox from [center_x, center_y, w, h] to [x1, y1, x2, y2] res_array = np.copy(boxes) box_pred_deep_copy = np.copy(boxes[:, :4]) - keypoints_pred_deep_copy = np.copy(boxes[:, -51:]) + keypoints_pred_deep_copy = np.copy(boxes[:, -POSE_NUM:]) res_box, res_keypoints = self.xywh2xyxy_with_keypoints( origin_h, origin_w, box_pred_deep_copy, keypoints_pred_deep_copy) res_array[:, :4] = res_box - res_array[:, -51:] = res_keypoints + res_array[:, -POSE_NUM:] = res_keypoints # clip the coordinates res_array[:, 0] = np.clip(res_array[:, 0], 0, origin_w - 1) res_array[:, 2] = np.clip(res_array[:, 2], 0, origin_w - 1) @@ -482,7 +483,7 @@ if __name__ == "__main__": try: print('batch size is', yolov8_wrapper.batch_size) - image_dir = "samples/" + image_dir = "images/" image_path_batches = get_img_path_batches(yolov8_wrapper.batch_size, image_dir) for i in range(10): diff --git a/yolov8/yolov8_seg_trt.py b/yolov8/yolov8_seg_trt.py index e0baec6..b31dd78 100644 --- a/yolov8/yolov8_seg_trt.py +++ b/yolov8/yolov8_seg_trt.py @@ -10,12 +10,15 @@ import threading import time import cv2 import numpy as np -import pycuda.autoinit +import pycuda.autoinit # noqa: F401 import pycuda.driver as cuda import tensorrt as trt CONF_THRESH = 0.5 IOU_THRESHOLD = 0.4 +POSE_NUM = 17 * 3 +DET_NUM = 6 +SEG_NUM = 32 def get_img_path_batches(batch_size, img_dir): @@ -36,7 +39,7 @@ def plot_one_box(x, img, color=None, label=None, line_thickness=None): """ description: Plots one bounding box on image img, this function comes from YoLov8 project. - param: + param: x: a box likes [x1,y1,x2,y2] img: a opencv image object color: color to draw rectangle, such as (0,255,0) @@ -122,18 +125,17 @@ class YoLov8TRT(object): self.bindings = bindings self.batch_size = engine.max_batch_size - #Data length - self.det_output_length = host_outputs[0].shape[0] + # Data length + self.det_output_length = host_outputs[0].shape[0] self.seg_output_length = host_outputs[1].shape[0] self.seg_w = int(self.input_w / 4) self.seg_h = int(self.input_h / 4) self.seg_c = int(self.seg_output_length / (self.seg_w * self.seg_w)) - self.det_row_output_length = self.seg_c + 6 + self.det_row_output_length = self.seg_c + DET_NUM + POSE_NUM # Draw mask self.colors_obj = Colors() - def infer(self, raw_image_generator): threading.Thread.__init__(self) # Make self the active context, pushing it on top of the context stack. @@ -141,7 +143,6 @@ class YoLov8TRT(object): # Restore stream = self.stream context = self.context - engine = self.engine host_inputs = self.host_inputs cuda_inputs = self.cuda_inputs host_outputs = self.host_outputs @@ -181,15 +182,18 @@ class YoLov8TRT(object): output_proto_mask = host_outputs[1] # Do postprocess for i in range(self.batch_size): - result_boxes, result_scores, result_classid,result_proto_coef = self.post_process( - output[i * 38001: (i + 1) * 38001], batch_origin_h[i], batch_origin_w[i] + result_boxes, result_scores, result_classid, result_proto_coef = self.post_process( + output[i * self.det_output_length: (i + 1) * self.det_output_length], batch_origin_h[i], + batch_origin_w[i] ) if result_proto_coef.shape[0] == 0: continue - result_masks = self.process_mask(output_proto_mask, result_proto_coef, result_boxes, batch_origin_h[i], batch_origin_w[i]) - - self.draw_mask(result_masks, colors_=[self.colors_obj(x, True) for x in result_classid],im_src=batch_image_raw[i]) + result_masks = self.process_mask(output_proto_mask, result_proto_coef, result_boxes, batch_origin_h[i], + batch_origin_w[i]) + + self.draw_mask(result_masks, colors_=[self.colors_obj(x, True) for x in result_classid], + im_src=batch_image_raw[i]) # Draw rectangles and labels on the original image for j in range(len(result_boxes)): @@ -301,7 +305,7 @@ class YoLov8TRT(object): """ description: postprocess the prediction param: - output: A numpy likes [num_boxes,cx,cy,w,h,conf,cls_id, cx,cy,w,h,conf,cls_id, ...] + output: A numpy likes [num_boxes,cx,cy,w,h,conf,cls_id, cx,cy,w,h,conf,cls_id, ...] origin_h: height of original image origin_w: width of original image return: @@ -312,22 +316,22 @@ class YoLov8TRT(object): # Get the num of boxes detected num = int(output[0]) # Reshape to a two dimentional ndarray - pred = np.reshape(output[1:], (-1, 38))[:num, :] + pred = np.reshape(output[1:], (-1, self.det_row_output_length))[:num, :] # Do nms boxes = self.non_max_suppression(pred, origin_h, origin_w, conf_thres=CONF_THRESH, nms_thres=IOU_THRESHOLD) result_boxes = boxes[:, :4] if len(boxes) else np.array([]) result_scores = boxes[:, 4] if len(boxes) else np.array([]) result_classid = boxes[:, 5] if len(boxes) else np.array([]) - result_proto_coef = boxes[:, 6:] if len(boxes) else np.array([]) - return result_boxes, result_scores, result_classid,result_proto_coef + result_proto_coef = boxes[:, DET_NUM:int(DET_NUM + SEG_NUM)] if len(boxes) else np.array([]) + return result_boxes, result_scores, result_classid, result_proto_coef def bbox_iou(self, box1, box2, x1y1x2y2=True): """ description: compute the IoU of two bounding boxes param: box1: A box coordinate (can be (x1, y1, x2, y2) or (x, y, w, h)) - box2: A box coordinate (can be (x1, y1, x2, y2) or (x, y, w, h)) + box2: A box coordinate (can be (x1, y1, x2, y2) or (x, y, w, h)) x1y1x2y2: select the coordinate format return: iou: computed iou @@ -349,8 +353,8 @@ class YoLov8TRT(object): inter_rect_x2 = np.minimum(b1_x2, b2_x2) inter_rect_y2 = np.minimum(b1_y2, b2_y2) # Intersection area - inter_area = np.clip(inter_rect_x2 - inter_rect_x1 + 1, 0, None) * \ - np.clip(inter_rect_y2 - inter_rect_y1 + 1, 0, None) + inter_area = (np.clip(inter_rect_x2 - inter_rect_x1 + 1, 0, None) + * np.clip(inter_rect_y2 - inter_rect_y1 + 1, 0, None)) # Union Area b1_area = (b1_x2 - b1_x1 + 1) * (b1_y2 - b1_y1 + 1) b2_area = (b2_x2 - b2_x1 + 1) * (b2_y2 - b2_y1 + 1) @@ -414,17 +418,17 @@ class YoLov8TRT(object): h = self.input_h x = int((self.input_w - w) / 2) y = 0 - crop = mask[y:y+h, x:x+w] + crop = mask[y:y + h, x:x + w] crop = cv2.resize(crop, (iw, ih)) - return crop - + return crop + def process_mask(self, output_proto_mask, result_proto_coef, result_boxes, ih, iw): """ description: Mask pred by yolov8 instance segmentation , - param: + param: output_proto_mask: prototype mask e.g. (32, 160, 160) for 640x640 input result_proto_coef: prototype mask coefficients (n, 32), n represents n results - result_boxes : + result_boxes : ih: rows of original image iw: cols of original image return: @@ -432,13 +436,15 @@ class YoLov8TRT(object): """ result_proto_masks = output_proto_mask.reshape(self.seg_c, self.seg_h, self.seg_w) c, mh, mw = result_proto_masks.shape - masks = self.sigmoid((result_proto_coef @ result_proto_masks.astype(np.float32).reshape(c, -1))).reshape(-1, mh, mw) - + print(result_proto_masks.shape) + print(result_proto_coef.shape) + masks = self.sigmoid((result_proto_coef @ result_proto_masks.astype(np.float32).reshape(c, -1))).reshape(-1, mh, + mw) mask_result = [] for mask, box in zip(masks, result_boxes): mask_s = np.zeros((ih, iw)) - crop_mask = self.scale_mask(mask, ih, iw) + crop_mask = self.scale_mask(mask, ih, iw) x1 = int(box[0]) y1 = int(box[1]) x2 = int(box[2]) @@ -455,7 +461,7 @@ class YoLov8TRT(object): def draw_mask(self, masks, colors_, im_src, alpha=0.5): """ description: Draw mask on image , - param: + param: masks : result_mask colors_: color to draw mask im_src : original image @@ -473,6 +479,7 @@ class YoLov8TRT(object): masks = (masks @ colors_).clip(0, 255) im_src[:] = masks * alpha + im_src * (1 - s * alpha) + class inferThread(threading.Thread): def __init__(self, yolov8_wrapper, image_path_batch): threading.Thread.__init__(self) @@ -498,6 +505,7 @@ class warmUpThread(threading.Thread): batch_image_raw, use_time = self.yolov8_wrapper.infer(self.yolov8_wrapper.get_raw_image_zeros()) print('warm_up->{}, time->{:.2f}ms'.format(batch_image_raw[0].shape, use_time * 1000)) + class Colors: def __init__(self): hexs = ('FF3838', 'FF9D97', 'FF701F', 'FFB21D', 'CFD231', '48F90A', @@ -515,10 +523,11 @@ class Colors: def hex2rgb(h): # rgb order (PIL) return tuple(int(h[1 + i:1 + i + 2], 16) for i in (0, 2, 4)) + if __name__ == "__main__": # load custom plugin and engine PLUGIN_LIBRARY = "build/libmyplugins.so" - engine_file_path = "yolov8s-seg.engine" + engine_file_path = "yolov8n-seg.engine" if len(sys.argv) > 1: engine_file_path = sys.argv[1]