From 46e72227f973f188467ffa5410d1db8c882394f4 Mon Sep 17 00:00:00 2001 From: Nicolas Mowen Date: Fri, 23 May 2025 06:05:33 -0600 Subject: [PATCH] Adjust config for state and object --- frigate/config/classification.py | 22 +++++++++++++++++----- 1 file changed, 17 insertions(+), 5 deletions(-) diff --git a/frigate/config/classification.py b/frigate/config/classification.py index 76e87a06a..94548ee01 100644 --- a/frigate/config/classification.py +++ b/frigate/config/classification.py @@ -34,16 +34,28 @@ class BirdClassificationConfig(FrigateBaseModel): ) -class TeachableMachineModelType(str, Enum): - object = "object" - state = "state" +class TeachableMachineStateCameraConfig(FrigateBaseModel): + crop: list[int, int, int, int] = Field( + title="Crop of image frame on this camera to run classification on." + ) + + +class TeachableMachineStateConfig(FrigateBaseModel): + cameras: Dict[str, TeachableMachineStateCameraConfig] = Field( + title="Cameras to run classification on." + ) + + +class TeachableMachineObjectConfig(FrigateBaseModel): + objects: list[str] = Field(title="Object types to classify.") class TeachableMachineConfig(FrigateBaseModel): enabled: bool = Field(default=True, title="Enable running the model.") model_path: str = Field(title="Path to teachable machine tflite model.") labelmap_path: str = Field(title="Path to teachable machine labelmap.") - model_type: TeachableMachineModelType = Field(title="Type of model.") + object_config: TeachableMachineObjectConfig | None = Field(default=None) + state_config: TeachableMachineStateConfig | None = Field(default=None) class ClassificationConfig(FrigateBaseModel): @@ -51,7 +63,7 @@ class ClassificationConfig(FrigateBaseModel): default_factory=BirdClassificationConfig, title="Bird classification config." ) teachable_machine: Dict[str, TeachableMachineConfig] = Field( - title="Teachable Machine Model Configs." + default={}, title="Teachable Machine Model Configs." )