diff --git a/.idea/.gitignore b/.idea/.gitignore
deleted file mode 100644
index 359bb530..00000000
--- a/.idea/.gitignore
+++ /dev/null
@@ -1,3 +0,0 @@
-# 默认忽略的文件
-/shelf/
-/workspace.xml
diff --git a/.idea/inspectionProfiles/Project_Default.xml b/.idea/inspectionProfiles/Project_Default.xml
deleted file mode 100644
index 28c863b6..00000000
--- a/.idea/inspectionProfiles/Project_Default.xml
+++ /dev/null
@@ -1,46 +0,0 @@
-
-
-
-
-
-
-
-
-
-
-
-
-
-
\ No newline at end of file
diff --git a/.idea/inspectionProfiles/profiles_settings.xml b/.idea/inspectionProfiles/profiles_settings.xml
deleted file mode 100644
index 105ce2da..00000000
--- a/.idea/inspectionProfiles/profiles_settings.xml
+++ /dev/null
@@ -1,6 +0,0 @@
-
-
-
-
-
-
\ No newline at end of file
diff --git a/.idea/misc.xml b/.idea/misc.xml
deleted file mode 100644
index d1e22ecb..00000000
--- a/.idea/misc.xml
+++ /dev/null
@@ -1,4 +0,0 @@
-
-
-
-
\ No newline at end of file
diff --git a/.idea/modules.xml b/.idea/modules.xml
deleted file mode 100644
index ecffae6d..00000000
--- a/.idea/modules.xml
+++ /dev/null
@@ -1,8 +0,0 @@
-
-
-
-
-
-
-
-
\ No newline at end of file
diff --git a/.idea/vcs.xml b/.idea/vcs.xml
deleted file mode 100644
index 94a25f7f..00000000
--- a/.idea/vcs.xml
+++ /dev/null
@@ -1,6 +0,0 @@
-
-
-
-
-
-
\ No newline at end of file
diff --git a/.idea/yolov7.iml b/.idea/yolov7.iml
deleted file mode 100644
index 8b8c3954..00000000
--- a/.idea/yolov7.iml
+++ /dev/null
@@ -1,12 +0,0 @@
-
-
-
-
-
-
-
-
-
-
-
-
\ No newline at end of file
diff --git a/README.md b/README.md
index 1892b1b0..38113bf9 100644
--- a/README.md
+++ b/README.md
@@ -153,9 +153,18 @@ python detect.py --weights yolov7.pt --conf 0.25 --img-size 640 --source inferen
## Export
-Use the args `--include-nms` can to export end to end onnx model which include the `EfficientNMS`.
+Tested with: Python 3.7.13 and Pytorch 1.12.0+cu113
+Pytorch to ONNX, use `--include-nms` flag for the end-to-end ONNX model with `EfficientNMS`.
```shell
-python models/export.py --weights yolov7.pt --grid --include-nms
+wget https://github.com/WongKinYiu/yolov7/releases/download/v0.1/yolov7-tiny.pt
+python export.py --weights yolov7-tiny.pt --grid --include-nms
+```
+
+ONNX to TensorRT
+```shell
+git clone https://github.com/Linaom1214/tensorrt-python.git
+cd tensorrt-python
+python export.py -o yolov7-tiny.onnx -e yolov7-tiny-nms.trt -p fp16
```
## Citation
diff --git a/export.py b/export.py
index fd4975f6..f9a6a89d 100644
--- a/export.py
+++ b/export.py
@@ -79,23 +79,17 @@ if __name__ == '__main__':
dynamic_axes={'images': {0: 'batch', 2: 'height', 3: 'width'}, # size(1,3,640,640)
'output': {0: 'batch', 2: 'y', 3: 'x'}} if opt.dynamic else None)
- if opt.include_nms:
- print('Registering NMS plugin...')
- mo = RegisterNMS(f)
- mo.register_nms()
- mo.save(f)
- else:
- # Checks
- onnx_model = onnx.load(f) # load onnx model
- onnx.checker.check_model(onnx_model) # check onnx model
- # print(onnx.helper.printable_graph(onnx_model.graph)) # print a human readable model
+ # Checks
+ onnx_model = onnx.load(f) # load onnx model
+ onnx.checker.check_model(onnx_model) # check onnx model
+ # print(onnx.helper.printable_graph(onnx_model.graph)) # print a human readable model
- # # Metadata
- # d = {'stride': int(max(model.stride))}
- # for k, v in d.items():
- # meta = onnx_model.metadata_props.add()
- # meta.key, meta.value = k, str(v)
- # onnx.save(onnx_model, f)
+ # # Metadata
+ # d = {'stride': int(max(model.stride))}
+ # for k, v in d.items():
+ # meta = onnx_model.metadata_props.add()
+ # meta.key, meta.value = k, str(v)
+ # onnx.save(onnx_model, f)
if opt.simplify:
try:
@@ -104,9 +98,17 @@ if __name__ == '__main__':
print('\nStarting to simplify ONNX...')
onnx_model, check = onnxsim.simplify(onnx_model)
assert check, 'assert check failed'
+ onnx.save(onnx_model, f)
except Exception as e:
print(f'Simplifier failure: {e}')
print('ONNX export success, saved as %s' % f)
+
+ if opt.include_nms:
+ print('Registering NMS plugin for ONNX...')
+ mo = RegisterNMS(f)
+ mo.register_nms()
+ mo.save(f)
+
except Exception as e:
print('ONNX export failure: %s' % e)
# CoreML export