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)