make the threshold configurable per region. fixes #31

blakeblackshear · May 11, 2019 · 3019b02 · 3019b02
1 parent 6900e14
commit 3019b02
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 3 deletions.
diff --git a/config/config.yml b/config/config.yml
@@ -19,11 +19,14 @@ cameras:
         x_offset: 0
         y_offset: 300
         min_person_area: 5000
+        threshold: 0.5
       - size: 400
         x_offset: 350
         y_offset: 250
         min_person_area: 2000
+        threshold: 0.5
       - size: 400
         x_offset: 750
         y_offset: 250
         min_person_area: 2000
+        threshold: 0.5
diff --git a/frigate/object_detection.py b/frigate/object_detection.py
@@ -38,7 +38,7 @@ def run(self):
             frame = self.prepped_frame_queue.get()
 
             # Actual detection.
-            objects = self.engine.DetectWithInputTensor(frame['frame'], threshold=0.5, top_k=3)
+            objects = self.engine.DetectWithInputTensor(frame['frame'], threshold=frame['region_threshold'], top_k=3)
             # parse and pass detected objects back to the camera
             parsed_objects = []
             for obj in objects:
@@ -59,7 +59,7 @@ def run(self):
 class FramePrepper(threading.Thread):
     def __init__(self, camera_name, shared_frame, frame_time, frame_ready, 
         frame_lock,
-        region_size, region_x_offset, region_y_offset,
+        region_size, region_x_offset, region_y_offset, region_threshold,
         prepped_frame_queue):
 
         threading.Thread.__init__(self)
@@ -71,6 +71,7 @@ def __init__(self, camera_name, shared_frame, frame_time, frame_ready,
         self.region_size = region_size
         self.region_x_offset = region_x_offset
         self.region_y_offset = region_y_offset
+        self.region_threshold = region_threshold
         self.prepped_frame_queue = prepped_frame_queue
 
     def run(self):
@@ -103,6 +104,7 @@ def run(self):
                     'frame_time': frame_time,
                     'frame': frame_expanded.flatten().copy(),
                     'region_size': self.region_size,
+                    'region_threshold': self.region_threshold,
                     'region_x_offset': self.region_x_offset,
                     'region_y_offset': self.region_y_offset
                 })

diff --git a/frigate/video.py b/frigate/video.py
@@ -159,13 +159,19 @@ def __init__(self, name, config, prepped_frame_queue, mqtt_client, mqtt_prefix):
         # for each region, create a separate thread to resize the region and prep for detection
         self.detection_prep_threads = []
         for region in self.config['regions']:
+            # set a default threshold of 0.5 if not defined
+            if not 'threshold' in region:
+                region['threshold'] = 0.5
+            if not isinstance(region['threshold'], float):
+                print('Threshold is not a float. Setting to 0.5 default.')
+                region['threshold'] = 0.5
             self.detection_prep_threads.append(FramePrepper(
                 self.name,
                 self.shared_frame_np,
                 self.shared_frame_time,
                 self.frame_ready,
                 self.frame_lock,
-                region['size'], region['x_offset'], region['y_offset'],
+                region['size'], region['x_offset'], region['y_offset'], region['threshold'],
                 prepped_frame_queue
             ))