mirror of https://github.com/WongKinYiu/yolov7.git
minor fix
parent
afdc86f519
commit
e5928b095a
|
@ -1,3 +0,0 @@
|
|||
# 默认忽略的文件
|
||||
/shelf/
|
||||
/workspace.xml
|
|
@ -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>
|
|
@ -1,6 +0,0 @@
|
|||
<component name="InspectionProjectProfileManager">
|
||||
<settings>
|
||||
<option name="USE_PROJECT_PROFILE" value="false" />
|
||||
<version value="1.0" />
|
||||
</settings>
|
||||
</component>
|
|
@ -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>
|
|
@ -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>
|
|
@ -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>
|
|
@ -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>
|
13
README.md
13
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
|
||||
|
|
34
export.py
34
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
|
||||
|
|
Loading…
Reference in New Issue