blakeblackshear.frigate/frigate/object_processing.py

501 lines
22 KiB
Python
Raw Normal View History

2020-02-16 04:07:54 +01:00
import copy
import base64
2020-11-04 13:31:25 +01:00
import datetime
import hashlib
import itertools
import json
2020-11-04 04:26:39 +01:00
import logging
import os
2020-08-02 15:46:36 +02:00
import queue
2020-11-04 13:31:25 +01:00
import threading
import time
2020-02-16 04:07:54 +01:00
from collections import Counter, defaultdict
2020-11-04 13:31:25 +01:00
from statistics import mean, median
from typing import Callable, Dict
import cv2
2020-02-16 04:07:54 +01:00
import matplotlib.pyplot as plt
2020-11-04 13:31:25 +01:00
import numpy as np
from frigate.config import FrigateConfig, CameraConfig
2020-11-04 13:31:25 +01:00
from frigate.edgetpu import load_labels
from frigate.util import SharedMemoryFrameManager, draw_box_with_label
2020-02-16 04:07:54 +01:00
2020-11-04 04:26:39 +01:00
logger = logging.getLogger(__name__)
2020-02-18 12:55:06 +01:00
PATH_TO_LABELS = '/labelmap.txt'
2020-02-16 04:07:54 +01:00
2020-02-16 15:49:43 +01:00
LABELS = load_labels(PATH_TO_LABELS)
2020-02-16 04:07:54 +01:00
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 on_edge(box, frame_shape):
if (
box[0] == 0 or
box[1] == 0 or
box[2] == frame_shape[1]-1 or
box[3] == frame_shape[0]-1
):
return True
def is_better_thumbnail(current_thumb, new_obj, frame_shape) -> bool:
# larger is better
# cutoff images are less ideal, but they should also be smaller?
# better scores are obviously better too
# if the new_thumb is on an edge, and the current thumb is not
if on_edge(new_obj['box'], frame_shape) and not on_edge(current_thumb['box'], frame_shape):
return False
# if the score is better by more than 5%
if new_obj['score'] > current_thumb['score']+.05:
return True
# if the area is 10% larger
if new_obj['area'] > current_thumb['area']*1.1:
return True
return False
class TrackedObject():
2020-11-11 23:55:50 +01:00
def __init__(self, camera, camera_config: CameraConfig, frame_cache, obj_data):
self.obj_data = obj_data
self.camera = camera
self.camera_config = camera_config
2020-11-11 23:55:50 +01:00
self.frame_cache = frame_cache
self.current_zones = []
self.entered_zones = set()
2020-11-10 04:11:27 +01:00
self.false_positive = True
self.top_score = self.computed_score = 0.0
self.thumbnail_data = None
self.frame = None
2020-11-25 19:06:01 +01:00
self.previous = None
self._snapshot_jpg_time = 0
ret, jpg = cv2.imencode('.jpg', np.zeros((300,300,3), np.uint8))
self._snapshot_jpg = jpg.tobytes()
# start the score history
self.score_history = [self.obj_data['score']]
2020-11-10 04:11:27 +01:00
def _is_false_positive(self):
# once a true positive, always a true positive
2020-11-10 04:11:27 +01:00
if not self.false_positive:
return False
threshold = self.camera_config.objects.filters[self.obj_data['label']].threshold
if self.computed_score < threshold:
return True
return False
def compute_score(self):
scores = self.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 update(self, current_frame_time, obj_data):
2020-11-25 19:06:01 +01:00
previous = self.to_dict()
self.obj_data.update(obj_data)
# if the object is not in the current frame, add a 0.0 to the score history
if self.obj_data['frame_time'] != current_frame_time:
self.score_history.append(0.0)
else:
self.score_history.append(self.obj_data['score'])
# only keep the last 10 scores
if len(self.score_history) > 10:
self.score_history = self.score_history[-10:]
# calculate if this is a false positive
self.computed_score = self.compute_score()
if self.computed_score > self.top_score:
self.top_score = self.computed_score
2020-11-10 04:11:27 +01:00
self.false_positive = self._is_false_positive()
if not self.false_positive:
# determine if this frame is a better thumbnail
2020-11-25 19:06:01 +01:00
if (
self.thumbnail_data is None
or is_better_thumbnail(self.thumbnail_data, self.obj_data, self.camera_config.frame_shape)
):
self.thumbnail_data = {
'frame_time': self.obj_data['frame_time'],
'box': self.obj_data['box'],
'area': self.obj_data['area'],
'region': self.obj_data['region'],
'score': self.obj_data['score']
}
2020-11-25 19:06:01 +01:00
self.previous = previous
# check zones
current_zones = []
bottom_center = (self.obj_data['centroid'][0], self.obj_data['box'][3])
# check each zone
for name, zone in self.camera_config.zones.items():
contour = zone.contour
# check if the object is in the zone
if (cv2.pointPolygonTest(contour, bottom_center, False) >= 0):
# if the object passed the filters once, dont apply again
if name in self.current_zones or not zone_filtered(self, zone.filters):
current_zones.append(name)
self.entered_zones.add(name)
self.current_zones = current_zones
def to_dict(self, include_thumbnail: bool = False):
return {
'id': self.obj_data['id'],
'camera': self.camera,
'frame_time': self.obj_data['frame_time'],
'label': self.obj_data['label'],
'top_score': self.top_score,
2020-11-10 04:11:27 +01:00
'false_positive': self.false_positive,
'start_time': self.obj_data['start_time'],
'end_time': self.obj_data.get('end_time', None),
'score': self.obj_data['score'],
'box': self.obj_data['box'],
'area': self.obj_data['area'],
'region': self.obj_data['region'],
'current_zones': self.current_zones.copy(),
'entered_zones': list(self.entered_zones).copy(),
'thumbnail': base64.b64encode(self.get_jpg_bytes()).decode('utf-8') if include_thumbnail else None
}
2020-11-12 00:44:51 +01:00
def get_jpg_bytes(self):
2020-11-15 00:30:24 +01:00
if self.thumbnail_data is None or self._snapshot_jpg_time == self.thumbnail_data['frame_time']:
return self._snapshot_jpg
2020-11-11 23:55:50 +01:00
if not self.thumbnail_data['frame_time'] in self.frame_cache:
logger.error(f"Unable to create thumbnail for {self.obj_data['id']}")
logger.error(f"Looking for frame_time of {self.thumbnail_data['frame_time']}")
2020-11-11 23:55:50 +01:00
logger.error(f"Thumbnail frames: {','.join([str(k) for k in self.frame_cache.keys()])}")
return self._snapshot_jpg
# TODO: crop first to avoid converting the entire frame?
snapshot_config = self.camera_config.snapshots
2020-11-11 23:55:50 +01:00
best_frame = cv2.cvtColor(self.frame_cache[self.thumbnail_data['frame_time']], cv2.COLOR_YUV2BGR_I420)
if snapshot_config.draw_bounding_boxes:
thickness = 2
color = COLOR_MAP[self.obj_data['label']]
box = self.thumbnail_data['box']
draw_box_with_label(best_frame, box[0], box[1], box[2], box[3], self.obj_data['label'],
f"{int(self.thumbnail_data['score']*100)}% {int(self.thumbnail_data['area'])}", thickness=thickness, color=color)
if snapshot_config.crop_to_region:
region = self.thumbnail_data['region']
best_frame = best_frame[region[1]:region[3], region[0]:region[2]]
if snapshot_config.height:
height = snapshot_config.height
width = int(height*best_frame.shape[1]/best_frame.shape[0])
best_frame = cv2.resize(best_frame, dsize=(width, height), interpolation=cv2.INTER_AREA)
if snapshot_config.show_timestamp:
time_to_show = datetime.datetime.fromtimestamp(self.thumbnail_data['frame_time']).strftime("%m/%d/%Y %H:%M:%S")
size = cv2.getTextSize(time_to_show, cv2.FONT_HERSHEY_SIMPLEX, fontScale=1, thickness=2)
text_width = size[0][0]
desired_size = max(200, 0.33*best_frame.shape[1])
font_scale = desired_size/text_width
cv2.putText(best_frame, time_to_show, (5, best_frame.shape[0]-7), cv2.FONT_HERSHEY_SIMPLEX,
fontScale=font_scale, color=(255, 255, 255), thickness=2)
ret, jpg = cv2.imencode('.jpg', best_frame)
if ret:
self._snapshot_jpg = jpg.tobytes()
return self._snapshot_jpg
def zone_filtered(obj: TrackedObject, object_config):
object_name = obj.obj_data['label']
if object_name in object_config:
obj_settings = object_config[object_name]
# if the min area is larger than the
# detected object, don't add it to detected objects
if obj_settings.min_area > obj.obj_data['area']:
return True
# if the detected object is larger than the
# max area, don't add it to detected objects
if obj_settings.max_area < obj.obj_data['area']:
return True
# if the score is lower than the threshold, skip
if obj_settings.threshold > obj.computed_score:
return True
return False
2020-09-07 19:17:42 +02:00
# Maintains the state of a camera
class CameraState():
def __init__(self, name, config, frame_manager):
self.name = name
self.config = config
self.camera_config = config.cameras[name]
2020-09-07 19:17:42 +02:00
self.frame_manager = frame_manager
2020-11-10 04:31:45 +01:00
self.best_objects: Dict[str, TrackedObject] = {}
2020-09-07 19:17:42 +02:00
self.object_status = defaultdict(lambda: 'OFF')
self.tracked_objects: Dict[str, TrackedObject] = {}
2020-11-11 23:55:50 +01:00
self.frame_cache = {}
2020-09-07 19:17:42 +02:00
self.zone_objects = defaultdict(lambda: [])
self._current_frame = np.zeros(self.camera_config.frame_shape_yuv, np.uint8)
self.current_frame_lock = threading.Lock()
2020-09-07 19:17:42 +02:00
self.current_frame_time = 0.0
self.previous_frame_id = None
self.callbacks = defaultdict(lambda: [])
2020-10-11 19:16:05 +02:00
def get_current_frame(self, draw=False):
with self.current_frame_lock:
2020-10-11 19:16:05 +02:00
frame_copy = np.copy(self._current_frame)
frame_time = self.current_frame_time
tracked_objects = {k: v.to_dict() for k,v in self.tracked_objects.items()}
2020-10-11 19:16:05 +02:00
frame_copy = cv2.cvtColor(frame_copy, cv2.COLOR_YUV2BGR_I420)
# draw on the frame
if draw:
# draw the bounding boxes on the frame
for obj in 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(frame_copy, 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(frame_copy, (region[0], region[1]), (region[2], region[3]), (0,255,0), 1)
if self.camera_config.snapshots.show_timestamp:
2020-10-11 19:16:05 +02:00
time_to_show = datetime.datetime.fromtimestamp(frame_time).strftime("%m/%d/%Y %H:%M:%S")
cv2.putText(frame_copy, time_to_show, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, fontScale=.8, color=(255, 255, 255), thickness=2)
if self.camera_config.snapshots.draw_zones:
for name, zone in self.camera_config.zones.items():
thickness = 8 if any([name in obj['current_zones'] for obj in tracked_objects.values()]) else 2
2020-11-03 15:15:58 +01:00
cv2.drawContours(frame_copy, [zone.contour], -1, zone.color, thickness)
2020-10-11 19:16:05 +02:00
return frame_copy
2020-11-25 17:37:41 +01:00
def finished(self, obj_id):
del self.tracked_objects[obj_id]
2020-09-07 19:17:42 +02:00
def on(self, event_type: str, callback: Callable[[Dict], None]):
self.callbacks[event_type].append(callback)
2020-11-11 23:55:50 +01:00
def update(self, frame_time, current_detections):
2020-09-07 19:17:42 +02:00
self.current_frame_time = frame_time
# get the new frame
2020-09-07 19:17:42 +02:00
frame_id = f"{self.name}{frame_time}"
current_frame = self.frame_manager.get(frame_id, self.camera_config.frame_shape_yuv)
2020-09-07 19:17:42 +02:00
2020-11-11 23:55:50 +01:00
current_ids = current_detections.keys()
2020-09-07 19:17:42 +02:00
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:
2020-11-11 23:55:50 +01:00
new_obj = self.tracked_objects[id] = TrackedObject(self.name, self.camera_config, self.frame_cache, current_detections[id])
2020-09-07 19:17:42 +02:00
# call event handlers
for c in self.callbacks['start']:
2020-11-25 19:06:01 +01:00
c(self.name, new_obj, frame_time)
2020-09-07 19:17:42 +02:00
for id in updated_ids:
updated_obj = self.tracked_objects[id]
2020-11-11 23:55:50 +01:00
updated_obj.update(frame_time, current_detections[id])
2020-11-10 04:11:27 +01:00
if (not updated_obj.false_positive
and updated_obj.thumbnail_data['frame_time'] == frame_time
2020-11-11 23:55:50 +01:00
and frame_time not in self.frame_cache):
self.frame_cache[frame_time] = np.copy(current_frame)
2020-09-07 19:17:42 +02:00
# call event handlers
for c in self.callbacks['update']:
2020-11-25 19:06:01 +01:00
c(self.name, updated_obj, frame_time)
2020-09-07 19:17:42 +02:00
for id in removed_ids:
# publish events to mqtt
removed_obj = self.tracked_objects[id]
2020-11-25 17:37:41 +01:00
if not 'end_time' in removed_obj.obj_data:
removed_obj.obj_data['end_time'] = frame_time
for c in self.callbacks['end']:
2020-11-25 19:06:01 +01:00
c(self.name, removed_obj, frame_time)
2020-09-07 19:17:42 +02:00
# TODO: can i switch to looking this up and only changing when an event ends?
# maybe make an api endpoint that pulls the thumbnail from the file system?
2020-09-07 19:17:42 +02:00
# maintain best objects
for obj in self.tracked_objects.values():
object_type = obj.obj_data['label']
# if the object's thumbnail is not from the current frame
2020-11-15 00:30:24 +01:00
if obj.false_positive or obj.thumbnail_data['frame_time'] != self.current_frame_time:
2020-09-07 19:17:42 +02:00
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 older than desired, use the new object
2020-11-10 04:31:45 +01:00
if (is_better_thumbnail(current_best.thumbnail_data, obj.thumbnail_data, self.camera_config.frame_shape)
or (now - current_best.thumbnail_data['frame_time']) > self.camera_config.best_image_timeout):
self.best_objects[object_type] = obj
2020-09-07 19:17:42 +02:00
for c in self.callbacks['snapshot']:
2020-11-25 19:06:01 +01:00
c(self.name, self.best_objects[object_type], frame_time)
2020-09-07 19:17:42 +02:00
else:
2020-11-10 04:31:45 +01:00
self.best_objects[object_type] = obj
2020-09-07 19:17:42 +02:00
for c in self.callbacks['snapshot']:
2020-11-25 19:06:01 +01:00
c(self.name, self.best_objects[object_type], frame_time)
2020-09-07 19:17:42 +02:00
# 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.obj_data['label']] += 1
2020-09-07 19:17:42 +02:00
# 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']:
2020-11-25 19:06:01 +01:00
c(self.name, self.best_objects[obj_name], frame_time)
2020-10-11 19:16:05 +02:00
# cleanup thumbnail frame cache
2020-11-10 04:11:27 +01:00
current_thumb_frames = set([obj.thumbnail_data['frame_time'] for obj in self.tracked_objects.values() if not obj.false_positive])
2020-11-11 23:55:50 +01:00
current_best_frames = set([obj.thumbnail_data['frame_time'] for obj in self.best_objects.values()])
thumb_frames_to_delete = [t for t in self.frame_cache.keys() if not t in current_thumb_frames and not t in current_best_frames]
for t in thumb_frames_to_delete:
2020-11-11 23:55:50 +01:00
del self.frame_cache[t]
2020-10-11 19:16:05 +02:00
with self.current_frame_lock:
self._current_frame = current_frame
if not self.previous_frame_id is None:
self.frame_manager.delete(self.previous_frame_id)
self.previous_frame_id = frame_id
2020-09-07 19:17:42 +02:00
2020-02-16 04:07:54 +01:00
class TrackedObjectProcessor(threading.Thread):
2020-11-25 17:37:41 +01:00
def __init__(self, config: FrigateConfig, client, topic_prefix, tracked_objects_queue, event_queue, event_processed_queue, stop_event):
2020-02-16 04:07:54 +01:00
threading.Thread.__init__(self)
2020-11-04 13:28:07 +01:00
self.name = "detected_frames_processor"
self.config = config
2020-02-16 04:07:54 +01:00
self.client = client
self.topic_prefix = topic_prefix
self.tracked_objects_queue = tracked_objects_queue
2020-07-09 13:57:16 +02:00
self.event_queue = event_queue
2020-11-25 17:37:41 +01:00
self.event_processed_queue = event_processed_queue
2020-08-02 15:46:36 +02:00
self.stop_event = stop_event
2020-09-07 19:17:42 +02:00
self.camera_states: Dict[str, CameraState] = {}
self.frame_manager = SharedMemoryFrameManager()
2020-09-07 19:17:42 +02:00
2020-11-25 19:06:01 +01:00
def start(camera, obj: TrackedObject, current_frame_time):
self.event_queue.put(('start', camera, obj.to_dict()))
2020-09-07 19:17:42 +02:00
2020-11-25 19:06:01 +01:00
def update(camera, obj: TrackedObject, current_frame_time):
if not obj.thumbnail_data is None and obj.thumbnail_data['frame_time'] == current_frame_time:
message = { 'before': obj.previous, 'after': obj.to_dict() }
self.client.publish(f"{self.topic_prefix}/events", json.dumps(message), retain=False)
2020-09-07 19:17:42 +02:00
2020-11-25 19:06:01 +01:00
def end(camera, obj: TrackedObject, current_frame_time):
message = { 'before': obj.previous, 'after': obj.to_dict() }
self.client.publish(f"{self.topic_prefix}/events", json.dumps(message), retain=False)
2020-11-10 04:11:27 +01:00
if self.config.cameras[camera].save_clips.enabled and not obj.false_positive:
thumbnail_file_name = f"{camera}-{obj.obj_data['id']}.jpg"
with open(os.path.join(self.config.save_clips.clips_dir, thumbnail_file_name), 'wb') as f:
2020-11-12 00:44:51 +01:00
f.write(obj.get_jpg_bytes())
if not obj.false_positive:
self.event_queue.put(('end', camera, obj.to_dict(include_thumbnail=True)))
2020-09-07 19:17:42 +02:00
2020-11-25 19:06:01 +01:00
def snapshot(camera, obj: TrackedObject, current_frame_time):
2020-11-12 00:44:51 +01:00
self.client.publish(f"{self.topic_prefix}/{camera}/{obj.obj_data['label']}/snapshot", obj.get_jpg_bytes(), retain=True)
2020-09-07 19:17:42 +02:00
def object_status(camera, object_name, status):
self.client.publish(f"{self.topic_prefix}/{camera}/{object_name}", status, retain=False)
for camera in self.config.cameras.keys():
camera_state = CameraState(camera, self.config, self.frame_manager)
2020-09-07 19:17:42 +02:00
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
# {
# 'zone_name': {
# 'person': ['camera_1', 'camera_2']
# }
# }
self.zone_data = defaultdict(lambda: defaultdict(lambda: set()))
2020-02-16 04:07:54 +01:00
def get_best(self, camera, label):
2020-11-12 00:44:51 +01:00
# TODO: need a lock here
camera_state = self.camera_states[camera]
if label in camera_state.best_objects:
best_obj = camera_state.best_objects[label]
best = best_obj.to_dict()
best['frame'] = camera_state.frame_cache[best_obj.thumbnail_data['frame_time']]
return best
2020-02-16 04:07:54 +01:00
else:
return {}
2020-02-16 04:07:54 +01:00
2020-10-11 19:16:05 +02:00
def get_current_frame(self, camera, draw=False):
return self.camera_states[camera].get_current_frame(draw)
2020-02-16 04:07:54 +01:00
def run(self):
while True:
2020-08-02 15:46:36 +02:00
if self.stop_event.is_set():
2020-11-04 04:26:39 +01:00
logger.info(f"Exiting object processor...")
2020-08-02 15:46:36 +02:00
break
try:
camera, frame_time, current_tracked_objects = self.tracked_objects_queue.get(True, 10)
except queue.Empty:
continue
2020-09-07 19:17:42 +02:00
camera_state = self.camera_states[camera]
camera_state.update(frame_time, current_tracked_objects)
# update zone status for each label
for zone in self.config.cameras[camera].zones.keys():
2020-09-07 19:17:42 +02:00
# get labels for current camera and all labels in current zone
2020-11-10 04:11:27 +01:00
labels_for_camera = set([obj.obj_data['label'] for obj in camera_state.tracked_objects.values() if zone in obj.current_zones and not obj.false_positive])
2020-09-07 19:17:42 +02:00
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
2020-07-25 14:44:07 +02:00
# 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)
2020-11-25 17:37:41 +01:00
# cleanup event finished queue
while not self.event_processed_queue.empty():
event_id, camera = self.event_processed_queue.get()
self.camera_states[camera].finished(event_id)