add model dimensions to config

This commit is contained in:
Blake Blackshear 2020-12-08 19:39:46 -06:00
parent 9c79392060
commit 5053305e17
2 changed files with 38 additions and 0 deletions

View File

@ -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:

View File

@ -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