apps/frigate-rockchip
1
0
Fork 0
mirror of https://github.com/MarcA711/frigate-rockchip.git synced 2025-10-15 05:10:39 +08:00
Code Issues Packages Projects Releases Wiki Activity

track and report all detected object types

Browse Source
This commit is contained in:
Blake Blackshear
2019-12-14 15:18:21 -06:00
parent 5c01720567
commit bee99ca6ff
7 changed files with 160 additions and 131 deletions
Download Patch File Download Diff File Expand all files Collapse all files

49
frigate/mqtt.py
View File

@@ -1,41 +1,46 @@
import json
import cv2
import threading
from collections import Counter, defaultdict
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_frames):
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.best_frames = best_frames
def run(self):
last_sent_payload = ""
current_object_status = defaultdict(lambda: 'OFF')
while True:
# initialize the payload
payload = {}
# wait until objects have been parsed
with self.objects_parsed:
self.objects_parsed.wait()
# add all the person scores in detected objects
# make a copy of detected objects
detected_objects = self._detected_objects.copy()
person_score = sum([obj['score'] for obj in detected_objects if obj['name'] == 'person'])
# if the person score is more than 100, set person to ON
payload['person'] = 'ON' if int(person_score*100) > 100 else 'OFF'
# send message for objects if different
new_payload = json.dumps(payload, sort_keys=True)
if new_payload != last_sent_payload:
last_sent_payload = new_payload
self.client.publish(self.topic_prefix+'/objects', new_payload, retain=False)
# send the snapshot over mqtt as well
if not self.best_person_frame.best_frame is None:
ret, jpg = cv2.imencode('.jpg', self.best_person_frame.best_frame)
if ret:
jpg_bytes = jpg.tobytes()
self.client.publish(self.topic_prefix+'/snapshot', jpg_bytes, retain=True)
# total up all scores by object type
obj_counter = Counter()
for obj in detected_objects:
obj_counter[obj['name']] += obj['score']
# report on detected objects
for obj_name, total_score in obj_counter.items():
new_status = 'ON' if int(total_score*100) > 100 else 'OFF'
if new_status != current_object_status[obj_name]:
current_object_status[obj_name] = new_status
self.client.publish(self.topic_prefix+'/'+obj_name, new_status, retain=False)
# send the snapshot over mqtt as well
if not self.best_frames.best_frames[obj_name] is None:
ret, jpg = cv2.imencode('.jpg', self.best_frames.best_frames[obj_name])
if ret:
jpg_bytes = jpg.tobytes()
self.client.publish(self.topic_prefix+'/'+obj_name+'/snapshot', jpg_bytes, retain=True)
# expire any objects that are ON and no longer detected
expired_objects = [obj_name for obj_name, status in current_object_status.items() if status == 'ON' and not obj_name in obj_counter]
for obj_name in expired_objects:
self.client.publish(self.topic_prefix+'/'+obj_name, 'OFF', retain=False)