mirror of
https://github.com/blakeblackshear/frigate.git
synced 2026-07-12 15:01:14 +03:00
Fix enabled flag for custom classification models (#23681)
Some checks are pending
CI / AMD64 Build (push) Waiting to run
CI / ARM Build (push) Waiting to run
CI / Jetson Jetpack 6 (push) Waiting to run
CI / AMD64 Extra Build (push) Blocked by required conditions
CI / ARM Extra Build (push) Blocked by required conditions
CI / Synaptics Build (push) Blocked by required conditions
CI / Assemble and push default build (push) Blocked by required conditions
Some checks are pending
CI / AMD64 Build (push) Waiting to run
CI / ARM Build (push) Waiting to run
CI / Jetson Jetpack 6 (push) Waiting to run
CI / AMD64 Extra Build (push) Blocked by required conditions
CI / ARM Extra Build (push) Blocked by required conditions
CI / Synaptics Build (push) Blocked by required conditions
CI / Assemble and push default build (push) Blocked by required conditions
* honor enabled flag for custom classification models for both startup and dynamically, even though the UI doesn't currently have a way to toggle dynamically * add test
This commit is contained in:
parent
da4037eb52
commit
5f6043aa92
@ -200,6 +200,9 @@ class EmbeddingMaintainer(threading.Thread):
|
||||
)
|
||||
|
||||
for model_config in self.config.classification.custom.values():
|
||||
if not model_config.enabled:
|
||||
continue
|
||||
|
||||
self.realtime_processors.append(
|
||||
CustomStateClassificationProcessor(
|
||||
self.config, model_config, self.requestor, self.metrics
|
||||
@ -332,6 +335,25 @@ class EmbeddingMaintainer(threading.Thread):
|
||||
for processor in self.post_processors:
|
||||
processor.update_config(topic, payload)
|
||||
|
||||
def _remove_custom_classification_processor(self, model_name: str) -> None:
|
||||
"""Shut down and drop any running processor for a custom model."""
|
||||
remaining = []
|
||||
for processor in self.realtime_processors:
|
||||
if (
|
||||
isinstance(
|
||||
processor,
|
||||
(
|
||||
CustomStateClassificationProcessor,
|
||||
CustomObjectClassificationProcessor,
|
||||
),
|
||||
)
|
||||
and processor.model_config.name == model_name
|
||||
):
|
||||
processor.shutdown()
|
||||
else:
|
||||
remaining.append(processor)
|
||||
self.realtime_processors = remaining
|
||||
|
||||
def _handle_custom_classification_update(
|
||||
self, topic: str, model_config: Any
|
||||
) -> None:
|
||||
@ -339,23 +361,7 @@ class EmbeddingMaintainer(threading.Thread):
|
||||
model_name = topic.split("/")[-1]
|
||||
|
||||
if model_config is None:
|
||||
remaining = []
|
||||
for processor in self.realtime_processors:
|
||||
if (
|
||||
isinstance(
|
||||
processor,
|
||||
(
|
||||
CustomStateClassificationProcessor,
|
||||
CustomObjectClassificationProcessor,
|
||||
),
|
||||
)
|
||||
and processor.model_config.name == model_name
|
||||
):
|
||||
processor.shutdown()
|
||||
else:
|
||||
remaining.append(processor)
|
||||
self.realtime_processors = remaining
|
||||
|
||||
self._remove_custom_classification_processor(model_name)
|
||||
logger.info(
|
||||
f"Successfully removed classification processor for model: {model_name}"
|
||||
)
|
||||
@ -363,6 +369,13 @@ class EmbeddingMaintainer(threading.Thread):
|
||||
|
||||
self.config.classification.custom[model_name] = model_config
|
||||
|
||||
# A disabled model must not run; tear down any existing processor and
|
||||
# do not register a new one.
|
||||
if not model_config.enabled:
|
||||
self._remove_custom_classification_processor(model_name)
|
||||
logger.info(f"Disabled classification processor for model: {model_name}")
|
||||
return
|
||||
|
||||
# Check if processor already exists
|
||||
for processor in self.realtime_processors:
|
||||
if isinstance(
|
||||
@ -702,7 +715,11 @@ class EmbeddingMaintainer(threading.Thread):
|
||||
and "license_plate" not in camera_config.objects.track
|
||||
)
|
||||
|
||||
if not dedicated_lpr_enabled and len(self.config.classification.custom) == 0:
|
||||
has_enabled_custom = any(
|
||||
c.enabled for c in self.config.classification.custom.values()
|
||||
)
|
||||
|
||||
if not dedicated_lpr_enabled and not has_enabled_custom:
|
||||
# no active features that use this data
|
||||
return
|
||||
|
||||
|
||||
106
frigate/test/test_classification_enabled.py
Normal file
106
frigate/test/test_classification_enabled.py
Normal file
@ -0,0 +1,106 @@
|
||||
"""Tests that disabled custom classification models are not registered or run."""
|
||||
|
||||
import sys
|
||||
import unittest
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
# Mock TFLite before importing the maintainer / classification modules
|
||||
_MOCK_MODULES = [
|
||||
"tflite_runtime",
|
||||
"tflite_runtime.interpreter",
|
||||
"ai_edge_litert",
|
||||
"ai_edge_litert.interpreter",
|
||||
]
|
||||
for mod in _MOCK_MODULES:
|
||||
if mod not in sys.modules:
|
||||
sys.modules[mod] = MagicMock()
|
||||
|
||||
from frigate.data_processing.real_time.custom_classification import ( # noqa: E402
|
||||
CustomObjectClassificationProcessor,
|
||||
)
|
||||
from frigate.embeddings.maintainer import EmbeddingMaintainer # noqa: E402
|
||||
|
||||
|
||||
class TestCustomClassificationEnabledGating(unittest.TestCase):
|
||||
"""A model with enabled: false must not keep a processor registered."""
|
||||
|
||||
def _make_maintainer(self) -> EmbeddingMaintainer:
|
||||
# Bypass the heavy __init__; only the attributes touched by the
|
||||
# config update path are needed for these tests.
|
||||
maintainer = EmbeddingMaintainer.__new__(EmbeddingMaintainer)
|
||||
maintainer.realtime_processors = []
|
||||
maintainer.config = MagicMock()
|
||||
maintainer.config.classification.custom = {}
|
||||
maintainer.requestor = MagicMock()
|
||||
maintainer.metrics = MagicMock()
|
||||
maintainer.event_metadata_publisher = MagicMock()
|
||||
return maintainer
|
||||
|
||||
def _make_model_config(self, name: str, enabled: bool) -> MagicMock:
|
||||
model_config = MagicMock()
|
||||
model_config.name = name
|
||||
model_config.enabled = enabled
|
||||
model_config.state_config = None
|
||||
return model_config
|
||||
|
||||
def _make_processor(self, name: str) -> MagicMock:
|
||||
processor = MagicMock(spec=CustomObjectClassificationProcessor)
|
||||
processor.model_config = MagicMock()
|
||||
processor.model_config.name = name
|
||||
return processor
|
||||
|
||||
def test_disabled_update_tears_down_existing_processor(self):
|
||||
"""Toggling a running model to disabled shuts down and drops its processor."""
|
||||
maintainer = self._make_maintainer()
|
||||
processor = self._make_processor("atli")
|
||||
maintainer.realtime_processors = [processor]
|
||||
|
||||
maintainer._handle_custom_classification_update(
|
||||
"config/classification/custom/atli",
|
||||
self._make_model_config("atli", enabled=False),
|
||||
)
|
||||
|
||||
processor.shutdown.assert_called_once()
|
||||
self.assertEqual(maintainer.realtime_processors, [])
|
||||
|
||||
def test_disabled_update_does_not_register_processor(self):
|
||||
"""A disabled model that has no processor is never registered."""
|
||||
maintainer = self._make_maintainer()
|
||||
|
||||
maintainer._handle_custom_classification_update(
|
||||
"config/classification/custom/atli",
|
||||
self._make_model_config("atli", enabled=False),
|
||||
)
|
||||
|
||||
self.assertEqual(maintainer.realtime_processors, [])
|
||||
|
||||
def test_disabled_update_leaves_other_processors_untouched(self):
|
||||
"""Disabling one model must not affect other running processors."""
|
||||
maintainer = self._make_maintainer()
|
||||
other = self._make_processor("simbi")
|
||||
maintainer.realtime_processors = [other]
|
||||
|
||||
maintainer._handle_custom_classification_update(
|
||||
"config/classification/custom/atli",
|
||||
self._make_model_config("atli", enabled=False),
|
||||
)
|
||||
|
||||
other.shutdown.assert_not_called()
|
||||
self.assertEqual(maintainer.realtime_processors, [other])
|
||||
|
||||
def test_removed_model_tears_down_processor(self):
|
||||
"""A None payload (model deleted) still shuts down its processor."""
|
||||
maintainer = self._make_maintainer()
|
||||
processor = self._make_processor("atli")
|
||||
maintainer.realtime_processors = [processor]
|
||||
|
||||
maintainer._handle_custom_classification_update(
|
||||
"config/classification/custom/atli", None
|
||||
)
|
||||
|
||||
processor.shutdown.assert_called_once()
|
||||
self.assertEqual(maintainer.realtime_processors, [])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Loading…
Reference in New Issue
Block a user