Added biuld TensorRT flag

Dockerfile.jetson

@@ -13,7 +13,7 @@ WORKDIR /app
 COPY requirements.txt requirements-jetson.txt ./
 RUN pip3 install --no-cache-dir -r requirements-jetson.txt
 COPY . .
-RUN python3 setup.py build_ext --inplace
+RUN BUILD_TENSORRT_EXTENSIONS=1 python3 setup.py build_ext --inplace
 ENV PYTHONPATH=/app/src
 RUN adduser --disabled-password --no-create-home --gecos "" appuser \
     && chown -R appuser /app

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),
 ]
 
+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",