From 5faa71461b875195602e3b569dc28c0684ade0ad Mon Sep 17 00:00:00 2001
From: Nicolas Mowen <nickmowen213@gmail.com>
Date: Sat, 7 Jun 2025 10:57:19 -0600
Subject: [PATCH] Combine base and arm trt detectors

---
 docker/tensorrt/Dockerfile.arm64 | 54 +++++++++++++++++++++++++++++-
 docker/tensorrt/Dockerfile.base  | 57 --------------------------------
 2 files changed, 53 insertions(+), 58 deletions(-)
 delete mode 100644 docker/tensorrt/Dockerfile.base

diff --git a/docker/tensorrt/Dockerfile.arm64 b/docker/tensorrt/Dockerfile.arm64
index 6e2e0280f..5b139acc6 100644
--- a/docker/tensorrt/Dockerfile.arm64
+++ b/docker/tensorrt/Dockerfile.arm64
@@ -2,8 +2,60 @@
 
 # https://askubuntu.com/questions/972516/debian-frontend-environment-variable
 ARG DEBIAN_FRONTEND=noninteractive
-
 ARG BASE_IMAGE
+ARG TRT_BASE=nvcr.io/nvidia/tensorrt:23.12-py3
+
+# Build TensorRT-specific library
+FROM ${TRT_BASE} AS trt-deps
+
+ARG TARGETARCH
+ARG COMPUTE_LEVEL
+
+RUN apt-get update \
+    && apt-get install -y git build-essential cuda-nvcc-* cuda-nvtx-* libnvinfer-dev libnvinfer-plugin-dev libnvparsers-dev libnvonnxparsers-dev \
+    && rm -rf /var/lib/apt/lists/*
+RUN --mount=type=bind,source=docker/tensorrt/detector/tensorrt_libyolo.sh,target=/tensorrt_libyolo.sh \
+    /tensorrt_libyolo.sh
+
+# COPY required individual CUDA deps
+RUN mkdir -p /usr/local/cuda-deps
+RUN if [ "$TARGETARCH" = "amd64" ]; then \
+    cp /usr/local/cuda-12.3/targets/x86_64-linux/lib/libcurand.so.* /usr/local/cuda-deps/ && \
+    cp /usr/local/cuda-12.3/targets/x86_64-linux/lib/libnvrtc.so.* /usr/local/cuda-deps/ && \
+    cd /usr/local/cuda-deps/ && \
+    for lib in libnvrtc.so.*; do \
+    if [[ "$lib" =~ libnvrtc.so\.([0-9]+\.[0-9]+\.[0-9]+) ]]; then \
+    version="${BASH_REMATCH[1]}"; \
+    ln -sf "libnvrtc.so.$version" libnvrtc.so; \
+    fi; \
+    done && \
+    for lib in libcurand.so.*; do \
+    if [[ "$lib" =~ libcurand.so\.([0-9]+\.[0-9]+\.[0-9]+\.[0-9]+) ]]; then \
+    version="${BASH_REMATCH[1]}"; \
+    ln -sf "libcurand.so.$version" libcurand.so; \
+    fi; \
+    done; \
+    fi
+
+# Frigate w/ TensorRT Support as separate image
+FROM deps AS tensorrt-base
+
+#Disable S6 Global timeout
+ENV S6_CMD_WAIT_FOR_SERVICES_MAXTIME=0
+
+# COPY TensorRT Model Generation Deps
+COPY --from=trt-deps /usr/local/lib/libyolo_layer.so /usr/local/lib/libyolo_layer.so
+COPY --from=trt-deps /usr/local/src/tensorrt_demos /usr/local/src/tensorrt_demos
+
+# COPY Individual CUDA deps folder
+COPY --from=trt-deps /usr/local/cuda-deps /usr/local/cuda
+
+COPY docker/tensorrt/detector/rootfs/ /
+ENV YOLO_MODELS=""
+
+HEALTHCHECK --start-period=600s --start-interval=5s --interval=15s --timeout=5s --retries=3 \
+    CMD curl --fail --silent --show-error http://127.0.0.1:5000/api/version || exit 1
+
 FROM ${BASE_IMAGE} AS build-wheels
 ARG DEBIAN_FRONTEND
 
diff --git a/docker/tensorrt/Dockerfile.base b/docker/tensorrt/Dockerfile.base
deleted file mode 100644
index 79a7f3c98..000000000
--- a/docker/tensorrt/Dockerfile.base
+++ /dev/null
@@ -1,57 +0,0 @@
-# syntax=docker/dockerfile:1.6
-
-# https://askubuntu.com/questions/972516/debian-frontend-environment-variable
-ARG DEBIAN_FRONTEND=noninteractive
-
-ARG TRT_BASE=nvcr.io/nvidia/tensorrt:23.12-py3
-
-# Build TensorRT-specific library
-FROM ${TRT_BASE} AS trt-deps
-
-ARG TARGETARCH
-ARG COMPUTE_LEVEL
-
-RUN apt-get update \
-    && apt-get install -y git build-essential cuda-nvcc-* cuda-nvtx-* libnvinfer-dev libnvinfer-plugin-dev libnvparsers-dev libnvonnxparsers-dev \
-    && rm -rf /var/lib/apt/lists/*
-RUN --mount=type=bind,source=docker/tensorrt/detector/tensorrt_libyolo.sh,target=/tensorrt_libyolo.sh \
-    /tensorrt_libyolo.sh
-
-# COPY required individual CUDA deps
-RUN mkdir -p /usr/local/cuda-deps
-RUN if [ "$TARGETARCH" = "amd64" ]; then \
-      cp /usr/local/cuda-12.3/targets/x86_64-linux/lib/libcurand.so.* /usr/local/cuda-deps/ && \
-      cp /usr/local/cuda-12.3/targets/x86_64-linux/lib/libnvrtc.so.* /usr/local/cuda-deps/ && \
-      cd /usr/local/cuda-deps/ && \
-      for lib in libnvrtc.so.*; do \
-        if [[ "$lib" =~ libnvrtc.so\.([0-9]+\.[0-9]+\.[0-9]+) ]]; then \
-          version="${BASH_REMATCH[1]}"; \
-          ln -sf "libnvrtc.so.$version" libnvrtc.so; \
-        fi; \
-      done && \
-      for lib in libcurand.so.*; do \
-        if [[ "$lib" =~ libcurand.so\.([0-9]+\.[0-9]+\.[0-9]+\.[0-9]+) ]]; then \
-          version="${BASH_REMATCH[1]}"; \
-          ln -sf "libcurand.so.$version" libcurand.so; \
-        fi; \
-      done; \
-    fi
-
-# Frigate w/ TensorRT Support as separate image
-FROM deps AS tensorrt-base
-
-#Disable S6 Global timeout
-ENV S6_CMD_WAIT_FOR_SERVICES_MAXTIME=0
-
-# COPY TensorRT Model Generation Deps
-COPY --from=trt-deps /usr/local/lib/libyolo_layer.so /usr/local/lib/libyolo_layer.so
-COPY --from=trt-deps /usr/local/src/tensorrt_demos /usr/local/src/tensorrt_demos
-
-# COPY Individual CUDA deps folder
-COPY --from=trt-deps /usr/local/cuda-deps /usr/local/cuda
-
-COPY docker/tensorrt/detector/rootfs/ /
-ENV YOLO_MODELS=""
-
-HEALTHCHECK --start-period=600s --start-interval=5s --interval=15s --timeout=5s --retries=3 \
-    CMD curl --fail --silent --show-error http://127.0.0.1:5000/api/version || exit 1