diff --git a/frigate/config.py b/frigate/config.py index c86be26ac..889f1bc45 100644 --- a/frigate/config.py +++ b/frigate/config.py @@ -12,7 +12,7 @@ from pydantic import BaseModel, Extra, Field, parse_obj_as, validator from pydantic.fields import PrivateAttr from frigate.const import CACHE_DIR, DEFAULT_DB_PATH, REGEX_CAMERA_NAME, YAML_EXT -from frigate.detectors import DetectorConfig, ModelConfig +from frigate.detectors import AudioModelConfig, DetectorConfig, ModelConfig from frigate.detectors.detector_config import BaseDetectorConfig from frigate.ffmpeg_presets import ( parse_preset_hardware_acceleration_decode, diff --git a/frigate/detectors/__init__.py b/frigate/detectors/__init__.py index 7465ed7c0..8e522ea3b 100644 --- a/frigate/detectors/__init__.py +++ b/frigate/detectors/__init__.py @@ -1,6 +1,11 @@ import logging -from .detector_config import InputTensorEnum, ModelConfig, PixelFormatEnum # noqa: F401 +from .detector_config import ( # noqa: F401 + AudioModelConfig, + InputTensorEnum, + ModelConfig, + PixelFormatEnum, +) from .detector_types import DetectorConfig, DetectorTypeEnum, api_types # noqa: F401 logger = logging.getLogger(__name__) diff --git a/frigate/detectors/detector_config.py b/frigate/detectors/detector_config.py index 6a179429a..5585a56f1 100644 --- a/frigate/detectors/detector_config.py +++ b/frigate/detectors/detector_config.py @@ -10,6 +10,7 @@ import requests from pydantic import BaseModel, Extra, Field from pydantic.fields import PrivateAttr +from frigate.const import AUDIO_DURATION, AUDIO_FORMAT, AUDIO_SAMPLE_RATE from frigate.plus import PlusApi from frigate.util import load_labels @@ -39,6 +40,18 @@ class ModelTypeEnum(str, Enum): yolov8 = "yolov8" +class AudioModelConfig(BaseModel): + path: str = Field( + default="/cpu_audio_model.tflite", title="Custom Object detection model path." + ) + labelmap_path: str = Field( + default="/audio-labelmap.txt", title="Label map for custom object detector." + ) + duration: float = Field(default=AUDIO_DURATION, title="Audio duration.") + format: str = Field(default=AUDIO_FORMAT, title="Audio format.") + sample_rate: int = Field(default=AUDIO_SAMPLE_RATE, title="Audio sample rate.") + + class ModelConfig(BaseModel): path: Optional[str] = Field(title="Custom Object detection model path.") labelmap_path: Optional[str] = Field(title="Label map for custom object detector.")