From 5053305e17be2a3c7501b5340d44324b6b07eef4 Mon Sep 17 00:00:00 2001
From: Blake Blackshear <blakeb@blakeshome.com>
Date: Tue, 8 Dec 2020 19:39:46 -0600
Subject: [PATCH] add model dimensions to config

---
 README.md         |  9 +++++++++
 frigate/config.py | 29 +++++++++++++++++++++++++++++
 2 files changed, 38 insertions(+)

diff --git a/README.md b/README.md
index ad2e9717b..f5e03518f 100644
--- a/README.md
+++ b/README.md
@@ -199,6 +199,13 @@ detectors:
     # Optional: device name as defined here: https://coral.ai/docs/edgetpu/multiple-edgetpu/#using-the-tensorflow-lite-python-api
     device: usb
 
+# Optional: model configuration
+model:
+  # Required: height of the trained model
+  height: 320
+  # Required: width of the trained model
+  width: 320
+
 # Required: mqtt configuration
 mqtt:
   # Required: host name
@@ -880,6 +887,8 @@ Models for both CPU and EdgeTPU (Coral) are bundled in the image. You can use yo
 - EdgeTPU Model: `/edgetpu_model.tflite`
 - Labels: `/labelmap.txt`
 
+You also need to update the model width/height in the config if they differ from the defaults.
+
 ### Customizing the Labelmap
 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. You must retain the same number of labels, but you can change the names. To change:
 
diff --git a/frigate/config.py b/frigate/config.py
index 6e4126b30..39e5c8af8 100644
--- a/frigate/config.py
+++ b/frigate/config.py
@@ -193,6 +193,10 @@ CAMERAS_SCHEMA = vol.Schema(vol.All(
 
 FRIGATE_CONFIG_SCHEMA = vol.Schema(
     {
+        vol.Optional('model', default={'width': 300, 'height': 300}): {
+            vol.Required('width'): int,
+            vol.Required('height'): int
+        },
         vol.Optional('detectors', default=DEFAULT_DETECTORS): DETECTORS_SCHEMA,
         'mqtt': MQTT_SCHEMA,
         vol.Optional('logger', default={'default': 'info', 'logs': {}}): {
@@ -210,6 +214,25 @@ FRIGATE_CONFIG_SCHEMA = vol.Schema(
     }
 )
 
+class ModelConfig():
+    def __init__(self, config):
+        self._width = config['width']
+        self._height = config['height']
+    
+    @property
+    def width(self):
+        return self._width
+    
+    @property
+    def height(self):
+        return self._height
+    
+    def to_dict(self):
+        return {
+            'width': self.width,
+            'height': self.height
+        }
+
 class DetectorConfig():
     def __init__(self, config):
         self._type = config['type']
@@ -756,6 +779,7 @@ class FrigateConfig():
 
         config = self._sub_env_vars(config)
 
+        self._model = ModelConfig(config['model'])
         self._detectors = { name: DetectorConfig(d) for name, d in config['detectors'].items() }
         self._mqtt = MqttConfig(config['mqtt'])
         self._save_clips = SaveClipsConfig(config['save_clips'])
@@ -787,6 +811,7 @@ class FrigateConfig():
     
     def to_dict(self):
         return {
+            'model': self.model.to_dict(),
             'detectors': {k: d.to_dict() for k, d in self.detectors.items()},
             'mqtt': self.mqtt.to_dict(),
             'save_clips': self.save_clips.to_dict(),
@@ -794,6 +819,10 @@ class FrigateConfig():
             'logger': self.logger.to_dict()
         }
     
+    @property
+    def model(self):
+        return self._model
+    
     @property
     def detectors(self) -> Dict[str, DetectorConfig]:
         return self._detectors