add ONNX deployment

Author: yjh0410

README.md

@@ -139,7 +139,7 @@ I have provided a bash file `train_ddp.sh` that enables DDP training. I hope som
 | Model         |   Backbone         | Scale | Epoch | AP<sup>val<br>0.5:0.95 | AP<sup>val<br>0.5 | FLOPs<br><sup>(G) | Params<br><sup>(M) | Weight |
 |---------------|--------------------|-------|-------|------------------------|-------------------|-------------------|--------------------|--------|
 | YOLOX-N       | CSPDarkNet-N       |  640  |  300  |         31.1           |       49.5        |   7.5             |   2.3              | [ckpt](https://github.com/yjh0410/PyTorch_YOLO_Tutorial/releases/download/yolo_tutorial_ckpt/yolox_n_coco.pth) |
-| YOLOX-S       | CSPDarkNet-S       |  640  |  300  |                        |                   |   27.1            |   9.0              |  |
+| YOLOX-S       | CSPDarkNet-S       |  640  |  300  |                        |                   |   26.8            |   8.9              |  |
 | YOLOX-M       | CSPDarkNet-M       |  640  |  300  |                        |                   |   74.3            |   25.4             |  |
 | YOLOX-L       | CSPDarkNet-L       |  640  |  300  |                        |                   |   155.4           |   54.2             |  |
 
@@ -360,3 +360,7 @@ python track.py --mode video \
 Results:
 
 ![image](./img_files/video_tracking_demo.gif)
+
+
+## Deployment
+1. [ONNX export and an ONNXRuntime](./deployment/ONNXRuntime/)

deployment/ONNXRuntime/README.md

@@ -0,0 +1,44 @@
+## YOLO ONNXRuntime
+
+
+### Convert Your Model to ONNX
+
+First, you should move to <FreeYOLO_HOME> by:
+```shell
+cd <FreeYOLO_HOME>
+cd tools/
+```
+Then, you can:
+
+1. Convert a standard YOLO model by:
+```shell
+python3 export_onnx.py -m yolov1 --weight ../weight/coco/yolov1/yolov1_coco.pth -nc 80 --img_size 640
+```
+
+Notes:
+* -n: specify a model name. The model name must be one of the [yolox-s,m,l,x and yolox-nano, yolox-tiny, yolov3]
+* -c: the model you have trained
+* -o: opset version, default 11. **However, if you will further convert your onnx model to [OpenVINO](https://github.com/Megvii-BaseDetection/YOLOX/demo/OpenVINO/), please specify the opset version to 10.**
+* --no-onnxsim: disable onnxsim
+* To customize an input shape for onnx model,  modify the following code in tools/export_onnx.py:
+
+    ```python
+    dummy_input = torch.randn(args.batch_size, 3, args.img_size, args.img_size)
+    ```
+
+### ONNXRuntime Demo
+
+Step1.
+```shell
+cd <YOLOX_HOME>/deployment/ONNXRuntime
+```
+
+Step2. 
+```shell
+python3 onnx_inference.py --weight ../../weights/onnx/11/yolov1.onnx -i ../test_image.jpg -s 0.3 --img_size 640
+```
+Notes:
+* --weight: your converted onnx model
+* -i: input_image
+* -s: score threshold for visualization.
+* --img_size: should be consistent with the shape you used for onnx convertion.

deployment/ONNXRuntime/onnx_inference.py

@@ -0,0 +1,87 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+# Copyright (c) Megvii, Inc. and its affiliates.
+
+import argparse
+import os
+
+import cv2
+import time
+import numpy as np
+import sys
+sys.path.append('../../')
+
+import onnxruntime
+from utils.misc import PreProcessor, PostProcessor
+from utils.vis_tools import visualize
+
+
+def make_parser():
+    parser = argparse.ArgumentParser("onnxruntime inference sample")
+    parser.add_argument("-m", "--model", type=str, default="../../weights/onnx/11/yolov1.onnx",
+                        help="Input your onnx model.")
+    parser.add_argument("-i", "--image_path", type=str, default='../test_image.jpg',
+                        help="Path to your input image.")
+    parser.add_argument("-o", "--output_dir", type=str, default='../../det_results/onnx/',
+                        help="Path to your output directory.")
+    parser.add_argument("-s", "--score_thr", type=float, default=0.35,
+                        help="Score threshould to filter the result.")
+    parser.add_argument("-size", "--img_size", type=int, default=640,
+                        help="Specify an input shape for inference.")
+    return parser
+
+
+if __name__ == '__main__':
+    args = make_parser().parse_args()
+
+    # class color for better visualization
+    np.random.seed(0)
+    class_colors = [(np.random.randint(255),
+                     np.random.randint(255),
+                     np.random.randint(255)) for _ in range(80)]
+
+    # preprocessor
+    prepocess = PreProcessor(img_size=args.img_size)
+
+    # postprocessor
+    postprocess = PostProcessor(num_classes=80, conf_thresh=args.score_thr, nms_thresh=0.5)
+
+    # read an image
+    input_shape = tuple([args.img_size, args.img_size])
+    origin_img = cv2.imread(args.image_path)
+
+    # preprocess
+    x, ratio = prepocess(origin_img)
+
+    t0 = time.time()
+    # inference
+    session = onnxruntime.InferenceSession(args.model)
+
+    ort_inputs = {session.get_inputs()[0].name: x[None, :, :, :]}
+    output = session.run(None, ort_inputs)
+    print("inference time: {:.1f} ms".format((time.time() - t0)*1000))
+
+    t0 = time.time()
+    # post process
+    bboxes, scores, labels = postprocess(output[0])
+    bboxes /= ratio
+    print("post-process time: {:.1f} ms".format((time.time() - t0)*1000))
+
+    # visualize detection
+    origin_img = visualize(
+        img=origin_img,
+        bboxes=bboxes,
+        scores=scores,
+        labels=labels,
+        vis_thresh=args.score_thr,
+        class_colors=class_colors
+        )
+
+    # show
+    cv2.imshow('onnx detection', origin_img)
+    cv2.waitKey(0)
+
+    # save results
+    os.makedirs(args.output_dir, exist_ok=True)
+    output_path = os.path.join(args.output_dir, os.path.basename(args.image_path))
+    cv2.imwrite(output_path, origin_img)

deployment/test_image.jpg

@@ -16,35 +16,36 @@ def build_model(args,
                 model_cfg,
                 device, 
                 num_classes=80, 
-                trainable=False):
+                trainable=False,
+                deploy=False):
     # YOLOv1    
     if args.model == 'yolov1':
         model, criterion = build_yolov1(
-            args, model_cfg, device, num_classes, trainable)
+            args, model_cfg, device, num_classes, trainable, deploy)
     # YOLOv2   
     elif args.model == 'yolov2':
         model, criterion = build_yolov2(
-            args, model_cfg, device, num_classes, trainable)
+            args, model_cfg, device, num_classes, trainable, deploy)
     # YOLOv3   
     elif args.model in ['yolov3', 'yolov3_t']:
         model, criterion = build_yolov3(
-            args, model_cfg, device, num_classes, trainable)
+            args, model_cfg, device, num_classes, trainable, deploy)
     # YOLOv4   
     elif args.model in ['yolov4', 'yolov4_t']:
         model, criterion = build_yolov4(
-            args, model_cfg, device, num_classes, trainable)
+            args, model_cfg, device, num_classes, trainable, deploy)
     # YOLOv5   
     elif args.model in ['yolov5_n', 'yolov5_s', 'yolov5_m', 'yolov5_l', 'yolov5_x']:
         model, criterion = build_yolov5(
-            args, model_cfg, device, num_classes, trainable)
+            args, model_cfg, device, num_classes, trainable, deploy)
     # YOLOv7
     elif args.model in ['yolov7_t', 'yolov7_l', 'yolov7_x']:
         model, criterion = build_yolov7(
-            args, model_cfg, device, num_classes, trainable)
+            args, model_cfg, device, num_classes, trainable, deploy)
     # YOLOX   
     elif args.model in ['yolox_n', 'yolox_s', 'yolox_m', 'yolox_l', 'yolox_x']:
         model, criterion = build_yolox(
-            args, model_cfg, device, num_classes, trainable)
+            args, model_cfg, device, num_classes, trainable, deploy)
 
     if trainable:
         # Load pretrained weight

models/detectors/__init__.py

@@ -9,7 +9,7 @@
 
 
 # build object detector
-def build_yolov1(args, cfg, device, num_classes=80, trainable=False):
+def build_yolov1(args, cfg, device, num_classes=80, trainable=False, deploy=False):
     print('==============================')
     print('Build {} ...'.format(args.model.upper()))
 
@@ -24,7 +24,8 @@ def build_yolov1(args, cfg, device, num_classes=80, trainable=False):
         num_classes = num_classes,
         conf_thresh = args.conf_thresh,
         nms_thresh = args.nms_thresh,
-        trainable = trainable
+        trainable = trainable,
+        deploy = deploy
         )
 
     # -------------- Initialize YOLO --------------

models/detectors/yolov1/build.py

@@ -18,7 +18,8 @@ def __init__(self,
                  num_classes=20,
                  conf_thresh=0.01,
                  nms_thresh=0.5,
-                 trainable=False):
+                 trainable=False,
+                 deploy=False):
         super(YOLOv1, self).__init__()
         # ------------------- Basic parameters -------------------
         self.cfg = cfg                                 # 模型配置文件
@@ -29,6 +30,7 @@ def __init__(self,
         self.conf_thresh = conf_thresh                 # 得分阈值
         self.nms_thresh = nms_thresh                   # NMS阈值
         self.stride = 32                               # 网络的最大步长
+        self.deploy = deploy
 
         # ------------------- Network Structure -------------------
         ## 主干网络
@@ -148,12 +150,18 @@ def inference(self, x):
         # 解算边界框, 并归一化边界框: [H*W, 4]
         bboxes = self.decode_boxes(reg_pred, fmp_size)
 
-        # 将预测放在cpu处理上，以便进行后处理
-        scores = scores.cpu().numpy()
-        bboxes = bboxes.cpu().numpy()
-        
-        # 后处理
-        bboxes, scores, labels = self.postprocess(bboxes, scores)
+        if self.deploy:
+            # [n_anchors_all, 4 + C]
+            outputs = torch.cat([bboxes, scores], dim=-1)
+
+            return outputs
+        else:
+            # 将预测放在cpu处理上，以便进行后处理
+            scores = scores.cpu().numpy()
+            bboxes = bboxes.cpu().numpy()
+            
+            # 后处理
+            bboxes, scores, labels = self.postprocess(bboxes, scores)
 
         return bboxes, scores, labels
 

models/detectors/yolov1/yolov1.py

@@ -9,7 +9,7 @@
 
 
 # build object detector
-def build_yolov2(args, cfg, device, num_classes=80, trainable=False):
+def build_yolov2(args, cfg, device, num_classes=80, trainable=False, deploy=False):
     print('==============================')
     print('Build {} ...'.format(args.model.upper()))
 
@@ -25,6 +25,7 @@ def build_yolov2(args, cfg, device, num_classes=80, trainable=False):
         conf_thresh=args.conf_thresh,
         nms_thresh=args.nms_thresh,
         topk=args.topk,
+        deploy=deploy
         )
 
     # -------------- Initialize YOLO --------------

models/detectors/yolov2/build.py

@@ -18,7 +18,8 @@ def __init__(self,
                  conf_thresh=0.01,
                  nms_thresh=0.5,
                  topk=100,
-                 trainable=False):
+                 trainable=False,
+                 deploy=False):
         super(YOLOv2, self).__init__()
         # ------------------- Basic parameters -------------------
         self.cfg = cfg                                 # 模型配置文件
@@ -29,8 +30,9 @@ def __init__(self,
         self.nms_thresh = nms_thresh                   # NMS阈值
         self.topk = topk                               # topk
         self.stride = 32                               # 网络的最大步长
+        self.deploy = deploy
         # ------------------- Anchor box -------------------
-        self.anchor_size = torch.as_tensor(cfg['anchor_size']).view(-1, 2) # [A, 2]
+        self.anchor_size = torch.as_tensor(cfg['anchor_size']).float().view(-1, 2) # [A, 2]
         self.num_anchors = self.anchor_size.shape[0]
 
         # ------------------- Network Structure -------------------
@@ -179,11 +181,19 @@ def inference(self, x):
         cls_pred = cls_pred[0]       # [H*W*A, NC]
         reg_pred = reg_pred[0]       # [H*W*A, 4]
 
-        # post process
-        bboxes, scores, labels = self.postprocess(
-            obj_pred, cls_pred, reg_pred, anchors)
+        if self.deploy:
+            scores = torch.sqrt(obj_pred.sigmoid() * cls_pred.sigmoid())
+            bboxes = self.decode_boxes(anchors, reg_pred)
+            # [n_anchors_all, 4 + C]
+            outputs = torch.cat([bboxes, scores], dim=-1)
 
-        return bboxes, scores, labels
+            return outputs
+        else:
+            # post process
+            bboxes, scores, labels = self.postprocess(
+                obj_pred, cls_pred, reg_pred, anchors)
+
+            return bboxes, scores, labels
 
 
     def forward(self, x):

models/detectors/yolov2/yolov2.py

@@ -9,7 +9,7 @@
 
 
 # build object detector
-def build_yolov3(args, cfg, device, num_classes=80, trainable=False):
+def build_yolov3(args, cfg, device, num_classes=80, trainable=False, deploy=False):
     print('==============================')
     print('Build {} ...'.format(args.model.upper()))
 
@@ -25,6 +25,7 @@ def build_yolov3(args, cfg, device, num_classes=80, trainable=False):
         conf_thresh=args.conf_thresh,
         nms_thresh=args.nms_thresh,
         topk=args.topk,
+        deploy = deploy
         )
 
     # -------------- Initialize YOLO --------------