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:
Nicolas Mowen
2025-08-17 11:27:42 -06:00
committed by GitHub
parent 856aab8e6e
commit 5a49d1f73c
10 changed files with 194 additions and 160 deletions

View File

@@ -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)