From aae70703f7fbc126fde4971dac767b6fb4e65d02 Mon Sep 17 00:00:00 2001
From: Alexey <AlexeyAB@users.noreply.github.com>
Date: Sat, 23 Jul 2022 04:54:18 +0300
Subject: [PATCH] Update add_nms.py

---
 utils/add_nms.py | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/utils/add_nms.py b/utils/add_nms.py
index 8cfa239..927355e 100644
--- a/utils/add_nms.py
+++ b/utils/add_nms.py
@@ -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}")
\ No newline at end of file
+        LOGGER.info(f"Saved ONNX model to {output_path}")