blakeblackshear.frigate/frigate/timeline.py

102 lines
3.1 KiB
Python
Raw Normal View History

"""Record events for object, audio, etc. detections."""
import logging
import threading
import queue
from enum import Enum
from frigate.config import FrigateConfig
from frigate.models import Timeline
from multiprocessing.queues import Queue
from multiprocessing.synchronize import Event as MpEvent
from frigate.util import to_relative_box
logger = logging.getLogger(__name__)
class TimelineSourceEnum(str, Enum):
# api = "api"
# audio = "audio"
tracked_object = "tracked_object"
class TimelineProcessor(threading.Thread):
"""Handle timeline queue and update DB."""
def __init__(
self,
config: FrigateConfig,
queue: Queue,
stop_event: MpEvent,
) -> None:
threading.Thread.__init__(self)
self.name = "timeline_processor"
self.config = config
self.queue = queue
self.stop_event = stop_event
def run(self) -> None:
while not self.stop_event.is_set():
try:
(
camera,
input_type,
event_type,
prev_event_data,
event_data,
) = self.queue.get(timeout=1)
except queue.Empty:
continue
if input_type == TimelineSourceEnum.tracked_object:
self.handle_object_detection(
camera, event_type, prev_event_data, event_data
)
def handle_object_detection(
self,
camera: str,
event_type: str,
prev_event_data: dict[any, any],
event_data: dict[any, any],
) -> None:
"""Handle object detection."""
camera_config = self.config.cameras[camera]
timeline_entry = {
Timeline.timestamp: event_data["frame_time"],
Timeline.camera: camera,
Timeline.source: "tracked_object",
Timeline.source_id: event_data["id"],
Timeline.data: {
"box": to_relative_box(
camera_config.detect.width,
camera_config.detect.height,
event_data["box"],
),
"label": event_data["label"],
"region": to_relative_box(
camera_config.detect.width,
camera_config.detect.height,
event_data["region"],
),
},
}
if event_type == "start":
timeline_entry[Timeline.class_type] = "visible"
Timeline.insert(timeline_entry).execute()
elif (
event_type == "update"
and prev_event_data["current_zones"] != event_data["current_zones"]
and len(event_data["current_zones"]) > 0
):
timeline_entry[Timeline.class_type] = "entered_zone"
timeline_entry[Timeline.data]["zones"] = event_data["current_zones"]
Timeline.insert(timeline_entry).execute()
elif event_type == "end":
timeline_entry[Timeline.class_type] = "gone"
Timeline.insert(timeline_entry).execute()