FROM nvcr.io/nvidia/cuda:12.6.2-cudnn-devel-ubuntu24.04 AS base_devel
FROM nvcr.io/nvidia/cuda:12.6.2-cudnn-runtime-ubuntu24.04 AS base

LABEL maintainer="nan2088"
WORKDIR /app



# 1) base
ENV PIP_BREAK_SYSTEM_PACKAGES=1
RUN apt-get update -y && apt-get install -y vim ninja-build zip curl python3-pip wget && apt-get clean && apt-get autoclean && apt-get autoremove && rm -rf /root/.cache/*
RUN ln -s /usr/bin/python3 /usr/bin/python # && python -m pip install --upgrade pip 
RUN pip config set global.index-url https://mirrors.aliyun.com/pypi/simple/

RUN pip --no-cache-dir install torch==2.9.1 numpy torchvision==0.24.1  -i https://mirrors.aliyun.com/pypi/simple &&\
    rm -rf /root/.cache/
# --index-url https://download.pytorch.org/whl/cu126 

FROM base AS thirdparty
RUN pip config set global.index-url https://mirrors.aliyun.com/pypi/simple/

# 2) TensorRT/opencv Installation
RUN wget https://codeload.github.com/opencv/opencv/zip/refs/tags/4.5.4 -O /app/opencv-4.5.4.zip

RUN unzip opencv-4.5.4.zip  && rm -rf /app/opencv-4.5.4.zip
RUN pip  --no-cache-dir install cmake
RUN cd opencv-4.5.4/ && mkdir build && cd build && \
        cmake -D CMAKE_BUILD_TYPE=Release \
            -D BUILD_WITH_DEBUG_INFO=OFF \
            -D CMAKE_INSTALL_PREFIX=/opencv_install \
            -D INSTALL_C_EXAMPLES=OFF \
            -D INSTALL_PYTHON_EXAMPLES=OFF \
            -DENABLE_NEON=OFF  \
            -D WITH_TBB=ON \
            -DBUILD_TBB=ON  \
            -DBUILD_WEBP=OFF \
            -D BUILD_ITT=OFF -D WITH_IPP=ON  \
            -D WITH_V4L=OFF \
            -D WITH_QT=OFF \
            -D WITH_OPENGL=OFF \
            -D BUILD_opencv_dnn=OFF \
            -DBUILD_opencv_java=OFF \
            -DBUILD_opencv_python2=OFF \
            -DBUILD_opencv_python3=ON \
            -D BUILD_NEW_PYTHON_SUPPORT=ON \
            -D BUILD_PYTHON_SUPPORT=ON \
            -D PYTHON_DEFAULT_EXECUTABLE=/usr/bin/python3 \
            -DBUILD_opencv_java_bindings_generator=OFF \
            -DBUILD_opencv_python_bindings_generator=ON \
            -D BUILD_EXAMPLES=OFF \
            -D WITH_OPENEXR=OFF \
            -DWITH_JPEG=ON  \
            -DBUILD_JPEG=ON\
            -D BUILD_JPEG_TURBO_DISABLE=OFF \
            -D BUILD_DOCS=OFF \
            -D BUILD_PERF_TESTS=OFF \
            -D BUILD_TESTS=OFF \
            -D BUILD_opencv_apps=OFF \
            -D BUILD_opencv_calib3d=OFF \
            -D BUILD_opencv_contrib=OFF \
            -D BUILD_opencv_features2d=OFF \
            -D BUILD_opencv_flann=OFF \
            -DBUILD_opencv_gapi=OFF \
            -D WITH_CUDA=OFF \
            -D WITH_CUDNN=OFF \
            -D OPENCV_DNN_CUDA=OFF \
            -D ENABLE_FAST_MATH=1 \
            -D WITH_CUBLAS=0 \
            -D BUILD_opencv_gpu=OFF \
            -D BUILD_opencv_ml=OFF \
            -D BUILD_opencv_nonfree=OFF \
            -D BUILD_opencv_objdetect=OFF \
            -D BUILD_opencv_photo=OFF \
            -D BUILD_opencv_stitching=OFF \
            -D BUILD_opencv_superres=OFF \
            -D BUILD_opencv_ts=OFF \
            -D BUILD_opencv_video=OFF \
            -D BUILD_videoio_plugins=OFF \
            -D BUILD_opencv_videostab=OFF \
            -DBUILD_EXAMPLES=OFF \
            -DBUILD_opencv_calib3d=OFF \
            -DBUILD_opencv_features2d=OFF\
            -DBUILD_opencv_flann=OFF\
            -DBUILD_opencv_ml=OFF\
            -DBUILD_opencv_videoio=OFF\
                .. && make -j4 && make install



# Please manually download the TensorRT installation package from https://github.com/NVIDIA/TensorRT/releases.
# Check if the file exists in the local directory
RUN if [ ! -f ./TensorRT-10.14.1.48.Linux.x86_64-gnu.cuda-12.9.tar.gz ]; then \
    # If not, download the file to the /app directory
    wget -P /app/ https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.14.1/tars/TensorRT-10.14.1.48.Linux.x86_64-gnu.cuda-12.9.tar.gz; \
else \
    # If the file exists, copy it to the /app directory
    cp ./TensorRT-10.14.1.48.Linux.x86_64-gnu.cuda-12.9.tar.gz /app/; \
fi

RUN tar -xzvf /app/TensorRT-10.14.1.48.Linux.x86_64-gnu.cuda-12.9.tar.gz -C /app/; 

# TensorRT-10.14.1.48/lib/libnvinfer_builder_resource_sm* :
# SM75: RTX 20 系列 (Turing)
# SM80/SM86: RTX 30 系列 (Ampere)
# SM89: RTX 40 系列 (Ada Lovelace)
# SM90: H100/Hopper
# SM100: Blackwell（下一代）
# SM120: RTX 50 系列（Blackwell 游戏卡）

RUN rm -rf /app/TensorRT-10.14.1.48/data && \
     rm /app/TensorRT-10.14.1.48/lib/*.a && \
     rm /app/TensorRT-10.14.1.48/lib/libnvinfer_builder_resource_ptx* && \
     rm /app/TensorRT-10.14.1.48/lib/libnvinfer_builder_resource_win* 

FROM  base_devel AS base_devel_removed
RUN rm /usr/local/cuda-12/targets/x86_64-linux/lib/*.a

RUN find /usr/local/cuda-12/targets/x86_64-linux/lib/ -type f \( -name 'lib*.a' ! -name 'libcudart_static.a' ! -name 'libcudadevrt.a' \) -delete

RUN rm /usr/local/cuda-12/targets/x86_64-linux/lib/libcusparse.so*
RUN rm /usr/local/cuda-12/targets/x86_64-linux/lib/libcusolverMg.so*
RUN rm /usr/local/cuda-12/targets/x86_64-linux/lib/libcusolver.so*
RUN rm /usr/local/cuda-12/targets/x86_64-linux/lib/libcufft.so*

# 3) TensorRT opencv cuda Installation
FROM base
COPY --from=thirdparty /app/TensorRT-10.14.1.48/targets/x86_64-linux-gnu/lib/ /usr/lib/
COPY --from=thirdparty /app/TensorRT-10.14.1.48/include/ /usr/include/
COPY --from=thirdparty /opencv_install/ /usr/local/
COPY --from=thirdparty /app/TensorRT-10.14.1.48/python/tensorrt-*-cp312-none-linux_x86_64.whl /app/
COPY --from=thirdparty /app/TensorRT-10.14.1.48/bin/trtexec /usr/bin/

COPY --from=base_devel_removed  /usr/local/cuda-12/include/ /usr/local/cuda-12/include/
COPY --from=base_devel_removed /usr/local/cuda-12/bin/ /usr/local/cuda-12/bin
COPY --from=base_devel_removed /usr/local/cuda-12/nvvm /usr/local/cuda-12/nvvm
COPY --from=base_devel_removed /usr/local/cuda-12/targets/x86_64-linux/lib /usr/local/cuda-12/lib64
RUN pip install /app/tensorrt-*-cp312-none-linux_x86_64.whl && rm /app/tensorrt-*-cp3*-none-linux_x86_64.whl

# git is needed for building PPL.CV
RUN apt-get update -y && apt-get install git -y && pip --no-cache-dir install cmake 
WORKDIR /app

# export BUILDKIT_STEP_LOG_MAX_SIZE=-1
# docker build -f DockerfileCuda12 -t torchpipe:base_trt-10.14 .