diff --git a/benchmark.py b/benchmark.py index 209578e74..3d0cacd87 100755 --- a/benchmark.py +++ b/benchmark.py @@ -3,7 +3,7 @@ from statistics import mean import multiprocessing as mp import numpy as np import datetime -from frigate.detectors import DetectorTypeEnum +from frigate.config import DetectorTypeEnum from frigate.object_detection import ( LocalObjectDetector, ObjectDetectProcess, diff --git a/frigate/config.py b/frigate/config.py index d3a2d2983..b22a7bf45 100644 --- a/frigate/config.py +++ b/frigate/config.py @@ -24,6 +24,7 @@ from frigate.util import ( get_ffmpeg_arg_list, escape_special_characters, load_config_with_no_duplicates, + load_labels, ) from frigate.ffmpeg_presets import ( parse_preset_hardware_acceleration, @@ -31,8 +32,8 @@ from frigate.ffmpeg_presets import ( parse_preset_output_record, parse_preset_output_rtmp, ) +from frigate.detectors import DetectorTypeEnum from frigate.version import VERSION -from frigate.detectors.config import ModelConfig, DetectorConfig logger = logging.getLogger(__name__) @@ -714,6 +715,70 @@ class DatabaseConfig(FrigateBaseModel): ) +class PixelFormatEnum(str, Enum): + rgb = "rgb" + bgr = "bgr" + yuv = "yuv" + + +class InputTensorEnum(str, Enum): + nchw = "nchw" + nhwc = "nhwc" + + +class ModelConfig(FrigateBaseModel): + path: Optional[str] = Field(title="Custom Object detection model path.") + labelmap_path: Optional[str] = Field(title="Label map for custom object detector.") + width: int = Field(default=320, title="Object detection model input width.") + height: int = Field(default=320, title="Object detection model input height.") + labelmap: Dict[int, str] = Field( + default_factory=dict, title="Labelmap customization." + ) + input_tensor: InputTensorEnum = Field( + default=InputTensorEnum.nhwc, title="Model Input Tensor Shape" + ) + input_pixel_format: PixelFormatEnum = Field( + default=PixelFormatEnum.rgb, title="Model Input Pixel Color Format" + ) + _merged_labelmap: Optional[Dict[int, str]] = PrivateAttr() + _colormap: Dict[int, Tuple[int, int, int]] = PrivateAttr() + + @property + def merged_labelmap(self) -> Dict[int, str]: + return self._merged_labelmap + + @property + def colormap(self) -> Dict[int, Tuple[int, int, int]]: + return self._colormap + + def __init__(self, **config): + super().__init__(**config) + + self._merged_labelmap = { + **load_labels(config.get("labelmap_path", "/labelmap.txt")), + **config.get("labelmap", {}), + } + + cmap = plt.cm.get_cmap("tab10", len(self._merged_labelmap.keys())) + + self._colormap = {} + for key, val in self._merged_labelmap.items(): + self._colormap[val] = tuple(int(round(255 * c)) for c in cmap(key)[:3]) + + +class DetectorConfig(BaseModel): + type: str = Field(default=DetectorTypeEnum.cpu, title="Detector Type") + device: Optional[str] = Field(default="usb", title="Device Type") + num_threads: Optional[int] = Field(default=3, title="Number of detection threads") + model: ModelConfig = Field( + default=None, title="Detector specific model configuration." + ) + + class Config: + extra = Extra.allow + arbitrary_types_allowed = True + + class LogLevelEnum(str, Enum): debug = "debug" info = "info" @@ -825,10 +890,10 @@ class FrigateConfig(FrigateBaseModel): ) ui: UIConfig = Field(default_factory=UIConfig, title="UI configuration.") model: ModelConfig = Field( - default_factory=ModelConfig, title="Default detection model configuration." + default_factory=ModelConfig, title="Detection model configuration." ) detectors: Dict[str, DetectorConfig] = Field( - default=DEFAULT_DETECTORS, + default={name: DetectorConfig(**d) for name, d in DEFAULT_DETECTORS.items()}, title="Detector hardware configuration.", ) logger: LoggerConfig = Field( diff --git a/frigate/detectors/__init__.py b/frigate/detectors/__init__.py index 741fa1829..47dfa39bc 100644 --- a/frigate/detectors/__init__.py +++ b/frigate/detectors/__init__.py @@ -2,13 +2,12 @@ import logging from .detection_api import DetectionApi from .detector_types import DetectorTypeEnum, api_types -from .config import ModelConfig, DetectorConfig logger = logging.getLogger(__name__) -def create_detector(detector_config: DetectorConfig): +def create_detector(detector_config): if detector_config.type == DetectorTypeEnum.cpu: logger.warning( "CPU detectors are not recommended and should only be used for testing or for trial purposes." diff --git a/frigate/detectors/config.py b/frigate/detectors/config.py deleted file mode 100644 index 3f666564b..000000000 --- a/frigate/detectors/config.py +++ /dev/null @@ -1,73 +0,0 @@ -import logging -from typing import Dict, List, Optional, Tuple, Union, Literal -from typing_extensions import Annotated - -import matplotlib.pyplot as plt -from pydantic import BaseModel, Extra, Field, validator, parse_obj_as -from pydantic.fields import PrivateAttr - -from frigate.enums import InputTensorEnum, PixelFormatEnum -from frigate.util import load_labels -from .detector_types import DetectorTypeEnum - - -logger = logging.getLogger(__name__) - - -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.") - width: int = Field(default=320, title="Object detection model input width.") - height: int = Field(default=320, title="Object detection model input height.") - labelmap: Dict[int, str] = Field( - default_factory=dict, title="Labelmap customization." - ) - input_tensor: InputTensorEnum = Field( - default=InputTensorEnum.nhwc, title="Model Input Tensor Shape" - ) - input_pixel_format: PixelFormatEnum = Field( - default=PixelFormatEnum.rgb, title="Model Input Pixel Color Format" - ) - _merged_labelmap: Optional[Dict[int, str]] = PrivateAttr() - _colormap: Dict[int, Tuple[int, int, int]] = PrivateAttr() - - @property - def merged_labelmap(self) -> Dict[int, str]: - return self._merged_labelmap - - @property - def colormap(self) -> Dict[int, Tuple[int, int, int]]: - return self._colormap - - def __init__(self, **config): - super().__init__(**config) - - self._merged_labelmap = { - **load_labels(config.get("labelmap_path", "/labelmap.txt")), - **config.get("labelmap", {}), - } - - cmap = plt.cm.get_cmap("tab10", len(self._merged_labelmap.keys())) - - self._colormap = {} - for key, val in self._merged_labelmap.items(): - self._colormap[val] = tuple(int(round(255 * c)) for c in cmap(key)[:3]) - - class Config: - extra = Extra.forbid - - -class DetectorConfig(BaseModel): - type: str = Field(default=DetectorTypeEnum.cpu, title="Detector Type") - model: ModelConfig = Field( - default=None, title="Detector specific model configuration." - ) - num_threads: Optional[int] = Field(default=3, title="Number of detection threads") - device: Optional[str] = Field(default="usb", title="Device Type") - - class Config: - extra = Extra.allow - arbitrary_types_allowed = True - - -DEFAULT_DETECTORS = parse_obj_as(Dict[str, DetectorConfig], {"cpu": {"type": "cpu"}}) diff --git a/frigate/enums.py b/frigate/enums.py deleted file mode 100644 index 89dcfff08..000000000 --- a/frigate/enums.py +++ /dev/null @@ -1,12 +0,0 @@ -from enum import Enum - - -class PixelFormatEnum(str, Enum): - rgb = "rgb" - bgr = "bgr" - yuv = "yuv" - - -class InputTensorEnum(str, Enum): - nchw = "nchw" - nhwc = "nhwc" diff --git a/frigate/object_detection.py b/frigate/object_detection.py index c0eb302bf..2fc080329 100644 --- a/frigate/object_detection.py +++ b/frigate/object_detection.py @@ -10,7 +10,7 @@ from abc import ABC, abstractmethod import numpy as np from setproctitle import setproctitle -from frigate.enums import InputTensorEnum +from frigate.config import InputTensorEnum from frigate.detectors import create_detector from frigate.util import EventsPerSecond, SharedMemoryFrameManager, listen, load_labels diff --git a/frigate/test/test_object_detector.py b/frigate/test/test_object_detector.py index 6c1bdb295..860f90518 100644 --- a/frigate/test/test_object_detector.py +++ b/frigate/test/test_object_detector.py @@ -3,9 +3,8 @@ from unittest.mock import Mock, patch import numpy as np -from frigate.enums import InputTensorEnum +from frigate.config import DetectorConfig, InputTensorEnum, ModelConfig from frigate.detectors import DetectorTypeEnum -from frigate.detectors.config import DetectorConfig, ModelConfig import frigate.detectors as detectors import frigate.object_detection diff --git a/frigate/video.py b/frigate/video.py index aae69194b..5e31dc457 100755 --- a/frigate/video.py +++ b/frigate/video.py @@ -14,9 +14,8 @@ import numpy as np import cv2 from setproctitle import setproctitle -from frigate.config import CameraConfig, DetectConfig +from frigate.config import CameraConfig, DetectConfig, PixelFormatEnum from frigate.const import CACHE_DIR -from frigate.enums import PixelFormatEnum from frigate.object_detection import RemoteObjectDetector from frigate.log import LogPipe from frigate.motion import MotionDetector