From fc676e7107bfe84586d745e25487978d890135d3 Mon Sep 17 00:00:00 2001 From: Nate Meyer Date: Wed, 26 Oct 2022 00:25:54 -0400 Subject: [PATCH] Change the model input tensor config to use an enumeration --- docs/docs/configuration/advanced.md | 3 ++- docs/docs/configuration/index.md | 2 +- frigate/config.py | 9 +++++++-- frigate/object_detection.py | 12 ++++-------- frigate/test/test_object_detector.py | 4 ++-- 5 files changed, 16 insertions(+), 14 deletions(-) diff --git a/docs/docs/configuration/advanced.md b/docs/docs/configuration/advanced.md index c2e93138a..a7e655e32 100644 --- a/docs/docs/configuration/advanced.md +++ b/docs/docs/configuration/advanced.md @@ -61,7 +61,7 @@ Custom models may also require different input tensor formats. The colorspace co ```yaml model: - input_tensor: ["B", "H", "W", "C"] + input_tensor: "nhwc" ``` The labelmap can be customized to your needs. A common reason to do this is to combine multiple object types that are easily confused when you don't need to be as granular such as car/truck. By default, truck is renamed to car because they are often confused. You cannot add new object types, but you can change the names of existing objects in the model. @@ -85,6 +85,7 @@ Note that if you rename objects in the labelmap, you will also need to update yo Included with Frigate is a build of ffmpeg that works for the vast majority of users. However, there exists some hardware setups which have incompatibilities with the included build. In this case, a docker volume mapping can be used to overwrite the included ffmpeg build with an ffmpeg build that works for your specific hardware setup. To do this: + 1. Download your ffmpeg build and uncompress to a folder on the host (let's use `/home/appdata/frigate/custom-ffmpeg` for this example). 2. Update your docker-compose or docker CLI to include `'/home/appdata/frigate/custom-ffmpeg':'/usr/lib/btbn-ffmpeg':'ro'` in the volume mappings. 3. Restart frigate and the custom version will be used if the mapping was done correctly. diff --git a/docs/docs/configuration/index.md b/docs/docs/configuration/index.md index e4dd7bbfd..bd8a93be0 100644 --- a/docs/docs/configuration/index.md +++ b/docs/docs/configuration/index.md @@ -101,7 +101,7 @@ model: # Valid values are rgb, bgr, or yuv. (default: shown below) input_pixel_format: rgb # Optional: Object detection model input tensor format (default: shown below) - input_tensor: ["B", "H", "W", "C"] + input_tensor: "nhwc" # Optional: Label name modifications. These are merged into the standard labelmap. labelmap: 2: vehicle diff --git a/frigate/config.py b/frigate/config.py index 005d0dfcf..90f89d204 100644 --- a/frigate/config.py +++ b/frigate/config.py @@ -693,6 +693,11 @@ class PixelFormatEnum(str, Enum): 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.") @@ -701,8 +706,8 @@ class ModelConfig(FrigateBaseModel): labelmap: Dict[int, str] = Field( default_factory=dict, title="Labelmap customization." ) - input_tensor: List[str] = Field( - default=["B", "H", "W", "C"], title="Model Input Tensor Shape" + 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" diff --git a/frigate/object_detection.py b/frigate/object_detection.py index 7dcc3088f..b4f2c2c60 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.config import DetectorTypeEnum +from frigate.config import DetectorTypeEnum, InputTensorEnum from frigate.detectors.edgetpu_tfl import EdgeTpuTfl from frigate.detectors.cpu_tfl import CpuTfl @@ -27,14 +27,10 @@ class ObjectDetector(ABC): def tensor_transform(desired_shape): # Currently this function only supports BHWC permutations - if desired_shape == ["B", "H", "W", "C"]: + if desired_shape == InputTensorEnum.nhwc: return None - else: - transform = [0] * 4 - transform[desired_shape.index("H")] = 1 - transform[desired_shape.index("W")] = 2 - transform[desired_shape.index("C")] = 3 - return tuple(transform) + elif desired_shape == InputTensorEnum.nchw: + return (0, 3, 1, 2) class LocalObjectDetector(ObjectDetector): diff --git a/frigate/test/test_object_detector.py b/frigate/test/test_object_detector.py index 4e85ebda7..f90f4d16c 100644 --- a/frigate/test/test_object_detector.py +++ b/frigate/test/test_object_detector.py @@ -2,7 +2,7 @@ import unittest from unittest.mock import patch import numpy as np -from frigate.config import DetectorTypeEnum, ModelConfig +from frigate.config import DetectorTypeEnum, InputTensorEnum, ModelConfig import frigate.object_detection @@ -66,7 +66,7 @@ class TestLocalObjectDetector(unittest.TestCase): TEST_DETECT_RESULT = np.ndarray([1, 2, 4, 8, 16, 32]) test_cfg = ModelConfig() - test_cfg.input_tensor = ["B", "C", "H", "W"] + test_cfg.input_tensor = InputTensorEnum.nchw test_obj_detect = frigate.object_detection.LocalObjectDetector( det_device=DetectorTypeEnum.cpu, model_config=test_cfg