mirror of
https://github.com/blakeblackshear/frigate.git
synced 2024-11-21 19:07:46 +01:00
351 lines
16 KiB
Python
351 lines
16 KiB
Python
import json
|
|
import hashlib
|
|
import datetime
|
|
import time
|
|
import copy
|
|
import cv2
|
|
import threading
|
|
import queue
|
|
import numpy as np
|
|
from collections import Counter, defaultdict
|
|
import itertools
|
|
import pyarrow.plasma as plasma
|
|
import matplotlib.pyplot as plt
|
|
from frigate.util import draw_box_with_label, PlasmaFrameManager
|
|
from frigate.edgetpu import load_labels
|
|
from typing import Callable, Dict
|
|
from statistics import mean, median
|
|
|
|
PATH_TO_LABELS = '/labelmap.txt'
|
|
|
|
LABELS = load_labels(PATH_TO_LABELS)
|
|
cmap = plt.cm.get_cmap('tab10', len(LABELS.keys()))
|
|
|
|
COLOR_MAP = {}
|
|
for key, val in LABELS.items():
|
|
COLOR_MAP[val] = tuple(int(round(255 * c)) for c in cmap(key)[:3])
|
|
|
|
def zone_filtered(obj, object_config):
|
|
object_name = obj['label']
|
|
object_filters = object_config.get('filters', {})
|
|
|
|
if object_name in object_filters:
|
|
obj_settings = object_filters[object_name]
|
|
|
|
# if the min area is larger than the
|
|
# detected object, don't add it to detected objects
|
|
if obj_settings.get('min_area',-1) > obj['area']:
|
|
return True
|
|
|
|
# if the detected object is larger than the
|
|
# max area, don't add it to detected objects
|
|
if obj_settings.get('max_area', 24000000) < obj['area']:
|
|
return True
|
|
|
|
# if the score is lower than the threshold, skip
|
|
if obj_settings.get('threshold', 0) > obj['computed_score']:
|
|
return True
|
|
|
|
return False
|
|
|
|
# Maintains the state of a camera
|
|
class CameraState():
|
|
def __init__(self, name, config, frame_manager):
|
|
self.name = name
|
|
self.config = config
|
|
self.frame_manager = frame_manager
|
|
|
|
self.best_objects = {}
|
|
self.object_status = defaultdict(lambda: 'OFF')
|
|
self.tracked_objects = {}
|
|
self.zone_objects = defaultdict(lambda: [])
|
|
self.current_frame = np.zeros((720,1280,3), np.uint8)
|
|
self.current_frame_time = 0.0
|
|
self.previous_frame_id = None
|
|
self.callbacks = defaultdict(lambda: [])
|
|
|
|
def false_positive(self, obj):
|
|
threshold = self.config['objects'].get('filters', {}).get(obj['label'], {}).get('threshold', 0.85)
|
|
if obj['computed_score'] < threshold:
|
|
return True
|
|
return False
|
|
|
|
def compute_score(self, obj):
|
|
scores = obj['score_history'][:]
|
|
# pad with zeros if you dont have at least 3 scores
|
|
if len(scores) < 3:
|
|
scores += [0.0]*(3 - len(scores))
|
|
return median(scores)
|
|
|
|
def on(self, event_type: str, callback: Callable[[Dict], None]):
|
|
self.callbacks[event_type].append(callback)
|
|
|
|
def update(self, frame_time, tracked_objects):
|
|
self.current_frame_time = frame_time
|
|
# get the new frame and delete the old frame
|
|
frame_id = f"{self.name}{frame_time}"
|
|
self.current_frame = self.frame_manager.get(frame_id)
|
|
if not self.previous_frame_id is None:
|
|
self.frame_manager.delete(self.previous_frame_id)
|
|
self.previous_frame_id = frame_id
|
|
|
|
current_ids = tracked_objects.keys()
|
|
previous_ids = self.tracked_objects.keys()
|
|
removed_ids = list(set(previous_ids).difference(current_ids))
|
|
new_ids = list(set(current_ids).difference(previous_ids))
|
|
updated_ids = list(set(current_ids).intersection(previous_ids))
|
|
|
|
for id in new_ids:
|
|
self.tracked_objects[id] = tracked_objects[id]
|
|
self.tracked_objects[id]['zones'] = []
|
|
|
|
# start the score history
|
|
self.tracked_objects[id]['score_history'] = [self.tracked_objects[id]['score']]
|
|
|
|
# calculate if this is a false positive
|
|
self.tracked_objects[id]['computed_score'] = self.compute_score(self.tracked_objects[id])
|
|
self.tracked_objects[id]['false_positive'] = self.false_positive(self.tracked_objects[id])
|
|
|
|
# call event handlers
|
|
for c in self.callbacks['start']:
|
|
c(self.name, tracked_objects[id])
|
|
|
|
for id in updated_ids:
|
|
self.tracked_objects[id].update(tracked_objects[id])
|
|
|
|
# if the object is not in the current frame, add a 0.0 to the score history
|
|
if self.tracked_objects[id]['frame_time'] != self.current_frame_time:
|
|
self.tracked_objects[id]['score_history'].append(0.0)
|
|
else:
|
|
self.tracked_objects[id]['score_history'].append(self.tracked_objects[id]['score'])
|
|
# only keep the last 10 scores
|
|
if len(self.tracked_objects[id]['score_history']) > 10:
|
|
self.tracked_objects[id]['score_history'] = self.tracked_objects[id]['score_history'][-10:]
|
|
|
|
# calculate if this is a false positive
|
|
self.tracked_objects[id]['computed_score'] = self.compute_score(self.tracked_objects[id])
|
|
self.tracked_objects[id]['false_positive'] = self.false_positive(self.tracked_objects[id])
|
|
|
|
# call event handlers
|
|
for c in self.callbacks['update']:
|
|
c(self.name, self.tracked_objects[id])
|
|
|
|
for id in removed_ids:
|
|
# publish events to mqtt
|
|
self.tracked_objects[id]['end_time'] = frame_time
|
|
for c in self.callbacks['end']:
|
|
c(self.name, self.tracked_objects[id])
|
|
del self.tracked_objects[id]
|
|
|
|
# check to see if the objects are in any zones
|
|
for obj in self.tracked_objects.values():
|
|
current_zones = []
|
|
bottom_center = (obj['centroid'][0], obj['box'][3])
|
|
# check each zone
|
|
for name, zone in self.config['zones'].items():
|
|
contour = zone['contour']
|
|
# check if the object is in the zone and not filtered
|
|
if (cv2.pointPolygonTest(contour, bottom_center, False) >= 0
|
|
and not zone_filtered(obj, zone.get('filters', {}))):
|
|
current_zones.append(name)
|
|
obj['zones'] = current_zones
|
|
|
|
# draw on the frame
|
|
if not self.current_frame is None:
|
|
# draw the bounding boxes on the frame
|
|
for obj in self.tracked_objects.values():
|
|
thickness = 2
|
|
color = COLOR_MAP[obj['label']]
|
|
|
|
if obj['frame_time'] != frame_time:
|
|
thickness = 1
|
|
color = (255,0,0)
|
|
|
|
# draw the bounding boxes on the frame
|
|
box = obj['box']
|
|
draw_box_with_label(self.current_frame, box[0], box[1], box[2], box[3], obj['label'], f"{int(obj['score']*100)}% {int(obj['area'])}", thickness=thickness, color=color)
|
|
# draw the regions on the frame
|
|
region = obj['region']
|
|
cv2.rectangle(self.current_frame, (region[0], region[1]), (region[2], region[3]), (0,255,0), 1)
|
|
|
|
if self.config['snapshots']['show_timestamp']:
|
|
time_to_show = datetime.datetime.fromtimestamp(frame_time).strftime("%m/%d/%Y %H:%M:%S")
|
|
cv2.putText(self.current_frame, time_to_show, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, fontScale=.8, color=(255, 255, 255), thickness=2)
|
|
|
|
if self.config['snapshots']['draw_zones']:
|
|
for name, zone in self.config['zones'].items():
|
|
thickness = 8 if any([name in obj['zones'] for obj in self.tracked_objects.values()]) else 2
|
|
cv2.drawContours(self.current_frame, [zone['contour']], -1, zone['color'], thickness)
|
|
|
|
# maintain best objects
|
|
for obj in self.tracked_objects.values():
|
|
object_type = obj['label']
|
|
# if the object wasn't seen on the current frame, skip it
|
|
if obj['frame_time'] != self.current_frame_time or obj['false_positive']:
|
|
continue
|
|
if object_type in self.best_objects:
|
|
current_best = self.best_objects[object_type]
|
|
now = datetime.datetime.now().timestamp()
|
|
# if the object is a higher score than the current best score
|
|
# or the current object is more than 1 minute old, use the new object
|
|
if obj['score'] > current_best['score'] or (now - current_best['frame_time']) > 60:
|
|
obj['frame'] = np.copy(self.current_frame)
|
|
self.best_objects[object_type] = obj
|
|
for c in self.callbacks['snapshot']:
|
|
c(self.name, self.best_objects[object_type])
|
|
else:
|
|
obj['frame'] = np.copy(self.current_frame)
|
|
self.best_objects[object_type] = obj
|
|
for c in self.callbacks['snapshot']:
|
|
c(self.name, self.best_objects[object_type])
|
|
|
|
# update overall camera state for each object type
|
|
obj_counter = Counter()
|
|
for obj in self.tracked_objects.values():
|
|
if not obj['false_positive']:
|
|
obj_counter[obj['label']] += 1
|
|
|
|
# report on detected objects
|
|
for obj_name, count in obj_counter.items():
|
|
new_status = 'ON' if count > 0 else 'OFF'
|
|
if new_status != self.object_status[obj_name]:
|
|
self.object_status[obj_name] = new_status
|
|
for c in self.callbacks['object_status']:
|
|
c(self.name, obj_name, new_status)
|
|
|
|
# expire any objects that are ON and no longer detected
|
|
expired_objects = [obj_name for obj_name, status in self.object_status.items() if status == 'ON' and not obj_name in obj_counter]
|
|
for obj_name in expired_objects:
|
|
self.object_status[obj_name] = 'OFF'
|
|
for c in self.callbacks['object_status']:
|
|
c(self.name, obj_name, 'OFF')
|
|
for c in self.callbacks['snapshot']:
|
|
c(self.name, self.best_objects[object_type])
|
|
|
|
|
|
class TrackedObjectProcessor(threading.Thread):
|
|
def __init__(self, camera_config, zone_config, client, topic_prefix, tracked_objects_queue, event_queue, stop_event):
|
|
threading.Thread.__init__(self)
|
|
self.camera_config = camera_config
|
|
self.zone_config = zone_config
|
|
self.client = client
|
|
self.topic_prefix = topic_prefix
|
|
self.tracked_objects_queue = tracked_objects_queue
|
|
self.event_queue = event_queue
|
|
self.stop_event = stop_event
|
|
self.camera_states: Dict[str, CameraState] = {}
|
|
self.plasma_client = PlasmaFrameManager(self.stop_event)
|
|
|
|
def start(camera, obj):
|
|
# publish events to mqtt
|
|
self.client.publish(f"{self.topic_prefix}/{camera}/events/start", json.dumps({x: obj[x] for x in obj if x not in ['frame']}), retain=False)
|
|
self.event_queue.put(('start', camera, obj))
|
|
|
|
def update(camera, obj):
|
|
pass
|
|
|
|
def end(camera, obj):
|
|
self.client.publish(f"{self.topic_prefix}/{camera}/events/end", json.dumps({x: obj[x] for x in obj if x not in ['frame']}), retain=False)
|
|
self.event_queue.put(('end', camera, obj))
|
|
|
|
def snapshot(camera, obj):
|
|
best_frame = cv2.cvtColor(obj['frame'], cv2.COLOR_RGB2BGR)
|
|
ret, jpg = cv2.imencode('.jpg', best_frame)
|
|
if ret:
|
|
jpg_bytes = jpg.tobytes()
|
|
self.client.publish(f"{self.topic_prefix}/{camera}/{obj['label']}/snapshot", jpg_bytes, retain=True)
|
|
|
|
def object_status(camera, object_name, status):
|
|
self.client.publish(f"{self.topic_prefix}/{camera}/{object_name}", status, retain=False)
|
|
|
|
for camera in self.camera_config.keys():
|
|
camera_state = CameraState(camera, self.camera_config[camera], self.plasma_client)
|
|
camera_state.on('start', start)
|
|
camera_state.on('update', update)
|
|
camera_state.on('end', end)
|
|
camera_state.on('snapshot', snapshot)
|
|
camera_state.on('object_status', object_status)
|
|
self.camera_states[camera] = camera_state
|
|
|
|
self.camera_data = defaultdict(lambda: {
|
|
'best_objects': {},
|
|
'object_status': defaultdict(lambda: defaultdict(lambda: 'OFF')),
|
|
'tracked_objects': {},
|
|
'current_frame': np.zeros((720,1280,3), np.uint8),
|
|
'current_frame_time': 0.0,
|
|
'object_id': None
|
|
})
|
|
# {
|
|
# 'zone_name': {
|
|
# 'person': ['camera_1', 'camera_2']
|
|
# }
|
|
# }
|
|
self.zone_data = defaultdict(lambda: defaultdict(lambda: set()))
|
|
|
|
# set colors for zones
|
|
zone_colors = {}
|
|
colors = plt.cm.get_cmap('tab10', len(self.zone_config.keys()))
|
|
for i, zone in enumerate(self.zone_config.keys()):
|
|
zone_colors[zone] = tuple(int(round(255 * c)) for c in colors(i)[:3])
|
|
|
|
# create zone contours
|
|
for zone_name, config in zone_config.items():
|
|
for camera, camera_zone_config in config.items():
|
|
camera_zone = {}
|
|
camera_zone['color'] = zone_colors[zone_name]
|
|
coordinates = camera_zone_config['coordinates']
|
|
if isinstance(coordinates, list):
|
|
camera_zone['contour'] = np.array([[int(p.split(',')[0]), int(p.split(',')[1])] for p in coordinates])
|
|
elif isinstance(coordinates, str):
|
|
points = coordinates.split(',')
|
|
camera_zone['contour'] = np.array([[int(points[i]), int(points[i+1])] for i in range(0, len(points), 2)])
|
|
else:
|
|
print(f"Unable to parse zone coordinates for {zone_name} - {camera}")
|
|
self.camera_config[camera]['zones'][zone_name] = camera_zone
|
|
|
|
def get_best(self, camera, label):
|
|
best_objects = self.camera_states[camera].best_objects
|
|
if label in best_objects:
|
|
return best_objects[label]['frame']
|
|
else:
|
|
return None
|
|
|
|
def get_current_frame(self, camera):
|
|
return self.camera_states[camera].current_frame
|
|
|
|
def run(self):
|
|
while True:
|
|
if self.stop_event.is_set():
|
|
print(f"Exiting object processor...")
|
|
break
|
|
|
|
try:
|
|
camera, frame_time, current_tracked_objects = self.tracked_objects_queue.get(True, 10)
|
|
except queue.Empty:
|
|
continue
|
|
|
|
camera_state = self.camera_states[camera]
|
|
|
|
camera_state.update(frame_time, current_tracked_objects)
|
|
|
|
# update zone status for each label
|
|
for zone in camera_state.config['zones'].keys():
|
|
# get labels for current camera and all labels in current zone
|
|
labels_for_camera = set([obj['label'] for obj in camera_state.tracked_objects.values() if zone in obj['zones']])
|
|
labels_to_check = labels_for_camera | set(self.zone_data[zone].keys())
|
|
# for each label in zone
|
|
for label in labels_to_check:
|
|
camera_list = self.zone_data[zone][label]
|
|
# remove or add the camera to the list for the current label
|
|
previous_state = len(camera_list) > 0
|
|
if label in labels_for_camera:
|
|
camera_list.add(camera_state.name)
|
|
elif camera_state.name in camera_list:
|
|
camera_list.remove(camera_state.name)
|
|
new_state = len(camera_list) > 0
|
|
# if the value is changing, send over MQTT
|
|
if previous_state == False and new_state == True:
|
|
self.client.publish(f"{self.topic_prefix}/{zone}/{label}", 'ON', retain=False)
|
|
elif previous_state == True and new_state == False:
|
|
self.client.publish(f"{self.topic_prefix}/{zone}/{label}", 'OFF', retain=False)
|