Sourcery refactored main branch

deploy/ONNX/OpenCV/yolo.py

@@ -36,18 +36,15 @@ def draw_label(input_image, label, left, top):
 
 
 def pre_process(input_image, net):
-	# Create a 4D blob from a frame.
-	blob = cv2.dnn.blobFromImage(input_image, 1/255, (INPUT_WIDTH, INPUT_HEIGHT), [0,0,0], 1, crop=False)
+    # Create a 4D blob from a frame.
+    blob = cv2.dnn.blobFromImage(input_image, 1/255, (INPUT_WIDTH, INPUT_HEIGHT), [0,0,0], 1, crop=False)
 
-	# Sets the input to the network.
-	net.setInput(blob)
+    # Sets the input to the network.
+    net.setInput(blob)
 
-	# Runs the forward pass to get output of the output layers.
-	output_layers = net.getUnconnectedOutLayersNames()
-	outputs = net.forward(output_layers)
-	# print(outputs[0].shape)
-
-	return outputs
+    # Runs the forward pass to get output of the output layers.
+    output_layers = net.getUnconnectedOutLayersNames()
+    return net.forward(output_layers)
 
 
 def post_process(input_image, outputs):

deploy/ONNX/OpenCV/yolo_video.py

@@ -32,15 +32,13 @@ def draw_label(im, label, x, y):
     cv2.putText(im, label, (x, y + dim[1]), FONT_FACE, FONT_SCALE, YELLOW, THICKNESS, cv2.LINE_AA)
 
 def pre_process(input_image, net):
-      # Create a 4D blob from a frame.
-      blob = cv2.dnn.blobFromImage(input_image, 1/255,  (INPUT_WIDTH, INPUT_HEIGHT), [0,0,0], 1, crop=False)
+    # Create a 4D blob from a frame.
+    blob = cv2.dnn.blobFromImage(input_image, 1/255,  (INPUT_WIDTH, INPUT_HEIGHT), [0,0,0], 1, crop=False)
 
-      # Sets the input to the network.
-      net.setInput(blob)
-
-      # Run the forward pass to get output of the output layers.
-      outputs = net.forward(net.getUnconnectedOutLayersNames())
-      return outputs
+    # Sets the input to the network.
+    net.setInput(blob)
+
+    return net.forward(net.getUnconnectedOutLayersNames())
 
 def post_process(input_image, outputs):
         # Lists to hold respective values while unwrapping.

deploy/ONNX/OpenCV/yolox.py

@@ -13,10 +13,7 @@ def __init__(self, model, classesFile, p6=False, confThreshold=0.5, nmsThreshold
         self.input_size = (640, 640)
         self.mean = (0.485, 0.456, 0.406)
         self.std = (0.229, 0.224, 0.225)
-        if not p6:
-            self.strides = [8, 16, 32]
-        else:
-            self.strides = [8, 16, 32, 64]
+        self.strides = [8, 16, 32, 64] if p6 else [8, 16, 32]
         self.confThreshold = confThreshold
         self.nmsThreshold = nmsThreshold
         self.objThreshold = objThreshold
@@ -98,17 +95,14 @@ def multiclass_nms(self, boxes, scores):
             valid_score_mask = cls_scores > self.confThreshold
             if valid_score_mask.sum() == 0:
                 continue
-            else:
-                valid_scores = cls_scores[valid_score_mask]
-                valid_boxes = boxes[valid_score_mask]
-                keep = self.nms(valid_boxes, valid_scores)
-                if len(keep) > 0:
-                    cls_inds = np.ones((len(keep), 1)) * cls_ind
-                    dets = np.concatenate([valid_boxes[keep], valid_scores[keep, None], cls_inds], 1)
-                    final_dets.append(dets)
-        if len(final_dets) == 0:
-            return None
-        return np.concatenate(final_dets, 0)
+            valid_scores = cls_scores[valid_score_mask]
+            valid_boxes = boxes[valid_score_mask]
+            keep = self.nms(valid_boxes, valid_scores)
+            if len(keep) > 0:
+                cls_inds = np.ones((len(keep), 1)) * cls_ind
+                dets = np.concatenate([valid_boxes[keep], valid_scores[keep, None], cls_inds], 1)
+                final_dets.append(dets)
+        return np.concatenate(final_dets, 0) if final_dets else None
 
     def vis(self, img, boxes, scores, cls_ids):
         for i in range(len(boxes)):

deploy/ONNX/eval_trt.py

@@ -60,8 +60,7 @@ def run(data,
                 None, device, False, save_dir)
 
     dataloader,pred_result = val.eval_trt(weights)
-    eval_result = val.eval_model(pred_result, dummy_model, dataloader, task)
-    return eval_result
+    return val.eval_model(pred_result, dummy_model, dataloader, task)
 
 
 def main(args):

deploy/ONNX/export_onnx.py

@@ -45,7 +45,10 @@
     # Check device
     cuda = args.device != 'cpu' and torch.cuda.is_available()
     device = torch.device(f'cuda:{args.device}' if cuda else 'cpu')
-    assert not (device.type == 'cpu' and args.half), '--half only compatible with GPU export, i.e. use --device 0'
+    assert (
+        device.type != 'cpu' or not args.half
+    ), '--half only compatible with GPU export, i.e. use --device 0'
+
     # Load PyTorch model
     model = load_checkpoint(args.weights, map_location=device, inplace=True, fuse=True)  # load FP32 model
     for layer in model.modules():
@@ -67,10 +70,6 @@
     dynamic_axes = None
     if args.dynamic_batch:
         args.batch_size = 'batch'
-        dynamic_axes = {
-            'images' :{
-                0:'batch',
-            },}
         if args.end2end:
             output_axes = {
                 'num_dets': {0: 'batch'},
@@ -82,8 +81,11 @@
             output_axes = {
                 'outputs': {0: 'batch'},
             }
-        dynamic_axes.update(output_axes)
-
+        dynamic_axes = {
+            'images': {
+                0: 'batch',
+            },
+        } | output_axes
 
     if args.end2end:
         from yolov6.models.end2end import End2End
@@ -134,16 +136,15 @@
 
     # Finish
     LOGGER.info('\nExport complete (%.2fs)' % (time.time() - t))
-    if args.end2end:
-        if not args.ort:
-            info = f'trtexec --onnx={export_file} --saveEngine={export_file.replace(".onnx",".engine")}'
-            if args.dynamic_batch:
-                LOGGER.info('Dynamic batch export should define min/opt/max batchsize\n'+
-                            'We set min/opt/max = 1/16/32 default!')
-                wandh = 'x'.join(list(map(str,args.img_size)))
-                info += (f' --minShapes=images:1x3x{wandh}'+
-                f' --optShapes=images:16x3x{wandh}'+
-                f' --maxShapes=images:32x3x{wandh}'+
-                f' --shapes=images:16x3x{wandh}')
-            LOGGER.info('\nYou can export tensorrt engine use trtexec tools.\nCommand is:')
-            LOGGER.info(info)
+    if args.end2end and not args.ort:
+        info = f'trtexec --onnx={export_file} --saveEngine={export_file.replace(".onnx",".engine")}'
+        if args.dynamic_batch:
+            LOGGER.info('Dynamic batch export should define min/opt/max batchsize\n'+
+                        'We set min/opt/max = 1/16/32 default!')
+            wandh = 'x'.join(list(map(str,args.img_size)))
+            info += (f' --minShapes=images:1x3x{wandh}'+
+            f' --optShapes=images:16x3x{wandh}'+
+            f' --maxShapes=images:32x3x{wandh}'+
+            f' --shapes=images:16x3x{wandh}')
+        LOGGER.info('\nYou can export tensorrt engine use trtexec tools.\nCommand is:')
+        LOGGER.info(info)

deploy/TensorRT/Processor.py

@@ -20,7 +20,7 @@ def torch_dtype_from_trt(dtype):
     elif dtype == trt.float32:
         return torch.float32
     else:
-        raise TypeError('%s is not supported by torch' % dtype)
+        raise TypeError(f'{dtype} is not supported by torch')
 
 
 def torch_device_from_trt(device):
@@ -29,7 +29,7 @@ def torch_device_from_trt(device):
     elif device == trt.TensorLocation.HOST:
         return torch.device('cpu')
     else:
-        return TypeError('%s is not supported by torch' % device)
+        return TypeError(f'{device} is not supported by torch')
 
 
 def get_input_shape(engine):
@@ -42,7 +42,7 @@ def get_input_shape(engine):
     elif len(binding_dims) == 3:
         return tuple(binding_dims[1:])
     else:
-        raise ValueError('bad dims of binding %s: %s' % (binding, str(binding_dims)))
+        raise ValueError(f'bad dims of binding {binding}: {str(binding_dims)}')
 
 def letterbox(im, new_shape=(640, 640), color=(114, 114, 114), auto=True, scaleup=False, stride=32, return_int=False):
     # Resize and pad image while meeting stride-multiple constraints
@@ -70,10 +70,7 @@ def letterbox(im, new_shape=(640, 640), color=(114, 114, 114), auto=True, scaleu
     top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))
     left, right = int(round(dw - 0.1)), int(round(dw + 0.1))
     im = cv2.copyMakeBorder(im, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color)  # add border
-    if not return_int:
-        return im, r, (dw, dh)
-    else:
-        return im, r, (left, top)
+    return (im, r, (left, top)) if return_int else (im, r, (dw, dh))
 
 
 class Processor():
@@ -89,8 +86,8 @@ def __init__(self, model, num_classes=80, num_layers=3, anchors=1, device=torch.
             self.engine = self.runtime.deserialize_cuda_engine(f.read())
         self.input_shape = get_input_shape(self.engine)
         self.bindings = OrderedDict()
-        self.input_names = list()
-        self.output_names = list()
+        self.input_names = []
+        self.output_names = []
         for index in range(self.engine.num_bindings):
             name = self.engine.get_binding_name(index)
             if self.engine.binding_is_input(index):
@@ -126,8 +123,7 @@ def __init__(self, model, num_classes=80, num_layers=3, anchors=1, device=torch.
     def detect(self, img):
         """Detect objects in the input image."""
         resized, _ = self.pre_process(img, self.input_shape)
-        outputs = self.inference(resized)
-        return outputs
+        return self.inference(resized)
 
     def pre_process(self, img_src, input_shape=None,):
         """Preprocess an image before TRT YOLO inferencing.
@@ -144,9 +140,7 @@ def inference(self, inputs):
         self.binding_addrs[self.input_names[0]] = int(inputs.data_ptr())
         #self.binding_addrs['x2paddle_image_arrays'] = int(inputs.data_ptr())
         self.context.execute_v2(list(self.binding_addrs.values()))
-        output = self.bindings[self.output_names[0]].data
-        #output = self.bindings['save_infer_model/scale_0.tmp_0'].data
-        return output
+        return self.bindings[self.output_names[0]].data
 
     def output_reformate(self, outputs):
         z = []
@@ -203,7 +197,6 @@ def non_max_suppression(self, prediction, conf_thres=0.25, iou_thres=0.45, class
         Returns:
              list of detections, echo item is one tensor with shape (num_boxes, 6), 6 is for [xyxy, conf, cls].
         """
-        num_classes = prediction.shape[2] - 5  # number of classes
         pred_candidates = prediction[..., 4] > conf_thres  # candidates
 
         # Check the parameters.
@@ -214,7 +207,7 @@ def non_max_suppression(self, prediction, conf_thres=0.25, iou_thres=0.45, class
         max_wh = 4096  # maximum box width and height
         max_nms = 30000  # maximum number of boxes put into torchvision.ops.nms()
         time_limit = 10.0  # quit the function when nms cost time exceed the limit time.
-        multi_label &= num_classes > 1  # multiple labels per box
+        multi_label &= prediction.shape[2] > 6
 
         tik = time.time()
         output = [torch.zeros((0, 6), device=prediction.device)] * prediction.shape[0]

deploy/TensorRT/calibrator.py

@@ -75,26 +75,27 @@ def __init__(self, batch_size, batch_num, calib_img_dir, input_w, input_h):
         self.input_w = input_w
         # self.img_list = [i.strip() for i in open('calib.txt').readlines()]
         self.img_list = glob.glob(os.path.join(calib_img_dir, "*.jpg"))
-        assert len(self.img_list) > self.batch_size * self.length, \
-            '{} must contains more than '.format(calib_img_dir) + str(self.batch_size * self.length) + ' images to calib'
-        print('found all {} images to calib.'.format(len(self.img_list)))
+        assert (
+            len(self.img_list) > self.batch_size * self.length
+        ), f'{calib_img_dir} must contains more than {str(self.batch_size * self.length)} images to calib'
+
+        print(f'found all {len(self.img_list)} images to calib.')
         self.calibration_data = np.zeros((self.batch_size, 3, input_h, input_w), dtype=np.float32)
 
     def reset(self):
         self.index = 0
 
     def next_batch(self):
-        if self.index < self.length:
-            for i in range(self.batch_size):
-                assert os.path.exists(self.img_list[i + self.index * self.batch_size]), 'not found!!'
-                img = cv2.imread(self.img_list[i + self.index * self.batch_size])
-                img = precess_image(img, self.input_h, 32)
-                self.calibration_data[i] = img
-
-            self.index += 1
-            return np.ascontiguousarray(self.calibration_data, dtype=np.float32)
-        else:
+        if self.index >= self.length:
             return np.array([])
+        for i in range(self.batch_size):
+            assert os.path.exists(self.img_list[i + self.index * self.batch_size]), 'not found!!'
+            img = cv2.imread(self.img_list[i + self.index * self.batch_size])
+            img = precess_image(img, self.input_h, 32)
+            self.calibration_data[i] = img
+
+        self.index += 1
+        return np.ascontiguousarray(self.calibration_data, dtype=np.float32)
 
     def __len__(self):
         return self.length

deploy/TensorRT/eval_yolo_trt.py

@@ -55,16 +55,15 @@ def parse_args():
     parser.add_argument('--force_no_pad', type=bool, default=True, help='for no extra pad in letterbox')
     parser.add_argument('-v', '--visualize', action="store_true", default=False, help='visualize demo')
     parser.add_argument('--num_imgs_to_visualize', type=int, default=10, help='number of images to visualize')
-    args = parser.parse_args()
-    return args
+    return parser.parse_args()
 
 
 def check_args(args):
     """Check and make sure command-line arguments are valid."""
     if not os.path.isdir(args.imgs_dir):
-        sys.exit('%s is not a valid directory' % args.imgs_dir)
+        sys.exit(f'{args.imgs_dir} is not a valid directory')
     if not os.path.isfile(args.annotations):
-        sys.exit('%s is not a valid file' % args.annotations)
+        sys.exit(f'{args.annotations} is not a valid file')
 
 
 def generate_results(processor, imgs_dir, jpgs, results_file, conf_thres, iou_thres, non_coco,
@@ -102,7 +101,7 @@ def generate_results(processor, imgs_dir, jpgs, results_file, conf_thres, iou_th
             image_ids.append(int(jpgs[idx].split('.')[0].split('_')[-1]))
             idx += 1
         output = processor.inference(imgs)
-        
+
         for j in range(len(shapes)):
             pred = processor.post_process(output[j].unsqueeze(0), shapes[j], conf_thres=conf_thres, iou_thres=iou_thres)
             if visualize and num_visualized < num_imgs_to_visualize:
@@ -113,10 +112,17 @@ def generate_results(processor, imgs_dir, jpgs, results_file, conf_thres, iou_th
                 w = float(p[2] - p[0])
                 h = float(p[3] - p[1])
                 s = float(p[4])
-                results.append({'image_id': image_ids[j],
-                                'category_id': coco91class[int(p[5])] if not non_coco else int(p[5]),
-                                'bbox': [round(x, 3) for x in [x, y, w, h]],
-                                'score': round(s, 5)})
+                results.append(
+                    {
+                        'image_id': image_ids[j],
+                        'category_id': int(p[5])
+                        if non_coco
+                        else coco91class[int(p[5])],
+                        'bbox': [round(x, 3) for x in [x, y, w, h]],
+                        'score': round(s, 5),
+                    }
+                )
+
 
                 if visualize and num_visualized < num_imgs_to_visualize:
                     cv2.rectangle(image, (int(x), int(y)), (int(x+w), int(y+h)), (255, 0, 0), 1)
@@ -141,10 +147,10 @@ def main():
 
         with open(args.model, 'wb') as f:
             f.write(engine.serialize())
-        print('Serialized the TensorRT engine to file: %s' % args.model)
+        print(f'Serialized the TensorRT engine to file: {args.model}')
 
     model_prefix = args.model.replace('.trt', '').split('/')[-1]
-    results_file = 'results_{}.json'.format(model_prefix)
+    results_file = f'results_{model_prefix}.json'
 
     # setup processor
     processor = Processor(model=args.model, scale_exact=args.scale_exact, return_int=args.letterbox_return_int, force_no_pad=args.force_no_pad)