diff --git a/frigate/mqtt.py b/frigate/mqtt.py index f6a9e57a4..155c28580 100644 --- a/frigate/mqtt.py +++ b/frigate/mqtt.py @@ -3,13 +3,14 @@ import cv2 import threading class MqttObjectPublisher(threading.Thread): - def __init__(self, client, topic_prefix, objects_parsed, detected_objects, best_person_frame): + def __init__(self, client, topic_prefix, objects_parsed, detected_objects, best_person_frame, label): threading.Thread.__init__(self) self.client = client self.topic_prefix = topic_prefix self.objects_parsed = objects_parsed self._detected_objects = detected_objects self.best_person_frame = best_person_frame + self.label = label def run(self): last_sent_payload = "" @@ -24,7 +25,7 @@ class MqttObjectPublisher(threading.Thread): # add all the person scores in detected objects detected_objects = self._detected_objects.copy() - person_score = sum([obj['score'] for obj in detected_objects if obj['name'] == 'person']) + person_score = sum([obj['score'] for obj in detected_objects if obj['name'] == self.label]) # if the person score is more than 100, set person to ON payload['person'] = 'ON' if int(person_score*100) > 100 else 'OFF' diff --git a/frigate/objects.py b/frigate/objects.py index 1ae63b3f3..0a7043a26 100644 --- a/frigate/objects.py +++ b/frigate/objects.py @@ -38,11 +38,12 @@ class ObjectCleaner(threading.Thread): # Maintains the frame and person with the highest score from the most recent # motion event class BestPersonFrame(threading.Thread): - def __init__(self, objects_parsed, recent_frames, detected_objects): + def __init__(self, objects_parsed, recent_frames, detected_objects, label): threading.Thread.__init__(self) self.objects_parsed = objects_parsed self.recent_frames = recent_frames self.detected_objects = detected_objects + self.label = label self.best_person = None self.best_frame = None @@ -55,7 +56,7 @@ class BestPersonFrame(threading.Thread): # make a copy of detected objects detected_objects = self.detected_objects.copy() - detected_people = [obj for obj in detected_objects if obj['name'] == 'person'] + detected_people = [obj for obj in detected_objects if obj['name'] == self.label] # get the highest scoring person new_best_person = max(detected_people, key=lambda x:x['score'], default=self.best_person) diff --git a/frigate/video.py b/frigate/video.py index 7a1d7ae70..54c7c532f 100644 --- a/frigate/video.py +++ b/frigate/video.py @@ -130,6 +130,7 @@ class Camera: self.frame_size = self.frame_shape[0] * self.frame_shape[1] * self.frame_shape[2] self.mqtt_client = mqtt_client self.mqtt_topic_prefix = '{}/{}'.format(mqtt_prefix, self.name) + self.label = config.get('label', 'person') # create a numpy array for the current frame in initialize to zeros self.current_frame = np.zeros(self.frame_shape, np.uint8) @@ -170,7 +171,7 @@ class Camera: self.frame_tracker.start() # start a thread to store the highest scoring recent person frame - self.best_person_frame = BestPersonFrame(self.objects_parsed, self.recent_frames, self.detected_objects) + self.best_person_frame = BestPersonFrame(self.objects_parsed, self.recent_frames, self.detected_objects, self.label) self.best_person_frame.start() # start a thread to expire objects from the detected objects list @@ -178,7 +179,7 @@ class Camera: self.object_cleaner.start() # start a thread to publish object scores (currently only person) - mqtt_publisher = MqttObjectPublisher(self.mqtt_client, self.mqtt_topic_prefix, self.objects_parsed, self.detected_objects, self.best_person_frame) + mqtt_publisher = MqttObjectPublisher(self.mqtt_client, self.mqtt_topic_prefix, self.objects_parsed, self.detected_objects, self.best_person_frame, self.label) mqtt_publisher.start() # create a watchdog thread for capture process @@ -255,7 +256,7 @@ class Camera: # Store object area to use in bounding box labels obj['area'] = (obj['xmax']-obj['xmin'])*(obj['ymax']-obj['ymin']) - if obj['name'] == 'person': + if obj['name'] == self.label: # find the matching region region = None for r in self.regions: