mirrors
/
yolov7
mirror of https://github.com/WongKinYiu/yolov7.git
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

Update add_nms.py

pull/314/head
Alexey 2022-07-23 04:54:18 +03:00 committed by GitHub
parent a5b853c3fa
commit aae70703f7
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 6 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
utils/add_nms.py
View File

@ -1,7 +1,11 @@
import numpy as np
import onnx
from onnx import shape_inference
import onnx_graphsurgeon as gs
try:
import onnx_graphsurgeon as gs
except Exception as e:
print('Import onnx_graphsurgeon failure: %s' % e)
import logging
LOGGER = logging.getLogger(__name__)
@ -148,4 +152,4 @@ class RegisterNMS(object):
self.graph.cleanup().toposort()
model = gs.export_onnx(self.graph)
onnx.save(model, output_path)
LOGGER.info(f"Saved ONNX model to {output_path}")
LOGGER.info(f"Saved ONNX model to {output_path}")