From af303cbf2a08b7548f249e6b5cd2006f8f43ec47 Mon Sep 17 00:00:00 2001 From: Blake Blackshear Date: Tue, 3 Nov 2020 08:15:58 -0600 Subject: [PATCH] create typed config classes --- frigate/__main__.py | 116 ++------- frigate/config.py | 466 ++++++++++++++++++++++++++++++++++- frigate/events.py | 17 +- frigate/http.py | 6 +- frigate/mqtt.py | 16 +- frigate/object_processing.py | 65 ++--- frigate/test/test_config.py | 169 ++++++++++++- frigate/util.py | 32 +++ frigate/video.py | 87 ++----- 9 files changed, 733 insertions(+), 241 deletions(-) diff --git a/frigate/__main__.py b/frigate/__main__.py index 1bfd6be27..95fff2cbd 100644 --- a/frigate/__main__.py +++ b/frigate/__main__.py @@ -7,100 +7,31 @@ import multiprocessing as mp from playhouse.sqlite_ext import SqliteExtDatabase from typing import Dict, List -from frigate.config import FRIGATE_CONFIG_SCHEMA +from frigate.config import FrigateConfig from frigate.edgetpu import EdgeTPUProcess from frigate.events import EventProcessor from frigate.http import create_app from frigate.models import Event from frigate.mqtt import create_mqtt_client from frigate.object_processing import TrackedObjectProcessor -from frigate.video import get_frame_shape, track_camera, get_ffmpeg_input, capture_camera +from frigate.video import track_camera, capture_camera from frigate.watchdog import FrigateWatchdog class FrigateApp(): def __init__(self): self.stop_event = mp.Event() - self.config: dict = None + self.config: FrigateConfig = None self.detection_queue = mp.Queue() - self.detectors: Dict[str: EdgeTPUProcess] = {} - self.detection_out_events: Dict[str: mp.Event] = {} + self.detectors: Dict[str, EdgeTPUProcess] = {} + self.detection_out_events: Dict[str, mp.Event] = {} self.detection_shms: List[mp.shared_memory.SharedMemory] = [] self.camera_metrics = {} def init_config(self): - # TODO: sub in FRIGATE_ENV vars - frigate_env_vars = {k: v for k, v in os.environ.items() if k.startswith('FRIGATE_')} config_file = os.environ.get('CONFIG_FILE', '/config/config.yml') + self.config = FrigateConfig(config_file=config_file) - with open(config_file) as f: - raw_config = f.read() - - if config_file.endswith(".yml"): - config = yaml.safe_load(raw_config) - elif config_file.endswith(".json"): - config = json.loads(raw_config) - - self.config = FRIGATE_CONFIG_SCHEMA(config) - - if 'password' in self.config['mqtt']: - self.config['mqtt']['password'] = self.config['mqtt']['password'].format(**frigate_env_vars) - - cache_dir = self.config['save_clips']['cache_dir'] - clips_dir = self.config['save_clips']['clips_dir'] - - if not os.path.exists(cache_dir) and not os.path.islink(cache_dir): - os.makedirs(cache_dir) - if not os.path.exists(clips_dir) and not os.path.islink(clips_dir): - os.makedirs(clips_dir) - - for camera_name, camera_config in self.config['cameras'].items(): - - # set shape - if 'width' in camera_config and 'height' in camera_config: - frame_shape = (camera_config['height'], camera_config['width'], 3) - else: - frame_shape = get_frame_shape(camera_config['ffmpeg']['input']) - - camera_config['frame_shape'] = frame_shape - - # build ffmpeg command - ffmpeg = camera_config['ffmpeg'] - ffmpeg_input = ffmpeg['input'].format(**frigate_env_vars) - ffmpeg_global_args = ffmpeg.get('global_args', self.config['ffmpeg']['global_args']) - ffmpeg_hwaccel_args = ffmpeg.get('hwaccel_args', self.config['ffmpeg']['hwaccel_args']) - ffmpeg_input_args = ffmpeg.get('input_args', self.config['ffmpeg']['input_args']) - ffmpeg_output_args = ffmpeg.get('output_args', self.config['ffmpeg']['output_args']) - if not camera_config.get('fps') is None: - ffmpeg_output_args = ["-r", str(camera_config['fps'])] + ffmpeg_output_args - if camera_config['save_clips']['enabled']: - ffmpeg_output_args = [ - "-f", - "segment", - "-segment_time", - "10", - "-segment_format", - "mp4", - "-reset_timestamps", - "1", - "-strftime", - "1", - "-c", - "copy", - "-an", - "-map", - "0", - f"{os.path.join(self.config['save_clips']['cache_dir'], camera_name)}-%Y%m%d%H%M%S.mp4" - ] + ffmpeg_output_args - ffmpeg_cmd = (['ffmpeg'] + - ffmpeg_global_args + - ffmpeg_hwaccel_args + - ffmpeg_input_args + - ['-i', ffmpeg_input] + - ffmpeg_output_args + - ['pipe:']) - - camera_config['ffmpeg_cmd'] = ffmpeg_cmd - + for camera_name in self.config.cameras.keys(): # create camera_metrics self.camera_metrics[camera_name] = { 'camera_fps': mp.Value('d', 0.0), @@ -118,10 +49,10 @@ class FrigateApp(): self.event_queue = mp.Queue() # Queue for cameras to push tracked objects to - self.detected_frames_queue = mp.Queue(maxsize=len(self.config['cameras'].keys())*2) + self.detected_frames_queue = mp.Queue(maxsize=len(self.config.cameras.keys())*2) def init_database(self): - self.db = SqliteExtDatabase(f"/{os.path.join(self.config['save_clips']['clips_dir'], 'frigate.db')}") + self.db = SqliteExtDatabase(f"/{os.path.join(self.config.save_clips.clips_dir, 'frigate.db')}") models = [Event] self.db.bind(models) self.db.create_tables(models, safe=True) @@ -130,38 +61,29 @@ class FrigateApp(): self.flask_app = create_app(self.config, self.db, self.camera_metrics, self.detectors, self.detected_frames_processor) def init_mqtt(self): - # TODO: create config class - mqtt_config = self.config['mqtt'] - self.mqtt_client = create_mqtt_client( - mqtt_config['host'], - mqtt_config['port'], - mqtt_config['topic_prefix'], - mqtt_config['client_id'], - mqtt_config.get('user'), - mqtt_config.get('password') - ) + self.mqtt_client = create_mqtt_client(self.config.mqtt) def start_detectors(self): - for name in self.config['cameras'].keys(): + for name in self.config.cameras.keys(): self.detection_out_events[name] = mp.Event() shm_in = mp.shared_memory.SharedMemory(name=name, create=True, size=300*300*3) shm_out = mp.shared_memory.SharedMemory(name=f"out-{name}", create=True, size=20*6*4) self.detection_shms.append(shm_in) self.detection_shms.append(shm_out) - for name, detector in self.config['detectors'].items(): - if detector['type'] == 'cpu': + for name, detector in self.config.detectors.items(): + if detector.type == 'cpu': self.detectors[name] = EdgeTPUProcess(self.detection_queue, out_events=self.detection_out_events, tf_device='cpu') - if detector['type'] == 'edgetpu': - self.detectors[name] = EdgeTPUProcess(self.detection_queue, out_events=self.detection_out_events, tf_device=detector['device']) + if detector.type == 'edgetpu': + self.detectors[name] = EdgeTPUProcess(self.detection_queue, out_events=self.detection_out_events, tf_device=detector.device) def start_detected_frames_processor(self): - self.detected_frames_processor = TrackedObjectProcessor(self.config['cameras'], self.mqtt_client, self.config['mqtt']['topic_prefix'], + self.detected_frames_processor = TrackedObjectProcessor(self.config.cameras, self.mqtt_client, self.config.mqtt.topic_prefix, self.detected_frames_queue, self.event_queue, self.stop_event) self.detected_frames_processor.start() def start_camera_processors(self): - for name, config in self.config['cameras'].items(): + for name, config in self.config.cameras.items(): camera_process = mp.Process(target=track_camera, args=(name, config, self.detection_queue, self.detection_out_events[name], self.detected_frames_queue, self.camera_metrics[name])) @@ -171,7 +93,7 @@ class FrigateApp(): print(f"Camera processor started for {name}: {camera_process.pid}") def start_camera_capture_processes(self): - for name, config in self.config['cameras'].items(): + for name, config in self.config.cameras.items(): capture_process = mp.Process(target=capture_camera, args=(name, config, self.camera_metrics[name])) capture_process.daemon = True @@ -199,7 +121,7 @@ class FrigateApp(): self.init_web_server() self.start_event_processor() self.start_watchdog() - self.flask_app.run(host='0.0.0.0', port=self.config['web_port'], debug=False) + self.flask_app.run(host='0.0.0.0', port=self.config.web_port, debug=False) self.stop() def stop(self): diff --git a/frigate/config.py b/frigate/config.py index 089167e1b..1184a21c3 100644 --- a/frigate/config.py +++ b/frigate/config.py @@ -1,5 +1,17 @@ +import base64 +import json +import os +import yaml + +from typing import Dict + +import cv2 +import matplotlib.pyplot as plt +import numpy as np import voluptuous as vol +from frigate.util import get_frame_shape + DETECTORS_SCHEMA = vol.Schema( { vol.Required(str): { @@ -66,12 +78,21 @@ FILTER_SCHEMA = vol.Schema( } ) -OBJECTS_SCHEMA = vol.Schema( +def filters_for_all_tracked_objects(object_config): + for tracked_object in object_config.get('track', ['person']): + if not 'filters' in object_config: + object_config['filters'] = {} + if not tracked_object in object_config['filters']: + object_config['filters'][tracked_object] = {} + return object_config + +OBJECTS_SCHEMA = vol.Schema(vol.All(filters_for_all_tracked_objects, { vol.Optional('track', default=['person']): [str], - 'filters': FILTER_SCHEMA.extend({vol.Optional('min_score', default=0.5): float}) + # TODO: this should populate filters for all tracked objects + vol.Optional('filters', default = {}): FILTER_SCHEMA.extend({ str: {vol.Optional('min_score', default=0.5): float}}) } -) +)) DEFAULT_CAMERA_MQTT = { 'crop_to_region': True @@ -99,8 +120,8 @@ CAMERAS_SCHEMA = vol.Schema( { str: { vol.Required('ffmpeg'): CAMERA_FFMPEG_SCHEMA, - 'height': int, - 'width': int, + vol.Required('height'): int, + vol.Required('width'): int, 'fps': int, 'mask': str, vol.Optional('best_image_timeout', default=60): int, @@ -140,3 +161,438 @@ FRIGATE_CONFIG_SCHEMA = vol.Schema( vol.Required('cameras', default={}): CAMERAS_SCHEMA } ) + +class DetectorConfig(): + def __init__(self, config): + self._type = config['type'] + self._device = config['device'] + + @property + def type(self): + return self._type + + @property + def device(self): + return self._device + + +class MqttConfig(): + def __init__(self, config): + self._host = config['host'] + self._port = config['port'] + self._topic_prefix = config['topic_prefix'] + self._client_id = config['client_id'] + self._user = config.get('user') + self._password = config.get('password') + + @property + def host(self): + return self._host + + @property + def port(self): + return self._port + + @property + def topic_prefix(self): + return self._topic_prefix + + @property + def client_id(self): + return self._client_id + + @property + def user(self): + return self._user + + @property + def password(self): + return self._password + +class SaveClipsConfig(): + def __init__(self, config): + self._max_seconds = config['max_seconds'] + self._clips_dir = config['clips_dir'] + self._cache_dir = config['cache_dir'] + + @property + def max_seconds(self): + return self._max_seconds + + @property + def clips_dir(self): + return self._clips_dir + + @property + def cache_dir(self): + return self._cache_dir + +class FfmpegConfig(): + def __init__(self, global_config, config): + self._input = config.get('input') + self._global_args = config.get('global_args', global_config['global_args']) + self._hwaccel_args = config.get('hwaccel_args', global_config['hwaccel_args']) + self._input_args = config.get('input_args', global_config['input_args']) + self._output_args = config.get('output_args', global_config['output_args']) + + @property + def input(self): + return self._input + + @property + def global_args(self): + return self._global_args + + @property + def hwaccel_args(self): + return self._hwaccel_args + + @property + def input_args(self): + return self._input_args + + @property + def output_args(self): + return self._output_args + +class FilterConfig(): + def __init__(self, config): + self._min_area = config['min_area'] + self._max_area = config['max_area'] + self._threshold = config['threshold'] + self._min_score = config.get('min_score') + + @property + def min_area(self): + return self._min_area + + @property + def max_area(self): + return self._max_area + + @property + def threshold(self): + return self._threshold + + @property + def min_score(self): + return self._min_score + +class ObjectConfig(): + def __init__(self, global_config, config): + self._track = config.get('track', global_config['track']) + if 'filters' in config: + self._filters = { name: FilterConfig(c) for name, c in config['filters'].items() } + else: + self._filters = { name: FilterConfig(c) for name, c in global_config['filters'].items() } + + @property + def track(self): + return self._track + + @property + def filters(self) -> Dict[str, FilterConfig]: + return self._filters + +class CameraSnapshotsConfig(): + def __init__(self, config): + self._show_timestamp = config['show_timestamp'] + self._draw_zones = config['draw_zones'] + self._draw_bounding_boxes = config['draw_bounding_boxes'] + + @property + def show_timestamp(self): + return self._show_timestamp + + @property + def draw_zones(self): + return self._draw_zones + + @property + def draw_bounding_boxes(self): + return self._draw_bounding_boxes + +class CameraSaveClipsConfig(): + def __init__(self, config): + self._enabled = config['enabled'] + self._pre_capture = config['pre_capture'] + self._objects = config.get('objects') + + @property + def enabled(self): + return self._enabled + + @property + def pre_capture(self): + return self._pre_capture + + @property + def objects(self): + return self._objects + +class CameraMqttConfig(): + def __init__(self, config): + self._crop_to_region = config['crop_to_region'] + self._snapshot_height = config.get('snapshot_height') + + @property + def crop_to_region(self): + return self._crop_to_region + + @property + def snapshot_height(self): + return self._snapshot_height + +class ZoneConfig(): + def __init__(self, name, config): + self._coordinates = config['coordinates'] + self._filters = { name: FilterConfig(c) for name, c in config['filters'].items() } + + if isinstance(self._coordinates, list): + self._contour = np.array([[int(p.split(',')[0]), int(p.split(',')[1])] for p in self._coordinates]) + elif isinstance(self._coordinates, str): + points = self._coordinates.split(',') + self._contour = np.array([[int(points[i]), int(points[i+1])] for i in range(0, len(points), 2)]) + else: + print(f"Unable to parse zone coordinates for {name}") + self._contour = np.array([]) + + self._color = (0,0,0) + + @property + def coordinates(self): + return self._coordinates + + @property + def contour(self): + return self._contour + + @contour.setter + def contour(self, val): + self._contour = val + + @property + def color(self): + return self._color + + @color.setter + def color(self, val): + self._color = val + + @property + def filters(self): + return self._filters + +class CameraConfig(): + def __init__(self, name, config, cache_dir, global_ffmpeg, global_objects): + self._name = name + self._ffmpeg = FfmpegConfig(global_ffmpeg, config['ffmpeg']) + self._height = config.get('height') + self._width = config.get('width') + self._frame_shape = (self._height, self._width) + self._frame_shape_yuv = (self._frame_shape[0]*3//2, self._frame_shape[1]) + self._fps = config.get('fps') + self._mask = self._create_mask(config.get('mask')) + self._best_image_timeout = config['best_image_timeout'] + self._mqtt = CameraMqttConfig(config['mqtt']) + self._zones = { name: ZoneConfig(name, z) for name, z in config['zones'].items() } + self._save_clips = CameraSaveClipsConfig(config['save_clips']) + self._snapshots = CameraSnapshotsConfig(config['snapshots']) + self._objects = ObjectConfig(global_objects, config.get('objects', {})) + + self._ffmpeg_cmd = self._get_ffmpeg_cmd(cache_dir) + + self._set_zone_colors(self._zones) + + def _create_mask(self, mask): + if mask: + if mask.startswith('base64,'): + img = base64.b64decode(mask[7:]) + np_img = np.fromstring(img, dtype=np.uint8) + mask_img = cv2.imdecode(np_img, cv2.IMREAD_GRAYSCALE) + elif mask.startswith('poly,'): + points = mask.split(',')[1:] + contour = np.array([[int(points[i]), int(points[i+1])] for i in range(0, len(points), 2)]) + mask_img = np.zeros(self.frame_shape, np.uint8) + mask_img[:] = 255 + cv2.fillPoly(mask_img, pts=[contour], color=(0)) + else: + mask_img = cv2.imread(f"/config/{mask}", cv2.IMREAD_GRAYSCALE) + else: + mask_img = None + + if mask_img is None or mask_img.size == 0: + mask_img = np.zeros(self.frame_shape, np.uint8) + mask_img[:] = 255 + + return mask_img + + def _get_ffmpeg_cmd(self, cache_dir): + ffmpeg_output_args = self.ffmpeg.output_args + if self.fps: + ffmpeg_output_args = ["-r", str(self.fps)] + ffmpeg_output_args + if self.save_clips.enabled: + ffmpeg_output_args = [ + "-f", + "segment", + "-segment_time", + "10", + "-segment_format", + "mp4", + "-reset_timestamps", + "1", + "-strftime", + "1", + "-c", + "copy", + "-an", + f"{os.path.join(cache_dir, self.name)}-%Y%m%d%H%M%S.mp4" + ] + ffmpeg_output_args + return (['ffmpeg'] + + self.ffmpeg.global_args + + self.ffmpeg.hwaccel_args + + self.ffmpeg.input_args + + ['-i', self.ffmpeg.input] + + ffmpeg_output_args + + ['pipe:']) + + def _set_zone_colors(self, zones: Dict[str, ZoneConfig]): + # set colors for zones + all_zone_names = zones.keys() + zone_colors = {} + colors = plt.cm.get_cmap('tab10', len(all_zone_names)) + for i, zone in enumerate(all_zone_names): + zone_colors[zone] = tuple(int(round(255 * c)) for c in colors(i)[:3]) + + for name, zone in zones.items(): + zone.color = zone_colors[name] + + @property + def name(self): + return self._name + + @property + def ffmpeg(self): + return self._ffmpeg + + @property + def height(self): + return self._height + + @property + def width(self): + return self._width + + @property + def fps(self): + return self._fps + + @property + def mask(self): + return self._mask + + @property + def best_image_timeout(self): + return self._best_image_timeout + + @property + def mqtt(self): + return self._mqtt + + @property + def zones(self)-> Dict[str, ZoneConfig]: + return self._zones + + @property + def save_clips(self): + return self._save_clips + + @property + def snapshots(self): + return self._snapshots + + @property + def objects(self): + return self._objects + + @property + def frame_shape(self): + return self._frame_shape + + @property + def frame_shape_yuv(self): + return self._frame_shape_yuv + + @property + def ffmpeg_cmd(self): + return self._ffmpeg_cmd + +class FrigateConfig(): + def __init__(self, config_file=None, config=None): + if config is None and config_file is None: + raise ValueError('config or config_file must be defined') + elif not config_file is None: + config = self._load_file(config_file) + + config = FRIGATE_CONFIG_SCHEMA(config) + + config = self._sub_env_vars(config) + + self._web_port = config['web_port'] + self._detectors = { name: DetectorConfig(d) for name, d in config['detectors'].items() } + self._mqtt = MqttConfig(config['mqtt']) + self._save_clips = SaveClipsConfig(config['save_clips']) + self._cameras = { name: CameraConfig(name, c, self._save_clips.cache_dir, config['ffmpeg'], config['objects']) for name, c in config['cameras'].items() } + + self._ensure_dirs() + + def _sub_env_vars(self, config): + frigate_env_vars = {k: v for k, v in os.environ.items() if k.startswith('FRIGATE_')} + + if 'password' in config['mqtt']: + config['mqtt']['password'] = config['mqtt']['password'].format(**frigate_env_vars) + + for camera in config['cameras'].values(): + camera['ffmpeg']['input'] = camera['ffmpeg']['input'].format(**frigate_env_vars) + + return config + + def _ensure_dirs(self): + cache_dir = self.save_clips.cache_dir + clips_dir = self.save_clips.clips_dir + + if not os.path.exists(cache_dir) and not os.path.islink(cache_dir): + os.makedirs(cache_dir) + if not os.path.exists(clips_dir) and not os.path.islink(clips_dir): + os.makedirs(clips_dir) + + def _load_file(self, config_file): + with open(config_file) as f: + raw_config = f.read() + + if config_file.endswith(".yml"): + config = yaml.safe_load(raw_config) + elif config_file.endswith(".json"): + config = json.loads(raw_config) + + return config + + @property + def web_port(self): + return self._web_port + + @property + def detectors(self) -> Dict[str, DetectorConfig]: + return self._detectors + + @property + def mqtt(self): + return self._mqtt + + @property + def save_clips(self): + return self._save_clips + + @property + def cameras(self) -> Dict[str, CameraConfig]: + return self._cameras \ No newline at end of file diff --git a/frigate/events.py b/frigate/events.py index 7dd4e4436..dc2743ea5 100644 --- a/frigate/events.py +++ b/frigate/events.py @@ -14,8 +14,8 @@ class EventProcessor(threading.Thread): def __init__(self, config, camera_processes, event_queue, stop_event): threading.Thread.__init__(self) self.config = config - self.cache_dir = self.config['save_clips']['cache_dir'] - self.clips_dir = self.config['save_clips']['clips_dir'] + self.cache_dir = self.config.save_clips.cache_dir + self.clips_dir = self.config.save_clips.clips_dir self.camera_processes = camera_processes self.cached_clips = {} self.event_queue = event_queue @@ -77,7 +77,7 @@ class EventProcessor(threading.Thread): earliest_event = datetime.datetime.now().timestamp() # if the earliest event exceeds the max seconds, cap it - max_seconds = self.config['save_clips']['max_seconds'] + max_seconds = self.config.save_clips.max_seconds if datetime.datetime.now().timestamp()-earliest_event > max_seconds: earliest_event = datetime.datetime.now().timestamp()-max_seconds @@ -163,15 +163,16 @@ class EventProcessor(threading.Thread): self.refresh_cache() - save_clips_config = self.config['cameras'][camera].get('save_clips', {}) + save_clips_config = self.config.cameras[camera].save_clips # if save clips is not enabled for this camera, just continue - if not save_clips_config.get('enabled', False): + if not save_clips_config.enabled: continue # if specific objects are listed for this camera, only save clips for them - if 'objects' in save_clips_config: - if not event_data['label'] in save_clips_config['objects']: + # TODO: default to all tracked objects rather than checking for None + if save_clips_config.objects: + if not event_data['label'] in save_clips_config.objects: continue if event_type == 'start': @@ -190,7 +191,7 @@ class EventProcessor(threading.Thread): ) if len(self.cached_clips) > 0 and not event_data['false_positive']: - self.create_clip(camera, event_data, save_clips_config.get('pre_capture', 30)) + self.create_clip(camera, event_data, save_clips_config.pre_capture) del self.events_in_process[event_data['id']] \ No newline at end of file diff --git a/frigate/http.py b/frigate/http.py index a24fe7387..58c0eee12 100644 --- a/frigate/http.py +++ b/frigate/http.py @@ -75,7 +75,7 @@ def stats(): @bp.route('//