From 46787b25e26458398501d181975c7fb039958705 Mon Sep 17 00:00:00 2001
From: Nate Meyer <nate.devel@gmail.com>
Date: Tue, 20 Jun 2023 00:16:58 -0400
Subject: [PATCH] Build trt libyolo_layer.so in container

---
 Dockerfile                 |  9 +++++++++
 docker/tensorrt_libyolo.sh | 10 ++++++++++
 2 files changed, 19 insertions(+)
 create mode 100755 docker/tensorrt_libyolo.sh

diff --git a/Dockerfile b/Dockerfile
index 0c244f901..cc872ccf4 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -253,8 +253,16 @@ FROM deps AS frigate
 WORKDIR /opt/frigate/
 COPY --from=rootfs / /
 
+# Build TensorRT-specific library
+FROM nvcr.io/nvidia/tensorrt:23.05-py3 AS trt-deps
+
+RUN --mount=type=bind,source=docker/tensorrt_libyolo.sh,target=/tensorrt_libyolo.sh \
+    /tensorrt_libyolo.sh
+
 # Frigate w/ TensorRT Support as separate image
 FROM frigate AS frigate-tensorrt
+
+COPY --from=trt-deps /usr/local/lib/libyolo_layer.so /usr/local/lib/libyolo_layer.so
 RUN --mount=type=bind,from=trt-wheels,source=/trt-wheels,target=/deps/trt-wheels \
     pip3 install -U /deps/trt-wheels/*.whl && \
     ln -s libnvrtc.so.11.2 /usr/local/lib/python3.9/dist-packages/nvidia/cuda_nvrtc/lib/libnvrtc.so && \
@@ -263,5 +271,6 @@ RUN --mount=type=bind,from=trt-wheels,source=/trt-wheels,target=/deps/trt-wheels
 # Dev Container w/ TRT
 FROM devcontainer AS devcontainer-trt
 
+COPY --from=trt-deps /usr/local/lib/libyolo_layer.so /usr/local/lib/libyolo_layer.so
 RUN --mount=type=bind,from=trt-wheels,source=/trt-wheels,target=/deps/trt-wheels \
     pip3 install -U /deps/trt-wheels/*.whl
diff --git a/docker/tensorrt_libyolo.sh b/docker/tensorrt_libyolo.sh
new file mode 100755
index 000000000..dc7a39ca1
--- /dev/null
+++ b/docker/tensorrt_libyolo.sh
@@ -0,0 +1,10 @@
+#!/bin/bash
+
+set -euxo pipefail
+
+# Clone tensorrt_demos repo
+git clone --depth 1 https://github.com/yeahme49/tensorrt_demos.git
+
+# Build libyolo
+cd ./tensorrt_demos/plugins && make all
+cp libyolo_layer.so /usr/local/lib/libyolo_layer.so
\ No newline at end of file