Added biuld TensorRT flag

setup.py

@@ -1,6 +1,7 @@
 from setuptools import setup, Extension
 from Cython.Build import cythonize
 import numpy as np
+import os
 
 SRC = "src"
 np_inc = [np.get_include(), SRC]
@@ -18,16 +19,22 @@ extensions = [
     Extension('inference', [f'{SRC}/inference.pyx'], include_dirs=np_inc),
 ]
 
-try:
-    import tensorrt  # pyright: ignore[reportMissingImports]
+build_tensorrt = os.environ.get("BUILD_TENSORRT_EXTENSIONS", "").lower() in ("1", "true", "yes")
+
+if not build_tensorrt:
+    try:
+        import tensorrt  # pyright: ignore[reportMissingImports]
+        build_tensorrt = True
+    except ImportError:
+        build_tensorrt = False
+
+if build_tensorrt:
     extensions.append(
         Extension('engines.tensorrt_engine', [f'{SRC}/engines/tensorrt_engine.pyx'], include_dirs=np_inc)
     )
     extensions.append(
         Extension('engines.jetson_tensorrt_engine', [f'{SRC}/engines/jetson_tensorrt_engine.pyx'], include_dirs=np_inc)
     )
-except ImportError:
-    pass
 
 setup(
     name="azaion.detections",