From ce90ae343c448d42965b270edad3d4edb0cc67e2 Mon Sep 17 00:00:00 2001
From: Blake Blackshear <blakeb@blakeshome.com>
Date: Sat, 6 Feb 2021 06:30:26 -0600
Subject: [PATCH] add global object mask

---
 docs/docs/configuration/cameras.md |  8 ++++++++
 frigate/config.py                  | 22 +++++++++++++++++---
 frigate/test/test_config.py        | 33 ++++++++++++++++++++++++++++++
 3 files changed, 60 insertions(+), 3 deletions(-)

diff --git a/docs/docs/configuration/cameras.md b/docs/docs/configuration/cameras.md
index 0c306e20c..7a822dc6b 100644
--- a/docs/docs/configuration/cameras.md
+++ b/docs/docs/configuration/cameras.md
@@ -108,6 +108,10 @@ objects:
   track:
     - person
     - car
+  # Optional: mask to prevent all object types from being detected in certain areas (default: no mask)
+  # Checks based on the bottom center of the bounding box of the object. 
+  # NOTE: This mask is COMBINED with the object type specific mask below
+  mask: 0,0,1000,0,1000,200,0,200
   filters:
     person:
       min_area: 5000
@@ -361,6 +365,10 @@ cameras:
       track:
         - person
         - car
+      # Optional: mask to prevent all object types from being detected in certain areas (default: no mask)
+      # Checks based on the bottom center of the bounding box of the object. 
+      # NOTE: This mask is COMBINED with the object type specific mask below
+      mask: 0,0,1000,0,1000,200,0,200
       filters:
         person:
           min_area: 5000
diff --git a/frigate/config.py b/frigate/config.py
index 4ea1fb6c3..0878f3c1f 100644
--- a/frigate/config.py
+++ b/frigate/config.py
@@ -131,6 +131,7 @@ def filters_for_all_tracked_objects(object_config):
 OBJECTS_SCHEMA = vol.Schema(vol.All(filters_for_all_tracked_objects,
     {
         'track': [str],
+        'mask': vol.Any(str, [str]),
         vol.Optional('filters', default = {}): FILTER_SCHEMA.extend(
             { 
                 str: {
@@ -512,12 +513,25 @@ class RecordConfig():
         }
 
 class FilterConfig():
-    def __init__(self, global_config, config, frame_shape=None):
+    def __init__(self, global_config, config, global_mask=None, frame_shape=None):
         self._min_area = config.get('min_area', global_config.get('min_area', 0))
         self._max_area = config.get('max_area', global_config.get('max_area', 24000000))
         self._threshold = config.get('threshold', global_config.get('threshold', 0.7))
         self._min_score = config.get('min_score', global_config.get('min_score', 0.5))
-        self._raw_mask = config.get('mask')
+
+        self._raw_mask = []
+        if global_mask:
+            if isinstance(global_mask, list):
+                self._raw_mask += global_mask
+            elif isinstance(global_mask, str):
+                self._raw_mask += [global_mask]
+
+        mask = config.get('mask')
+        if mask:
+            if isinstance(mask, list):
+                self._raw_mask += mask
+            elif isinstance(mask, str):
+                self._raw_mask += [mask]
         self._mask = create_mask(frame_shape, self._raw_mask) if self._raw_mask else None
 
     @property
@@ -552,7 +566,8 @@ class FilterConfig():
 class ObjectConfig():
     def __init__(self, global_config, config, frame_shape):
         self._track = config.get('track', global_config.get('track', DEFAULT_TRACKED_OBJECTS))
-        self._filters = { name: FilterConfig(global_config.get('filters').get(name, {}), config.get('filters').get(name, {}), frame_shape) for name in self._track }
+        self._raw_mask = config.get('mask')
+        self._filters = { name: FilterConfig(global_config['filters'].get(name, {}), config['filters'].get(name, {}), self._raw_mask, frame_shape) for name in self._track }
 
     @property
     def track(self):
@@ -565,6 +580,7 @@ class ObjectConfig():
     def to_dict(self):
         return {
             'track': self.track,
+            'mask': self._raw_mask,
             'filters': { k: f.to_dict() for k, f in self.filters.items() }
         }
 
diff --git a/frigate/test/test_config.py b/frigate/test/test_config.py
index c906a7c0c..b18c858ff 100644
--- a/frigate/test/test_config.py
+++ b/frigate/test/test_config.py
@@ -160,6 +160,39 @@ class TestConfig(TestCase):
         assert('dog' in frigate_config.cameras['back'].objects.filters)
         assert(frigate_config.cameras['back'].objects.filters['dog'].threshold == 0.7)
     
+    def test_global_object_mask(self):
+        config = {
+            'mqtt': {
+                'host': 'mqtt'
+            },
+            'objects': {
+                'track': ['person', 'dog']
+            },
+            'cameras': {
+                'back': {
+                    'ffmpeg': {
+                        'inputs': [
+                            { 'path': 'rtsp://10.0.0.1:554/video', 'roles': ['detect'] }
+                        ]
+                    },
+                    'height': 1080,
+                    'width': 1920,
+                    'objects': {
+                        'mask': '0,0,1,1,0,1',
+                        'filters': {
+                            'dog': {
+                                'mask': '1,1,1,1,1,1'
+                            }
+                        }
+                    }
+                }
+            }
+        }
+        frigate_config = FrigateConfig(config=config)
+        assert('dog' in frigate_config.cameras['back'].objects.filters)
+        assert(len(frigate_config.cameras['back'].objects.filters['dog']._raw_mask) == 2)
+        assert(len(frigate_config.cameras['back'].objects.filters['person']._raw_mask) == 1)
+    
     def test_ffmpeg_params_global(self):
         config = {
             'ffmpeg': {