blakeblackshear.frigate/frigate/object_processing.py

868 lines
32 KiB
Python
Raw Normal View History

import base64
2021-08-16 15:02:04 +02:00
import copy
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
import numpy as np
from frigate.config import CameraConfig, SnapshotsConfig, RecordConfig, FrigateConfig
2021-08-16 15:02:04 +02:00
from frigate.const import CACHE_DIR, CLIPS_DIR, RECORD_DIR
2020-11-04 13:31:25 +01:00
from frigate.edgetpu import load_labels
2021-06-15 22:19:49 +02:00
from frigate.util import (
SharedMemoryFrameManager,
2021-08-16 15:02:04 +02:00
calculate_region,
2021-06-15 22:19:49 +02:00
draw_box_with_label,
draw_timestamp,
)
2020-02-16 04:07:54 +01:00
2020-11-04 04:26:39 +01:00
logger = logging.getLogger(__name__)
2021-02-17 14:23:32 +01:00
def on_edge(box, frame_shape):
if (
2021-02-17 14:23:32 +01:00
box[0] == 0
or box[1] == 0
or box[2] == frame_shape[1] - 1
or box[3] == frame_shape[0] - 1
):
return True
2021-02-17 14:23:32 +01:00
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
2021-02-17 14:23:32 +01:00
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%
2021-02-17 14:23:32 +01:00
if new_obj["score"] > current_thumb["score"] + 0.05:
return True
# if the area is 10% larger
2021-02-17 14:23:32 +01:00
if new_obj["area"] > current_thumb["area"] * 1.1:
return True
return False
2021-02-17 14:23:32 +01:00
class TrackedObject:
2021-08-16 15:02:04 +02:00
def __init__(
self, camera, colormap, camera_config: CameraConfig, frame_cache, obj_data
):
self.obj_data = obj_data
self.camera = camera
2021-08-16 15:02:04 +02:00
self.colormap = colormap
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.has_clip = False
self.has_snapshot = False
self.top_score = self.computed_score = 0.0
self.thumbnail_data = None
self.last_updated = 0
self.last_published = 0
self.frame = None
2020-12-20 14:19:18 +01:00
self.previous = self.to_dict()
# start the score history
2021-02-17 14:23:32 +01:00
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
2021-02-17 14:23:32 +01:00
threshold = self.camera_config.objects.filters[self.obj_data["label"]].threshold
return self.computed_score < threshold
def compute_score(self):
scores = self.score_history[:]
# pad with zeros if you dont have at least 3 scores
if len(scores) < 3:
2021-02-17 14:23:32 +01:00
scores += [0.0] * (3 - len(scores))
return median(scores)
def update(self, current_frame_time, obj_data):
2020-12-20 14:19:18 +01:00
significant_update = False
2021-07-07 14:02:36 +02:00
zone_change = False
self.obj_data.update(obj_data)
# if the object is not in the current frame, add a 0.0 to the score history
2021-02-17 14:23:32 +01:00
if self.obj_data["frame_time"] != current_frame_time:
self.score_history.append(0.0)
else:
2021-02-17 14:23:32 +01:00
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
2021-02-17 14:23:32 +01:00
if self.thumbnail_data is None or is_better_thumbnail(
self.thumbnail_data, self.obj_data, self.camera_config.frame_shape
2020-11-25 19:06:01 +01:00
):
self.thumbnail_data = {
2021-02-17 14:23:32 +01:00
"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-12-20 14:19:18 +01:00
significant_update = True
# check zones
current_zones = []
2021-02-17 14:23:32 +01:00
bottom_center = (self.obj_data["centroid"][0], self.obj_data["box"][3])
# check each zone
for name, zone in self.camera_config.zones.items():
# if the zone is not for this object type, skip
if len(zone.objects) > 0 and not self.obj_data["label"] in zone.objects:
continue
contour = zone.contour
# check if the object is in the zone
2021-02-17 14:23:32 +01:00
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)
2020-12-20 14:19:18 +01:00
# if the zones changed, signal an update
if not self.false_positive and set(self.current_zones) != set(current_zones):
2021-07-07 14:02:36 +02:00
zone_change = True
2020-12-20 14:19:18 +01:00
self.current_zones = current_zones
2021-07-07 14:02:36 +02:00
return (significant_update, zone_change)
def to_dict(self, include_thumbnail: bool = False):
2021-06-19 15:40:28 +02:00
snapshot_time = (
self.thumbnail_data["frame_time"]
if not self.thumbnail_data is None
else 0.0
)
2021-02-23 02:51:31 +01:00
event = {
"id": self.obj_data["id"],
"camera": self.camera,
"frame_time": self.obj_data["frame_time"],
2021-06-19 15:40:28 +02:00
"snapshot_time": snapshot_time,
"label": self.obj_data["label"],
"top_score": self.top_score,
"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(),
"has_clip": self.has_clip,
"has_snapshot": self.has_snapshot,
}
2021-02-23 02:51:31 +01:00
if include_thumbnail:
event["thumbnail"] = base64.b64encode(self.get_thumbnail()).decode("utf-8")
2021-02-23 02:51:31 +01:00
return event
2021-01-15 14:52:53 +01:00
def get_thumbnail(self):
2021-02-17 14:23:32 +01:00
if (
self.thumbnail_data is None
or self.thumbnail_data["frame_time"] not in self.frame_cache
2021-02-17 14:23:32 +01:00
):
ret, jpg = cv2.imencode(".jpg", np.zeros((175, 175, 3), np.uint8))
2021-02-17 14:23:32 +01:00
jpg_bytes = self.get_jpg_bytes(
timestamp=False, bounding_box=False, crop=True, height=175
)
if jpg_bytes:
return jpg_bytes
else:
2021-02-17 14:23:32 +01:00
ret, jpg = cv2.imencode(".jpg", np.zeros((175, 175, 3), np.uint8))
return jpg.tobytes()
2021-02-17 14:23:32 +01:00
2021-06-19 15:38:42 +02:00
def get_clean_png(self):
if self.thumbnail_data is None:
return None
try:
best_frame = cv2.cvtColor(
self.frame_cache[self.thumbnail_data["frame_time"]],
cv2.COLOR_YUV2BGR_I420,
)
except KeyError:
logger.warning(
f"Unable to create clean png because frame {self.thumbnail_data['frame_time']} is not in the cache"
)
return None
ret, png = cv2.imencode(".png", best_frame)
if ret:
return png.tobytes()
else:
return None
2021-02-17 14:23:32 +01:00
def get_jpg_bytes(
2021-07-02 14:47:03 +02:00
self, timestamp=False, bounding_box=False, crop=False, height=None, quality=70
2021-02-17 14:23:32 +01:00
):
2021-01-15 14:52:53 +01:00
if self.thumbnail_data is None:
return None
2021-02-17 14:23:32 +01:00
try:
2021-02-17 14:23:32 +01:00
best_frame = cv2.cvtColor(
self.frame_cache[self.thumbnail_data["frame_time"]],
cv2.COLOR_YUV2BGR_I420,
)
except KeyError:
2021-02-17 14:23:32 +01:00
logger.warning(
f"Unable to create jpg because frame {self.thumbnail_data['frame_time']} is not in the cache"
)
return None
2021-02-17 14:23:32 +01:00
if bounding_box:
thickness = 2
2021-08-16 15:02:04 +02:00
color = self.colormap[self.obj_data["label"]]
# draw the bounding boxes on the frame
2021-02-17 14:23:32 +01:00
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 crop:
2021-02-17 14:23:32 +01:00
box = self.thumbnail_data["box"]
region = calculate_region(
best_frame.shape, box[0], box[1], box[2], box[3], 1.1
)
best_frame = best_frame[region[1] : region[3], region[0] : region[2]]
if height:
2021-02-17 14:23:32 +01:00
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 timestamp:
2021-06-24 07:45:27 +02:00
color = self.camera_config.timestamp_style.color
draw_timestamp(
2021-06-15 22:19:49 +02:00
best_frame,
self.thumbnail_data["frame_time"],
2021-06-15 22:19:49 +02:00
self.camera_config.timestamp_style.format,
font_effect=self.camera_config.timestamp_style.effect,
font_thickness=self.camera_config.timestamp_style.thickness,
2021-09-04 23:39:56 +02:00
font_color=(color.blue, color.green, color.red),
2021-06-15 22:19:49 +02:00
position=self.camera_config.timestamp_style.position,
2021-02-17 14:23:32 +01:00
)
2021-07-02 14:47:03 +02:00
ret, jpg = cv2.imencode(
".jpg", best_frame, [int(cv2.IMWRITE_JPEG_QUALITY), quality]
)
if ret:
return jpg.tobytes()
else:
return None
2021-02-17 14:23:32 +01:00
def zone_filtered(obj: TrackedObject, object_config):
2021-02-17 14:23:32 +01:00
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
2021-02-17 14:23:32 +01:00
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
2021-02-17 14:23:32 +01:00
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
2021-02-17 14:23:32 +01:00
2020-09-07 19:17:42 +02:00
# Maintains the state of a camera
2021-02-17 14:23:32 +01:00
class CameraState:
2021-07-13 15:51:15 +02:00
def __init__(
self, name, config: FrigateConfig, frame_manager: SharedMemoryFrameManager
):
2020-09-07 19:17:42 +02:00
self.name = name
self.config = config
2021-06-24 07:51:41 +02:00
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] = {}
self.object_counts = defaultdict(int)
self.tracked_objects: Dict[str, TrackedObject] = {}
2020-11-11 23:55:50 +01:00
self.frame_cache = {}
self.zone_objects = defaultdict(list)
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
2021-01-12 14:00:08 +01:00
self.motion_boxes = []
self.regions = []
2020-09-07 19:17:42 +02:00
self.previous_frame_id = None
self.callbacks = defaultdict(list)
2020-09-07 19:17:42 +02:00
def get_current_frame(self, draw_options={}):
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
2021-02-17 14:23:32 +01:00
tracked_objects = {k: v.to_dict() for k, v in self.tracked_objects.items()}
motion_boxes = self.motion_boxes.copy()
regions = self.regions.copy()
2020-10-11 19:16:05 +02:00
frame_copy = cv2.cvtColor(frame_copy, cv2.COLOR_YUV2BGR_I420)
# draw on the frame
2021-02-17 14:23:32 +01:00
if draw_options.get("bounding_boxes"):
2020-10-11 19:16:05 +02:00
# draw the bounding boxes on the frame
for obj in tracked_objects.values():
if obj["frame_time"] == frame_time:
thickness = 2
2021-08-16 15:02:04 +02:00
color = self.config.model.colormap[obj["label"]]
else:
2020-10-11 19:16:05 +02:00
thickness = 1
2021-02-17 14:23:32 +01:00
color = (255, 0, 0)
2020-10-11 19:16:05 +02:00
# draw the bounding boxes on the frame
2021-02-17 14:23:32 +01:00
box = obj["box"]
draw_box_with_label(
frame_copy,
box[0],
box[1],
box[2],
box[3],
obj["label"],
f"{obj['score']:.0%} {int(obj['area'])}",
2021-02-17 14:23:32 +01:00
thickness=thickness,
color=color,
)
2021-02-17 14:23:32 +01:00
if draw_options.get("regions"):
for region in regions:
2021-02-17 14:23:32 +01:00
cv2.rectangle(
frame_copy,
(region[0], region[1]),
(region[2], region[3]),
(0, 255, 0),
2,
)
2021-02-17 14:23:32 +01:00
if draw_options.get("zones"):
for name, zone in self.camera_config.zones.items():
2021-02-17 14:23:32 +01:00
thickness = (
8
if any(
name in obj["current_zones"] for obj in tracked_objects.values()
2021-02-17 14:23:32 +01:00
)
else 2
)
cv2.drawContours(frame_copy, [zone.contour], -1, zone.color, thickness)
2021-02-17 14:23:32 +01:00
if draw_options.get("mask"):
2021-07-13 15:51:15 +02:00
mask_overlay = np.where(self.camera_config.motion.mask == [0])
2021-02-17 14:23:32 +01:00
frame_copy[mask_overlay] = [0, 0, 0]
2021-02-17 14:23:32 +01:00
if draw_options.get("motion_boxes"):
for m_box in motion_boxes:
2021-02-17 14:23:32 +01:00
cv2.rectangle(
frame_copy,
(m_box[0], m_box[1]),
(m_box[2], m_box[3]),
(0, 0, 255),
2,
)
2021-02-17 14:23:32 +01:00
if draw_options.get("timestamp"):
2021-06-24 23:02:46 +02:00
color = self.camera_config.timestamp_style.color
draw_timestamp(
2021-02-17 14:23:32 +01:00
frame_copy,
frame_time,
2021-06-15 22:19:49 +02:00
self.camera_config.timestamp_style.format,
font_effect=self.camera_config.timestamp_style.effect,
font_thickness=self.camera_config.timestamp_style.thickness,
2021-09-04 23:39:56 +02:00
font_color=(color.blue, color.green, color.red),
2021-06-15 22:19:49 +02:00
position=self.camera_config.timestamp_style.position,
2021-02-17 14:23:32 +01:00
)
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)
def update(self, frame_time, current_detections, motion_boxes, regions):
# get the new frame
2020-09-07 19:17:42 +02:00
frame_id = f"{self.name}{frame_time}"
2021-02-17 14:23:32 +01:00
current_frame = self.frame_manager.get(
frame_id, self.camera_config.frame_shape_yuv
)
2020-09-07 19:17:42 +02:00
tracked_objects = self.tracked_objects.copy()
current_ids = set(current_detections.keys())
previous_ids = set(tracked_objects.keys())
removed_ids = previous_ids.difference(current_ids)
new_ids = current_ids.difference(previous_ids)
updated_ids = current_ids.intersection(previous_ids)
2020-09-07 19:17:42 +02:00
for id in new_ids:
new_obj = tracked_objects[id] = TrackedObject(
2021-08-16 15:02:04 +02:00
self.name,
self.config.model.colormap,
self.camera_config,
self.frame_cache,
current_detections[id],
2021-02-17 14:23:32 +01:00
)
2020-09-07 19:17:42 +02:00
# call event handlers
2021-02-17 14:23:32 +01:00
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 = tracked_objects[id]
2021-07-07 14:02:36 +02:00
significant_update, zone_change = updated_obj.update(
frame_time, current_detections[id]
)
2020-12-20 14:19:18 +01:00
if significant_update:
# ensure this frame is stored in the cache
2021-02-17 14:23:32 +01:00
if (
updated_obj.thumbnail_data["frame_time"] == frame_time
and frame_time not in self.frame_cache
):
2020-12-20 14:19:18 +01:00
self.frame_cache[frame_time] = np.copy(current_frame)
2021-02-17 14:23:32 +01:00
updated_obj.last_updated = frame_time
2021-02-17 14:23:32 +01:00
# if it has been more than 5 seconds since the last publish
2021-07-07 14:02:36 +02:00
# and the last update is greater than the last publish or
# the object has changed zones
2021-02-17 14:23:32 +01:00
if (
frame_time - updated_obj.last_published > 5
and updated_obj.last_updated > updated_obj.last_published
2021-07-07 14:02:36 +02:00
) or zone_change:
2020-12-20 14:19:18 +01:00
# call event handlers
2021-02-17 14:23:32 +01:00
for c in self.callbacks["update"]:
2020-12-20 14:19:18 +01:00
c(self.name, updated_obj, frame_time)
updated_obj.last_published = frame_time
2020-09-07 19:17:42 +02:00
for id in removed_ids:
# publish events to mqtt
removed_obj = tracked_objects[id]
2021-02-17 14:23:32 +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?
2020-09-07 19:17:42 +02:00
# maintain best objects
for obj in tracked_objects.values():
2021-02-17 14:23:32 +01:00
object_type = obj.obj_data["label"]
# if the object's thumbnail is not from the current frame
if obj.false_positive or obj.thumbnail_data["frame_time"] != 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
2021-02-17 14:23:32 +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
):
2020-11-10 04:31:45 +01:00
self.best_objects[object_type] = obj
2021-02-17 14:23:32 +01: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
2021-02-17 14:23:32 +01: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(
obj.obj_data["label"]
for obj in tracked_objects.values()
if not obj.false_positive
)
2020-09-07 19:17:42 +02:00
# report on detected objects
for obj_name, count in obj_counter.items():
if count != self.object_counts[obj_name]:
self.object_counts[obj_name] = count
2021-02-17 14:23:32 +01:00
for c in self.callbacks["object_status"]:
c(self.name, obj_name, count)
2020-09-07 19:17:42 +02:00
# expire any objects that are >0 and no longer detected
2021-02-17 14:23:32 +01:00
expired_objects = [
obj_name
for obj_name, count in self.object_counts.items()
if count > 0 and obj_name not in obj_counter
2021-02-17 14:23:32 +01:00
]
2020-09-07 19:17:42 +02:00
for obj_name in expired_objects:
self.object_counts[obj_name] = 0
2021-02-17 14:23:32 +01:00
for c in self.callbacks["object_status"]:
c(self.name, obj_name, 0)
2021-02-17 14:23:32 +01:00
for c in self.callbacks["snapshot"]:
2020-11-25 19:06:01 +01:00
c(self.name, self.best_objects[obj_name], frame_time)
# cleanup thumbnail frame cache
current_thumb_frames = {
obj.thumbnail_data["frame_time"]
for obj in tracked_objects.values()
if not obj.false_positive
}
current_best_frames = {
obj.thumbnail_data["frame_time"] for obj in self.best_objects.values()
}
2021-02-17 14:23:32 +01:00
thumb_frames_to_delete = [
t
for t in self.frame_cache.keys()
if t not in current_thumb_frames and t not in current_best_frames
2021-02-17 14:23:32 +01:00
]
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.tracked_objects = tracked_objects
self.current_frame_time = frame_time
self.motion_boxes = motion_boxes
self.regions = regions
2020-10-11 19:16:05 +02:00
self._current_frame = current_frame
if self.previous_frame_id is not None:
self.frame_manager.close(self.previous_frame_id)
2020-10-11 19:16:05 +02:00
self.previous_frame_id = frame_id
2020-09-07 19:17:42 +02:00
2021-02-17 14:23:32 +01:00
2020-02-16 04:07:54 +01:00
class TrackedObjectProcessor(threading.Thread):
2021-02-17 14:23:32 +01:00
def __init__(
self,
config: FrigateConfig,
client,
topic_prefix,
tracked_objects_queue,
event_queue,
event_processed_queue,
video_output_queue,
2021-02-17 14:23:32 +01:00
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
self.video_output_queue = video_output_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):
2021-02-17 14:23:32 +01:00
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):
obj.has_snapshot = self.should_save_snapshot(camera, obj)
obj.has_clip = self.should_retain_recording(camera, obj)
2020-12-20 14:19:18 +01:00
after = obj.to_dict()
2021-02-17 14:23:32 +01:00
message = {
"before": obj.previous,
"after": after,
"type": "new" if obj.previous["false_positive"] else "update",
}
self.client.publish(
f"{self.topic_prefix}/events", json.dumps(message), retain=False
)
2020-12-20 14:19:18 +01:00
obj.previous = after
self.event_queue.put(
("update", camera, obj.to_dict(include_thumbnail=True))
)
2020-09-07 19:17:42 +02:00
2020-11-25 19:06:01 +01:00
def end(camera, obj: TrackedObject, current_frame_time):
# populate has_snapshot
obj.has_snapshot = self.should_save_snapshot(camera, obj)
obj.has_clip = self.should_retain_recording(camera, obj)
# write the snapshot to disk
if obj.has_snapshot:
snapshot_config: SnapshotsConfig = self.config.cameras[camera].snapshots
jpg_bytes = obj.get_jpg_bytes(
timestamp=snapshot_config.timestamp,
bounding_box=snapshot_config.bounding_box,
crop=snapshot_config.crop,
height=snapshot_config.height,
quality=snapshot_config.quality,
)
if jpg_bytes is None:
logger.warning(f"Unable to save snapshot for {obj.obj_data['id']}.")
else:
with open(
os.path.join(CLIPS_DIR, f"{camera}-{obj.obj_data['id']}.jpg"),
"wb",
) as j:
j.write(jpg_bytes)
# write clean snapshot if enabled
if snapshot_config.clean_copy:
png_bytes = obj.get_clean_png()
if png_bytes is None:
logger.warning(
f"Unable to save clean snapshot for {obj.obj_data['id']}."
)
else:
with open(
os.path.join(
CLIPS_DIR,
f"{camera}-{obj.obj_data['id']}-clean.png",
),
"wb",
) as p:
p.write(png_bytes)
2020-11-26 03:22:54 +01:00
if not obj.false_positive:
2021-02-17 14:23:32 +01:00
message = {
"before": obj.previous,
"after": obj.to_dict(),
"type": "end",
}
self.client.publish(
f"{self.topic_prefix}/events", json.dumps(message), retain=False
)
self.event_queue.put(("end", camera, obj.to_dict(include_thumbnail=True)))
2021-02-17 14:23:32 +01:00
2020-11-25 19:06:01 +01:00
def snapshot(camera, obj: TrackedObject, current_frame_time):
mqtt_config = self.config.cameras[camera].mqtt
if mqtt_config.enabled and self.should_mqtt_snapshot(camera, obj):
jpg_bytes = obj.get_jpg_bytes(
timestamp=mqtt_config.timestamp,
bounding_box=mqtt_config.bounding_box,
crop=mqtt_config.crop,
2021-02-17 14:23:32 +01:00
height=mqtt_config.height,
2021-07-02 14:47:03 +02:00
quality=mqtt_config.quality,
)
2021-02-09 14:35:41 +01:00
if jpg_bytes is None:
2021-02-17 14:23:32 +01:00
logger.warning(
f"Unable to send mqtt snapshot for {obj.obj_data['id']}."
)
2021-02-09 14:35:41 +01:00
else:
2021-02-17 14:23:32 +01:00
self.client.publish(
f"{self.topic_prefix}/{camera}/{obj.obj_data['label']}/snapshot",
jpg_bytes,
retain=True,
)
2020-09-07 19:17:42 +02:00
def object_status(camera, object_name, status):
2021-02-17 14:23:32 +01:00
self.client.publish(
f"{self.topic_prefix}/{camera}/{object_name}", status, retain=False
)
2020-09-07 19:17:42 +02:00
for camera in self.config.cameras.keys():
camera_state = CameraState(camera, self.config, self.frame_manager)
2021-02-17 14:23:32 +01: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)
2020-09-07 19:17:42 +02:00
self.camera_states[camera] = camera_state
# {
# 'zone_name': {
# 'person': {
# 'camera_1': 2,
# 'camera_2': 1
# }
2020-09-07 19:17:42 +02:00
# }
# }
self.zone_data = defaultdict(lambda: defaultdict(dict))
def should_save_snapshot(self, camera, obj: TrackedObject):
if obj.false_positive:
return False
snapshot_config: SnapshotsConfig = self.config.cameras[camera].snapshots
if not snapshot_config.enabled:
return False
# if there are required zones and there is no overlap
required_zones = snapshot_config.required_zones
if len(required_zones) > 0 and not obj.entered_zones & set(required_zones):
2021-02-17 14:23:32 +01:00
logger.debug(
f"Not creating snapshot for {obj.obj_data['id']} because it did not enter required zones"
)
return False
return True
def should_retain_recording(self, camera, obj: TrackedObject):
if obj.false_positive:
return False
record_config: RecordConfig = self.config.cameras[camera].record
# Recording is disabled
if not record_config.enabled:
return False
# If there are required zones and there is no overlap
required_zones = record_config.events.required_zones
if len(required_zones) > 0 and not set(obj.entered_zones) & set(required_zones):
logger.debug(
f"Not creating clip for {obj.obj_data['id']} because it did not enter required zones"
)
return False
# If the required objects are not present
if (
record_config.events.objects is not None
and obj.obj_data["label"] not in record_config.events.objects
):
logger.debug(
f"Not creating clip for {obj.obj_data['id']} because it did not contain required objects"
)
return False
return True
def should_mqtt_snapshot(self, camera, obj: TrackedObject):
# if there are required zones and there is no overlap
required_zones = self.config.cameras[camera].mqtt.required_zones
if len(required_zones) > 0 and not obj.entered_zones & set(required_zones):
2021-02-17 14:23:32 +01:00
logger.debug(
f"Not sending mqtt for {obj.obj_data['id']} because it did not enter required zones"
)
return False
return True
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.thumbnail_data.copy()
2021-02-17 14:23:32 +01:00
best["frame"] = camera_state.frame_cache.get(
best_obj.thumbnail_data["frame_time"]
)
2020-11-12 00:44:51 +01:00
return best
2020-02-16 04:07:54 +01:00
else:
return {}
def get_current_frame(self, camera, draw_options={}):
return self.camera_states[camera].get_current_frame(draw_options)
2020-02-16 04:07:54 +01:00
def run(self):
while not self.stop_event.is_set():
2020-08-02 15:46:36 +02:00
try:
2021-02-17 14:23:32 +01:00
(
camera,
frame_time,
current_tracked_objects,
motion_boxes,
regions,
) = self.tracked_objects_queue.get(True, 10)
2020-08-02 15:46:36 +02:00
except queue.Empty:
continue
2020-09-07 19:17:42 +02:00
camera_state = self.camera_states[camera]
2021-02-17 14:23:32 +01:00
camera_state.update(
frame_time, current_tracked_objects, motion_boxes, regions
)
2020-09-07 19:17:42 +02:00
self.video_output_queue.put(
(
camera,
frame_time,
current_tracked_objects,
motion_boxes,
regions,
)
2021-05-23 04:02:26 +02:00
)
# update zone counts for each label
# for each zone in the current camera
for zone in self.config.cameras[camera].zones.keys():
# count labels for the camera in the zone
obj_counter = Counter(
obj.obj_data["label"]
for obj in camera_state.tracked_objects.values()
if zone in obj.current_zones and not obj.false_positive
)
# update counts and publish status
for label in set(self.zone_data[zone].keys()) | set(obj_counter.keys()):
# if we have previously published a count for this zone/label
zone_label = self.zone_data[zone][label]
if camera in zone_label:
current_count = sum(zone_label.values())
2021-02-17 14:23:32 +01:00
zone_label[camera] = (
obj_counter[label] if label in obj_counter else 0
)
new_count = sum(zone_label.values())
if new_count != current_count:
2021-02-17 14:23:32 +01:00
self.client.publish(
f"{self.topic_prefix}/{zone}/{label}",
new_count,
retain=False,
)
# if this is a new zone/label combo for this camera
else:
if label in obj_counter:
zone_label[camera] = obj_counter[label]
2021-02-17 14:23:32 +01:00
self.client.publish(
f"{self.topic_prefix}/{zone}/{label}",
obj_counter[label],
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()
2020-11-25 17:37:41 +01:00
self.camera_states[camera].finished(event_id)
logger.info(f"Exiting object processor...")