repaire python infer code

yolov8/yolov8_det_trt.py

@@ -10,7 +10,7 @@
 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
 
@@ -36,7 +36,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)
@@ -129,7 +129,7 @@ def infer(self, raw_image_generator):
         # Restore
         stream = self.stream
         context = self.context
-        engine = self.engine
+        engine = self.engine  # noqa: F841
         host_inputs = self.host_inputs
         cuda_inputs = self.cuda_inputs
         host_outputs = self.host_outputs
@@ -167,7 +167,7 @@ def infer(self, raw_image_generator):
         # 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 * 89001: (i + 1) * 89001], 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 +279,7 @@ def post_process(self, output, origin_h, origin_w):
         """
         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:
@@ -290,7 +290,9 @@ def post_process(self, output, origin_h, origin_w):
         # 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, 89))[: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 +305,7 @@ 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
@@ -325,8 +327,7 @@ def bbox_iou(self, box1, box2, x1y1x2y2=True):
         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)  # noqa: E501
         # 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 +404,7 @@ def run(self):
 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]

yolov8/yolov8_seg_trt.py

@@ -10,7 +10,7 @@
 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
 
@@ -36,7 +36,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,26 +122,25 @@ def __init__(self, engine_file_path):
         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 + 6 + 51
 
         # 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.
         self.ctx.push()
         # Restore
         stream = self.stream
         context = self.context
-        engine = self.engine
+        engine = self.engine  # noqa: F841
         host_inputs = self.host_inputs
         cuda_inputs = self.cuda_inputs
         host_outputs = self.host_outputs
@@ -181,15 +180,18 @@ def infer(self, raw_image_generator):
         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 +303,7 @@ def post_process(self, output, origin_h, origin_w):
         """
         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 +314,22 @@ def post_process(self, output, origin_h, origin_w):
         # 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[:, 6:38] 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 +351,8 @@ def bbox_iou(self, box1, box2, x1y1x2y2=True):
         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)  # noqa: E501
         # 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,31 +416,33 @@ def scale_mask(self, mask, ih, iw):
             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:
             mask_result: (n, ih, iw)
         """
         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 +459,7 @@ def process_mask(self, output_proto_mask, result_proto_coef, result_boxes, ih, i
     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 +477,7 @@ def draw_mask(self, masks, colors_, im_src, alpha=0.5):
         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 +503,7 @@ def run(self):
         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 +521,11 @@ def __call__(self, i, bgr=False):
     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]
@@ -552,7 +559,7 @@ def hex2rgb(h):  # rgb order (PIL)
     try:
         print('batch size is', yolov8_wrapper.batch_size)
 
-        image_dir = "images/"
+        image_dir = "samples/"
         image_path_batches = get_img_path_batches(yolov8_wrapper.batch_size, image_dir)
 
         for i in range(10):