forked from blakeblackshear/frigate
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Nvidia TensorRT detector (blakeblackshear#4718)
* Initial WIP dockerfile and scripts to add tensorrt support * Add tensorRT detector * WIP attempt to install TensorRT 8.5 * Updates to detector for cuda python library * TensorRT Cuda library rework WIP Does not run * Fixes from rebase to detector factory * Fix parsing output memory pointer * Handle TensorRT logs with the python logger * Use non-async interface and convert input data to float32. Detection runs without error. * Make TensorRT a separate build from the base Frigate image. * Add script and documentation for generating TRT Models * Add support for TensorRT devcontainer * Add labelmap to trt model script and docs. Cleanup of old scripts. * Update detect to normalize input tensor using model input type * Add config for selecting GPU. Fix Async inference. Update documentation. * Update some CUDA libraries to clean up version warning * Add CI stage to build TensorRT tag * Add note in docs for image tag and model support
- Loading branch information
Showing
9 changed files
with
515 additions
and
16 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,37 @@ | ||
#!/bin/bash | ||
|
||
set -euxo pipefail | ||
|
||
CUDA_HOME=/usr/local/cuda | ||
LD_LIBRARY_PATH=${LD_LIBRARY_PATH}:/usr/local/cuda/lib64:/usr/local/cuda/extras/CUPTI/lib64 | ||
OUTPUT_FOLDER=/tensorrt_models | ||
echo "Generating the following TRT Models: ${YOLO_MODELS:="yolov4-tiny-288,yolov4-tiny-416,yolov7-tiny-416"}" | ||
|
||
# Create output folder | ||
mkdir -p ${OUTPUT_FOLDER} | ||
|
||
# Install packages | ||
pip install --upgrade pip && pip install onnx==1.9.0 protobuf==3.20.3 | ||
|
||
# Clone tensorrt_demos repo | ||
git clone --depth 1 https://github.com/yeahme49/tensorrt_demos.git /tensorrt_demos | ||
|
||
# Build libyolo | ||
cd /tensorrt_demos/plugins && make all | ||
cp libyolo_layer.so ${OUTPUT_FOLDER}/libyolo_layer.so | ||
|
||
# Download yolo weights | ||
cd /tensorrt_demos/yolo && ./download_yolo.sh | ||
|
||
# Build trt engine | ||
cd /tensorrt_demos/yolo | ||
|
||
for model in ${YOLO_MODELS//,/ } | ||
do | ||
python3 yolo_to_onnx.py -m ${model} | ||
python3 onnx_to_tensorrt.py -m ${model} | ||
cp /tensorrt_demos/yolo/${model}.trt ${OUTPUT_FOLDER}/${model}.trt; | ||
done | ||
|
||
# Download Labelmap | ||
wget -q https://github.com/openvinotoolkit/open_model_zoo/raw/master/data/dataset_classes/coco_91cl.txt -O ${OUTPUT_FOLDER}/coco_91cl.txt |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.