Mirrors/blakeblackshear.frigate
1
0
Fork 0
mirror of https://github.com/blakeblackshear/frigate.git synced 2024-11-26 19:06:11 +01:00
Code Issues Projects Releases Wiki Activity

Detection adjustments (#14329)

Browse Source
This commit is contained in:
Nicolas Mowen 2024-10-13 20:34:51 -06:00 committed by GitHub
parent 4ca267ea17
commit 9adffa1ef5
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 21 additions and 28 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

16
frigate/detectors/detector_config.py
View File

@ -158,17 +158,13 @@ class ModelConfig(BaseModel):
def create_colormap(self, enabled_labels: set[str]) -> None: def create_colormap(self, enabled_labels: set[str]) -> None:
"""Get a list of colors for enabled labels that aren't attributes.""" """Get a list of colors for enabled labels that aren't attributes."""
colors = generate_color_palette( enabled_trackable_labels = list(
len( filter(lambda label: label not in self._all_attributes, enabled_labels)
list(
filter(
lambda label: label not in self._all_attributes, enabled_labels
) )
) colors = generate_color_palette(len(enabled_trackable_labels))
) self._colormap = {
) label: color for label, color in zip(enabled_trackable_labels, colors)
}
self._colormap = {label: color for label, color in zip(enabled_labels, colors)}
model_config = ConfigDict(extra="forbid", protected_namespaces=()) model_config = ConfigDict(extra="forbid", protected_namespaces=())

23
frigate/util/model.py
View File

@ -25,28 +25,23 @@ def get_ort_providers(
], ],
) )
providers = ort.get_available_providers() providers = []
options = [] options = []
for provider in providers: for provider in ort.get_available_providers():
if provider == "TensorrtExecutionProvider": if provider == "CUDAExecutionProvider":
os.makedirs("/config/model_cache/tensorrt/ort/trt-engines", exist_ok=True) providers.append(provider)
if not requires_fp16 or os.environ.get("USE_FP_16", "True") != "False":
options.append( options.append(
{ {
"arena_extend_strategy": "kSameAsRequested", "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: elif provider == "TensorrtExecutionProvider":
options.append({}) # TensorrtExecutionProvider uses too much memory without options to control it
pass
elif provider == "OpenVINOExecutionProvider": elif provider == "OpenVINOExecutionProvider":
os.makedirs("/config/model_cache/openvino/ort", exist_ok=True) os.makedirs("/config/model_cache/openvino/ort", exist_ok=True)
providers.append(provider)
options.append( options.append(
{ {
"arena_extend_strategy": "kSameAsRequested", "arena_extend_strategy": "kSameAsRequested",
@ -55,12 +50,14 @@ def get_ort_providers(
} }
) )
elif provider == "CPUExecutionProvider": elif provider == "CPUExecutionProvider":
providers.append(provider)
options.append( options.append(
{ {
"arena_extend_strategy": "kSameAsRequested", "arena_extend_strategy": "kSameAsRequested",
} }
) )
else: else:
providers.append(provider)
options.append({}) options.append({})
return (providers, options) return (providers, options)