From de251c2c217a4d1e53ee2a469d3c579b5883840f Mon Sep 17 00:00:00 2001
From: Nate Meyer <nate.devel@gmail.com>
Date: Thu, 29 Dec 2022 00:45:39 -0500
Subject: [PATCH] Update detect to normalize input tensor using model input
 type

---
 docs/docs/configuration/detectors.md  |  1 +
 frigate/detectors/plugins/tensorrt.py | 21 +++++++++++++++------
 2 files changed, 16 insertions(+), 6 deletions(-)

diff --git a/docs/docs/configuration/detectors.md b/docs/docs/configuration/detectors.md
index 6e29b3f05..2900c7711 100644
--- a/docs/docs/configuration/detectors.md
+++ b/docs/docs/configuration/detectors.md
@@ -166,6 +166,7 @@ To generate the model files, create a new folder to save the models, download th
 ```bash
 mkdir trt-models
 wget https://github.com/blakeblackshear/frigate/raw/master/docker/tensorrt_models.sh
+chmod +x tensorrt_models.sh
 docker run --gpus=all --rm -it -v `pwd`/trt-models:/tensorrt_models -v `pwd`/tensorrt_models.sh:/tensorrt_models.sh nvcr.io/nvidia/tensorrt:22.07-py3 /tensorrt_models.sh
 ```
 
diff --git a/frigate/detectors/plugins/tensorrt.py b/frigate/detectors/plugins/tensorrt.py
index 8f107a2f3..a1f73d90f 100644
--- a/frigate/detectors/plugins/tensorrt.py
+++ b/frigate/detectors/plugins/tensorrt.py
@@ -5,7 +5,7 @@ import numpy as np
 
 try:
     import tensorrt as trt
-    from cuda import cuda, cudart
+    from cuda import cuda
 
     TRT_SUPPORT = True
 except ModuleNotFoundError as e:
@@ -72,7 +72,7 @@ class HostDeviceMem(object):
 
 class TensorRtDetector(DetectionApi):
     type_key = DETECTOR_KEY
-    # class LocalObjectDetector(ObjectDetector):
+
     def _load_engine(self, model_path):
         try:
             ctypes.cdll.LoadLibrary(
@@ -100,9 +100,15 @@ class TensorRtDetector(DetectionApi):
         assert self.engine.binding_is_input(binding)
         binding_dims = self.engine.get_binding_shape(binding)
         if len(binding_dims) == 4:
-            return tuple(binding_dims[2:])
+            return (
+                tuple(binding_dims[2:]),
+                trt.nptype(self.engine.get_binding_dtype(binding)),
+            )
         elif len(binding_dims) == 3:
-            return tuple(binding_dims[1:])
+            return (
+                tuple(binding_dims[1:]),
+                trt.nptype(self.engine.get_binding_dtype(binding)),
+            )
         else:
             raise ValueError(
                 "bad dims of binding %s: %s" % (binding, str(binding_dims))
@@ -249,10 +255,13 @@ class TensorRtDetector(DetectionApi):
         # 2..5 - a value between 0 and 1 of the box: [top, left, bottom, right]
 
         # normalize
-        tensor_input = tensor_input.astype(np.float32)
+        if self.input_shape[-1] != trt.int8:
+            tensor_input = tensor_input.astype(self.input_shape[-1])
         tensor_input /= 255.0
 
-        self.inputs[0].host = np.ascontiguousarray(tensor_input.astype(np.float32))
+        self.inputs[0].host = np.ascontiguousarray(
+            tensor_input.astype(self.input_shape[-1])
+        )
         trt_outputs = self._do_inference()
 
         raw_detections = self._postprocess_yolo(trt_outputs, self.conf_th)