add global object mask

This commit is contained in:
Blake Blackshear 2021-02-06 06:30:26 -06:00
parent a99f360a64
commit ce90ae343c
3 changed files with 60 additions and 3 deletions

View File

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

View File

@ -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() }
}

View File

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