Add `--include torchscript onnx coreml` argument (#3137)

models/export.py

-"""Exports a YOLOv5 *.pt model to ONNX and TorchScript formats
+"""Exports a YOLOv5 *.pt model to TorchScript, ONNX, CoreML formats
 
 Usage:
-    $ export PYTHONPATH="$PWD" && python models/export.py --weights yolov5s.pt --img 640 --batch 1
+    $ python path/to/models/export.py --weights yolov5s.pt --img 640 --batch 1
 """
 
 import argparse
@@ -27,6 +27,7 @@ if __name__ == '__main__':
     parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='image size')  # height, width
     parser.add_argument('--batch-size', type=int, default=1, help='batch size')
     parser.add_argument('--device', default='cpu', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
+    parser.add_argument('--include', nargs='+', default=['torchscript', 'onnx', 'coreml'], help='include formats')
     parser.add_argument('--half', action='store_true', help='FP16 half-precision export')
     parser.add_argument('--inplace', action='store_true', help='set YOLOv5 Detect() inplace=True')
     parser.add_argument('--train', action='store_true', help='model.train() mode')
@@ -35,6 +36,7 @@ if __name__ == '__main__':
     parser.add_argument('--simplify', action='store_true', help='simplify ONNX model')  # ONNX-only
     opt = parser.parse_args()
     opt.img_size *= 2 if len(opt.img_size) == 1 else 1  # expand
+    opt.include = [x.lower() for x in opt.include]
     print(opt)
     set_logging()
     t = time.time()
@@ -47,7 +49,7 @@ if __name__ == '__main__':
     # Checks
     gs = int(max(model.stride))  # grid size (max stride)
     opt.img_size = [check_img_size(x, gs) for x in opt.img_size]  # verify img_size are gs-multiples
-    assert not (opt.device.lower() == "cpu" and opt.half), '--half only compatible with GPU export, i.e. use --device 0'
+    assert not (opt.device.lower() == 'cpu' and opt.half), '--half only compatible with GPU export, i.e. use --device 0'
 
     # Input
     img = torch.zeros(opt.batch_size, 3, *opt.img_size).to(device)  # image size(1,3,320,192) iDetection
@@ -74,6 +76,7 @@ if __name__ == '__main__':
     print(f"\n{colorstr('PyTorch:')} starting from {opt.weights} ({file_size(opt.weights):.1f} MB)")
 
     # TorchScript export -----------------------------------------------------------------------------------------------
+    if 'torchscript' in opt.include or 'coreml' in opt.include:
         prefix = colorstr('TorchScript:')
         try:
             print(f'\n{prefix} starting export with torch {torch.__version__}...')
@@ -85,6 +88,7 @@ if __name__ == '__main__':
             print(f'{prefix} export failure: {e}')
 
     # ONNX export ------------------------------------------------------------------------------------------------------
+    if 'onnx' in opt.include:
         prefix = colorstr('ONNX:')
         try:
             import onnx
@@ -107,7 +111,8 @@ if __name__ == '__main__':
                     import onnxsim
 
                     print(f'{prefix} simplifying with onnx-simplifier {onnxsim.__version__}...')
-                model_onnx, check = onnxsim.simplify(model_onnx,
+                    model_onnx, check = onnxsim.simplify(
+                        model_onnx,
                         dynamic_input_shape=opt.dynamic,
                         input_shapes={'images': list(img.shape)} if opt.dynamic else None)
                     assert check, 'assert check failed'
@@ -119,12 +124,13 @@ if __name__ == '__main__':
             print(f'{prefix} export failure: {e}')
 
     # CoreML export ----------------------------------------------------------------------------------------------------
+    if 'coreml' in opt.include:
         prefix = colorstr('CoreML:')
         try:
             import coremltools as ct
 
             print(f'{prefix} starting export with coremltools {ct.__version__}...')
-        model = ct.convert(ts, inputs=[ct.ImageType(name='image', shape=img.shape, scale=1 / 255.0, bias=[0, 0, 0])])
+            model = ct.convert(ts, inputs=[ct.ImageType('image', shape=img.shape, scale=1 / 255.0, bias=[0, 0, 0])])
             f = opt.weights.replace('.pt', '.mlmodel')  # filename
             model.save(f)
             print(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')