trt5 legacy api

1 files changed, 13 insertions, 4 deletions

diff --git a/run_engine.py b/run_engine.py
index 700494d..6ac2ec9 100644
--- a/run_engine.py
+++ b/run_engine.py
@@ -2,7 +2,11 @@ import os
 import sys
 from random import randint
 import numpy as np
-import tensorrt
+
+if 'TensorRT-5.0' in os.getenv('LD_LIBRARY_PATH'):
+    import tensorrt.legacy
+else:
+    import tensorrt
 
 try:
     from PIL import Image
@@ -19,9 +23,14 @@ pip(3) install tensorrt[examples]
     exit(1)
 
 try:
-    import tensorrt as trt
-    from tensorrt.parsers import caffeparser
-    from tensorrt.parsers import onnxparser    
+    if 'TensorRT-5.0' in os.getenv('LD_LIBRARY_PATH'):
+        import tensorrt.legacy as trt
+        from tensorrt.legacy.parsers import caffeparser
+        from tensorrt.legacy.parsers import onnxparser
+    else:
+        import tensorrt as trt
+        from tensorrt.parsers import caffeparser
+        from tensorrt.parsers import onnxparser
 except ImportError as err:
     sys.stderr.write("""ERROR: failed to import module ({})
 Please make sure you have the TensorRT Library installed