minor fix

pull/280/head
AlexeyAB84 2022-07-22 17:04:38 +03:00
parent afdc86f519
commit e5928b095a
9 changed files with 29 additions and 103 deletions

3
.idea/.gitignore vendored
View File

@ -1,3 +0,0 @@
# 默认忽略的文件
/shelf/
/workspace.xml

View File

@ -1,46 +0,0 @@
<component name="InspectionProjectProfileManager">
<profile version="1.0">
<option name="myName" value="Project Default" />
<inspection_tool class="PyPackageRequirementsInspection" enabled="true" level="WARNING" enabled_by_default="true">
<option name="ignoredPackages">
<value>
<list size="18">
<item index="0" class="java.lang.String" itemvalue="onnxruntime" />
<item index="1" class="java.lang.String" itemvalue="onnx-simplifier" />
<item index="2" class="java.lang.String" itemvalue="scipy" />
<item index="3" class="java.lang.String" itemvalue="thop" />
<item index="4" class="java.lang.String" itemvalue="opencv-python" />
<item index="5" class="java.lang.String" itemvalue="torch" />
<item index="6" class="java.lang.String" itemvalue="numpy" />
<item index="7" class="java.lang.String" itemvalue="torchvision" />
<item index="8" class="java.lang.String" itemvalue="tqdm" />
<item index="9" class="java.lang.String" itemvalue="pandas" />
<item index="10" class="java.lang.String" itemvalue="tensorboard" />
<item index="11" class="java.lang.String" itemvalue="seaborn" />
<item index="12" class="java.lang.String" itemvalue="matplotlib" />
<item index="13" class="java.lang.String" itemvalue="Cython" />
<item index="14" class="java.lang.String" itemvalue="pycocotools" />
<item index="15" class="java.lang.String" itemvalue="h5py" />
<item index="16" class="java.lang.String" itemvalue="opencv_python" />
<item index="17" class="java.lang.String" itemvalue="Pillow" />
</list>
</value>
</option>
</inspection_tool>
<inspection_tool class="PyPep8NamingInspection" enabled="true" level="WEAK WARNING" enabled_by_default="true">
<option name="ignoredErrors">
<list>
<option value="N806" />
<option value="N801" />
</list>
</option>
</inspection_tool>
<inspection_tool class="PyUnresolvedReferencesInspection" enabled="true" level="WARNING" enabled_by_default="true">
<option name="ignoredIdentifiers">
<list>
<option value="tkinter.*" />
</list>
</option>
</inspection_tool>
</profile>
</component>

View File

@ -1,6 +0,0 @@
<component name="InspectionProjectProfileManager">
<settings>
<option name="USE_PROJECT_PROFILE" value="false" />
<version value="1.0" />
</settings>
</component>

View File

@ -1,4 +0,0 @@
<?xml version="1.0" encoding="UTF-8"?>
<project version="4">
<component name="ProjectRootManager" version="2" project-jdk-name="Python 3.8" project-jdk-type="Python SDK" />
</project>

View File

@ -1,8 +0,0 @@
<?xml version="1.0" encoding="UTF-8"?>
<project version="4">
<component name="ProjectModuleManager">
<modules>
<module fileurl="file://$PROJECT_DIR$/.idea/yolov7.iml" filepath="$PROJECT_DIR$/.idea/yolov7.iml" />
</modules>
</component>
</project>

View File

@ -1,6 +0,0 @@
<?xml version="1.0" encoding="UTF-8"?>
<project version="4">
<component name="VcsDirectoryMappings">
<mapping directory="$PROJECT_DIR$" vcs="Git" />
</component>
</project>

View File

@ -1,12 +0,0 @@
<?xml version="1.0" encoding="UTF-8"?>
<module type="PYTHON_MODULE" version="4">
<component name="NewModuleRootManager">
<content url="file://$MODULE_DIR$" />
<orderEntry type="inheritedJdk" />
<orderEntry type="sourceFolder" forTests="false" />
</component>
<component name="PyDocumentationSettings">
<option name="format" value="PLAIN" />
<option name="myDocStringFormat" value="Plain" />
</component>
</module>

View File

@ -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

View File

@ -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