mirror of
https://github.com/blakeblackshear/frigate.git
synced 2026-04-19 23:08:08 +02:00
Enable mypy for track and fix typing errors (#19529)
* Enable mypy for track * WIP cleaning up tracked object * Fix tracked object typing * Fix typing and imports of centroid tracker * Cleanup typing * Cleanup * Formatting * Fix imports * Don't specify callable type * Type out json setting
This commit is contained in:
@@ -6,6 +6,7 @@ import queue
|
||||
import threading
|
||||
from collections import defaultdict
|
||||
from enum import Enum
|
||||
from multiprocessing import Queue as MpQueue
|
||||
from multiprocessing.synchronize import Event as MpEvent
|
||||
from typing import Any
|
||||
|
||||
@@ -39,6 +40,7 @@ from frigate.const import (
|
||||
)
|
||||
from frigate.events.types import EventStateEnum, EventTypeEnum
|
||||
from frigate.models import Event, ReviewSegment, Timeline
|
||||
from frigate.ptz.autotrack import PtzAutoTrackerThread
|
||||
from frigate.track.tracked_object import TrackedObject
|
||||
from frigate.util.image import SharedMemoryFrameManager
|
||||
|
||||
@@ -56,10 +58,10 @@ class TrackedObjectProcessor(threading.Thread):
|
||||
self,
|
||||
config: FrigateConfig,
|
||||
dispatcher: Dispatcher,
|
||||
tracked_objects_queue,
|
||||
ptz_autotracker_thread,
|
||||
stop_event,
|
||||
):
|
||||
tracked_objects_queue: MpQueue,
|
||||
ptz_autotracker_thread: PtzAutoTrackerThread,
|
||||
stop_event: MpEvent,
|
||||
) -> None:
|
||||
super().__init__(name="detected_frames_processor")
|
||||
self.config = config
|
||||
self.dispatcher = dispatcher
|
||||
@@ -98,8 +100,12 @@ class TrackedObjectProcessor(threading.Thread):
|
||||
# }
|
||||
# }
|
||||
# }
|
||||
self.zone_data = defaultdict(lambda: defaultdict(dict))
|
||||
self.active_zone_data = defaultdict(lambda: defaultdict(dict))
|
||||
self.zone_data: dict[str, dict[str, Any]] = defaultdict(
|
||||
lambda: defaultdict(dict)
|
||||
)
|
||||
self.active_zone_data: dict[str, dict[str, Any]] = defaultdict(
|
||||
lambda: defaultdict(dict)
|
||||
)
|
||||
|
||||
for camera in self.config.cameras.keys():
|
||||
self.create_camera_state(camera)
|
||||
@@ -107,7 +113,7 @@ class TrackedObjectProcessor(threading.Thread):
|
||||
def create_camera_state(self, camera: str) -> None:
|
||||
"""Creates a new camera state."""
|
||||
|
||||
def start(camera: str, obj: TrackedObject, frame_name: str):
|
||||
def start(camera: str, obj: TrackedObject, frame_name: str) -> None:
|
||||
self.event_sender.publish(
|
||||
(
|
||||
EventTypeEnum.tracked_object,
|
||||
@@ -118,7 +124,7 @@ class TrackedObjectProcessor(threading.Thread):
|
||||
)
|
||||
)
|
||||
|
||||
def update(camera: str, obj: TrackedObject, frame_name: str):
|
||||
def update(camera: str, obj: TrackedObject, frame_name: str) -> None:
|
||||
obj.has_snapshot = self.should_save_snapshot(camera, obj)
|
||||
obj.has_clip = self.should_retain_recording(camera, obj)
|
||||
after = obj.to_dict()
|
||||
@@ -139,10 +145,10 @@ class TrackedObjectProcessor(threading.Thread):
|
||||
)
|
||||
)
|
||||
|
||||
def autotrack(camera: str, obj: TrackedObject, frame_name: str):
|
||||
def autotrack(camera: str, obj: TrackedObject, frame_name: str) -> None:
|
||||
self.ptz_autotracker_thread.ptz_autotracker.autotrack_object(camera, obj)
|
||||
|
||||
def end(camera: str, obj: TrackedObject, frame_name: str):
|
||||
def end(camera: str, obj: TrackedObject, frame_name: str) -> None:
|
||||
# populate has_snapshot
|
||||
obj.has_snapshot = self.should_save_snapshot(camera, obj)
|
||||
obj.has_clip = self.should_retain_recording(camera, obj)
|
||||
@@ -211,7 +217,7 @@ class TrackedObjectProcessor(threading.Thread):
|
||||
|
||||
return False
|
||||
|
||||
def camera_activity(camera, activity):
|
||||
def camera_activity(camera: str, activity: dict[str, Any]) -> None:
|
||||
last_activity = self.camera_activity.get(camera)
|
||||
|
||||
if not last_activity or activity != last_activity:
|
||||
@@ -229,7 +235,7 @@ class TrackedObjectProcessor(threading.Thread):
|
||||
camera_state.on("camera_activity", camera_activity)
|
||||
self.camera_states[camera] = camera_state
|
||||
|
||||
def should_save_snapshot(self, camera, obj: TrackedObject):
|
||||
def should_save_snapshot(self, camera: str, obj: TrackedObject) -> bool:
|
||||
if obj.false_positive:
|
||||
return False
|
||||
|
||||
@@ -252,7 +258,7 @@ class TrackedObjectProcessor(threading.Thread):
|
||||
|
||||
return True
|
||||
|
||||
def should_retain_recording(self, camera: str, obj: TrackedObject):
|
||||
def should_retain_recording(self, camera: str, obj: TrackedObject) -> bool:
|
||||
if obj.false_positive:
|
||||
return False
|
||||
|
||||
@@ -272,7 +278,7 @@ class TrackedObjectProcessor(threading.Thread):
|
||||
|
||||
return True
|
||||
|
||||
def should_mqtt_snapshot(self, camera, obj: TrackedObject):
|
||||
def should_mqtt_snapshot(self, camera: str, obj: TrackedObject) -> bool:
|
||||
# object never changed position
|
||||
if obj.is_stationary():
|
||||
return False
|
||||
@@ -287,7 +293,9 @@ class TrackedObjectProcessor(threading.Thread):
|
||||
|
||||
return True
|
||||
|
||||
def update_mqtt_motion(self, camera, frame_time, motion_boxes):
|
||||
def update_mqtt_motion(
|
||||
self, camera: str, frame_time: float, motion_boxes: list
|
||||
) -> None:
|
||||
# publish if motion is currently being detected
|
||||
if motion_boxes:
|
||||
# only send ON if motion isn't already active
|
||||
@@ -313,11 +321,15 @@ class TrackedObjectProcessor(threading.Thread):
|
||||
# reset the last_motion so redundant `off` commands aren't sent
|
||||
self.last_motion_detected[camera] = 0
|
||||
|
||||
def get_best(self, camera, label):
|
||||
def get_best(self, camera: str, label: str) -> dict[str, Any]:
|
||||
# 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]
|
||||
|
||||
if not best_obj.thumbnail_data:
|
||||
return {}
|
||||
|
||||
best = best_obj.thumbnail_data.copy()
|
||||
best["frame"] = camera_state.frame_cache.get(
|
||||
best_obj.thumbnail_data["frame_time"]
|
||||
@@ -340,7 +352,7 @@ class TrackedObjectProcessor(threading.Thread):
|
||||
|
||||
return self.camera_states[camera].get_current_frame(draw_options)
|
||||
|
||||
def get_current_frame_time(self, camera) -> int:
|
||||
def get_current_frame_time(self, camera: str) -> float:
|
||||
"""Returns the latest frame time for a given camera."""
|
||||
return self.camera_states[camera].current_frame_time
|
||||
|
||||
@@ -348,7 +360,7 @@ class TrackedObjectProcessor(threading.Thread):
|
||||
self, event_id: str, sub_label: str | None, score: float | None
|
||||
) -> None:
|
||||
"""Update sub label for given event id."""
|
||||
tracked_obj: TrackedObject = None
|
||||
tracked_obj: TrackedObject | None = None
|
||||
|
||||
for state in self.camera_states.values():
|
||||
tracked_obj = state.tracked_objects.get(event_id)
|
||||
@@ -357,7 +369,7 @@ class TrackedObjectProcessor(threading.Thread):
|
||||
break
|
||||
|
||||
try:
|
||||
event: Event = Event.get(Event.id == event_id)
|
||||
event: Event | None = Event.get(Event.id == event_id)
|
||||
except DoesNotExist:
|
||||
event = None
|
||||
|
||||
@@ -368,12 +380,12 @@ class TrackedObjectProcessor(threading.Thread):
|
||||
tracked_obj.obj_data["sub_label"] = (sub_label, score)
|
||||
|
||||
if event:
|
||||
event.sub_label = sub_label
|
||||
event.sub_label = sub_label # type: ignore[assignment]
|
||||
data = event.data
|
||||
if sub_label is None:
|
||||
data["sub_label_score"] = None
|
||||
data["sub_label_score"] = None # type: ignore[index]
|
||||
elif score is not None:
|
||||
data["sub_label_score"] = score
|
||||
data["sub_label_score"] = score # type: ignore[index]
|
||||
event.data = data
|
||||
event.save()
|
||||
|
||||
@@ -402,7 +414,7 @@ class TrackedObjectProcessor(threading.Thread):
|
||||
objects_list = []
|
||||
sub_labels = set()
|
||||
events = Event.select(Event.id, Event.label, Event.sub_label).where(
|
||||
Event.id.in_(detection_ids)
|
||||
Event.id.in_(detection_ids) # type: ignore[call-arg, misc]
|
||||
)
|
||||
for det_event in events:
|
||||
if det_event.sub_label:
|
||||
@@ -431,13 +443,11 @@ class TrackedObjectProcessor(threading.Thread):
|
||||
f"Updated sub_label for event {event_id} in review segment {review_segment.id}"
|
||||
)
|
||||
|
||||
except ReviewSegment.DoesNotExist:
|
||||
except DoesNotExist:
|
||||
logger.debug(
|
||||
f"No review segment found with event ID {event_id} when updating sub_label"
|
||||
)
|
||||
|
||||
return True
|
||||
|
||||
def set_object_attribute(
|
||||
self,
|
||||
event_id: str,
|
||||
@@ -446,7 +456,7 @@ class TrackedObjectProcessor(threading.Thread):
|
||||
score: float | None,
|
||||
) -> None:
|
||||
"""Update attribute for given event id."""
|
||||
tracked_obj: TrackedObject = None
|
||||
tracked_obj: TrackedObject | None = None
|
||||
|
||||
for state in self.camera_states.values():
|
||||
tracked_obj = state.tracked_objects.get(event_id)
|
||||
@@ -455,7 +465,7 @@ class TrackedObjectProcessor(threading.Thread):
|
||||
break
|
||||
|
||||
try:
|
||||
event: Event = Event.get(Event.id == event_id)
|
||||
event: Event | None = Event.get(Event.id == event_id)
|
||||
except DoesNotExist:
|
||||
event = None
|
||||
|
||||
@@ -470,16 +480,14 @@ class TrackedObjectProcessor(threading.Thread):
|
||||
|
||||
if event:
|
||||
data = event.data
|
||||
data[field_name] = field_value
|
||||
data[field_name] = field_value # type: ignore[index]
|
||||
if field_value is None:
|
||||
data[f"{field_name}_score"] = None
|
||||
data[f"{field_name}_score"] = None # type: ignore[index]
|
||||
elif score is not None:
|
||||
data[f"{field_name}_score"] = score
|
||||
data[f"{field_name}_score"] = score # type: ignore[index]
|
||||
event.data = data
|
||||
event.save()
|
||||
|
||||
return True
|
||||
|
||||
def save_lpr_snapshot(self, payload: tuple) -> None:
|
||||
# save the snapshot image
|
||||
(frame, event_id, camera) = payload
|
||||
@@ -638,7 +646,7 @@ class TrackedObjectProcessor(threading.Thread):
|
||||
)
|
||||
self.ongoing_manual_events.pop(event_id)
|
||||
|
||||
def force_end_all_events(self, camera: str, camera_state: CameraState):
|
||||
def force_end_all_events(self, camera: str, camera_state: CameraState) -> None:
|
||||
"""Ends all active events on camera when disabling."""
|
||||
last_frame_name = camera_state.previous_frame_id
|
||||
for obj_id, obj in list(camera_state.tracked_objects.items()):
|
||||
@@ -656,7 +664,7 @@ class TrackedObjectProcessor(threading.Thread):
|
||||
{"enabled": False, "motion": 0, "objects": []},
|
||||
)
|
||||
|
||||
def run(self):
|
||||
def run(self) -> None:
|
||||
while not self.stop_event.is_set():
|
||||
# check for config updates
|
||||
updated_topics = self.camera_config_subscriber.check_for_updates()
|
||||
@@ -698,11 +706,14 @@ class TrackedObjectProcessor(threading.Thread):
|
||||
|
||||
# check for sub label updates
|
||||
while True:
|
||||
(raw_topic, payload) = self.sub_label_subscriber.check_for_update(
|
||||
timeout=0
|
||||
)
|
||||
update = self.sub_label_subscriber.check_for_update(timeout=0)
|
||||
|
||||
if not raw_topic:
|
||||
if not update:
|
||||
break
|
||||
|
||||
(raw_topic, payload) = update
|
||||
|
||||
if not raw_topic or not payload:
|
||||
break
|
||||
|
||||
topic = str(raw_topic)
|
||||
|
||||
Reference in New Issue
Block a user