diff --git a/Dockerfile.jetson b/Dockerfile.jetson index 8da6048..08261b0 100644 --- a/Dockerfile.jetson +++ b/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 diff --git a/setup.py b/setup.py index 5dcb156..b00d2be 100644 --- a/setup.py +++ b/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",