From 9adffa1ef5b2b8e18ce791faf6c245e86c0d5785 Mon Sep 17 00:00:00 2001
From: Nicolas Mowen <nickmowen213@gmail.com>
Date: Sun, 13 Oct 2024 20:34:51 -0600
Subject: [PATCH] Detection adjustments (#14329)

---
 frigate/detectors/detector_config.py | 16 +++++---------
 frigate/util/model.py                | 33 +++++++++++++---------------
 2 files changed, 21 insertions(+), 28 deletions(-)

diff --git a/frigate/detectors/detector_config.py b/frigate/detectors/detector_config.py
index 11f08a86c..90937d8f4 100644
--- a/frigate/detectors/detector_config.py
+++ b/frigate/detectors/detector_config.py
@@ -158,17 +158,13 @@ class ModelConfig(BaseModel):
 
     def create_colormap(self, enabled_labels: set[str]) -> None:
         """Get a list of colors for enabled labels that aren't attributes."""
-        colors = generate_color_palette(
-            len(
-                list(
-                    filter(
-                        lambda label: label not in self._all_attributes, enabled_labels
-                    )
-                )
-            )
+        enabled_trackable_labels = list(
+            filter(lambda label: label not in self._all_attributes, enabled_labels)
         )
-
-        self._colormap = {label: color for label, color in zip(enabled_labels, colors)}
+        colors = generate_color_palette(len(enabled_trackable_labels))
+        self._colormap = {
+            label: color for label, color in zip(enabled_trackable_labels, colors)
+        }
 
     model_config = ConfigDict(extra="forbid", protected_namespaces=())
 
diff --git a/frigate/util/model.py b/frigate/util/model.py
index 008f5169a..685cd34ec 100644
--- a/frigate/util/model.py
+++ b/frigate/util/model.py
@@ -25,28 +25,23 @@ def get_ort_providers(
             ],
         )
 
-    providers = ort.get_available_providers()
+    providers = []
     options = []
 
-    for provider in providers:
-        if provider == "TensorrtExecutionProvider":
-            os.makedirs("/config/model_cache/tensorrt/ort/trt-engines", exist_ok=True)
-
-            if not requires_fp16 or os.environ.get("USE_FP_16", "True") != "False":
-                options.append(
-                    {
-                        "arena_extend_strategy": "kSameAsRequested",
-                        "trt_fp16_enable": requires_fp16,
-                        "trt_timing_cache_enable": True,
-                        "trt_engine_cache_enable": True,
-                        "trt_timing_cache_path": "/config/model_cache/tensorrt/ort",
-                        "trt_engine_cache_path": "/config/model_cache/tensorrt/ort/trt-engines",
-                    }
-                )
-            else:
-                options.append({})
+    for provider in ort.get_available_providers():
+        if provider == "CUDAExecutionProvider":
+            providers.append(provider)
+            options.append(
+                {
+                    "arena_extend_strategy": "kSameAsRequested",
+                }
+            )
+        elif provider == "TensorrtExecutionProvider":
+            # TensorrtExecutionProvider uses too much memory without options to control it
+            pass
         elif provider == "OpenVINOExecutionProvider":
             os.makedirs("/config/model_cache/openvino/ort", exist_ok=True)
+            providers.append(provider)
             options.append(
                 {
                     "arena_extend_strategy": "kSameAsRequested",
@@ -55,12 +50,14 @@ def get_ort_providers(
                 }
             )
         elif provider == "CPUExecutionProvider":
+            providers.append(provider)
             options.append(
                 {
                     "arena_extend_strategy": "kSameAsRequested",
                 }
             )
         else:
+            providers.append(provider)
             options.append({})
 
     return (providers, options)