Clean up tracked object to pass model data

This commit is contained in:
Nicolas Mowen 2024-10-16 13:21:44 -06:00
parent 8281989aa0
commit acc25f78ed

View File

@ -20,6 +20,7 @@ from frigate.comms.inter_process import InterProcessRequestor
from frigate.config import ( from frigate.config import (
CameraConfig, CameraConfig,
FrigateConfig, FrigateConfig,
ModelConfig,
MqttConfig, MqttConfig,
RecordConfig, RecordConfig,
SnapshotsConfig, SnapshotsConfig,
@ -109,8 +110,7 @@ def is_better_thumbnail(label, current_thumb, new_obj, frame_shape) -> bool:
class TrackedObject: class TrackedObject:
def __init__( def __init__(
self, self,
camera, model_config: ModelConfig,
colormap,
camera_config: CameraConfig, camera_config: CameraConfig,
frame_cache, frame_cache,
obj_data: dict[str, any], obj_data: dict[str, any],
@ -120,8 +120,8 @@ class TrackedObject:
del obj_data["score_history"] del obj_data["score_history"]
self.obj_data = obj_data self.obj_data = obj_data
self.camera = camera self.colormap = model_config.colormap
self.colormap = colormap self.logos = model_config.all_attribute_logos
self.camera_config = camera_config self.camera_config = camera_config
self.frame_cache = frame_cache self.frame_cache = frame_cache
self.zone_presence: dict[str, int] = {} self.zone_presence: dict[str, int] = {}
@ -245,9 +245,7 @@ class TrackedObject:
# populate the sub_label for object with highest scoring logo # populate the sub_label for object with highest scoring logo
if self.obj_data["label"] in ["car", "package", "person"]: if self.obj_data["label"] in ["car", "package", "person"]:
recognized_logos = { recognized_logos = {
k: self.attributes[k] k: self.attributes[k] for k in self.logos if k in self.attributes
for k in ["ups", "fedex", "amazon"]
if k in self.attributes
} }
if len(recognized_logos) > 0: if len(recognized_logos) > 0:
max_logo = max(recognized_logos, key=recognized_logos.get) max_logo = max(recognized_logos, key=recognized_logos.get)
@ -291,7 +289,7 @@ class TrackedObject:
def to_dict(self, include_thumbnail: bool = False): def to_dict(self, include_thumbnail: bool = False):
event = { event = {
"id": self.obj_data["id"], "id": self.obj_data["id"],
"camera": self.camera, "camera": self.camera_config.name,
"frame_time": self.obj_data["frame_time"], "frame_time": self.obj_data["frame_time"],
"snapshot": self.thumbnail_data, "snapshot": self.thumbnail_data,
"label": self.obj_data["label"], "label": self.obj_data["label"],
@ -707,8 +705,7 @@ class CameraState:
for id in new_ids: for id in new_ids:
new_obj = tracked_objects[id] = TrackedObject( new_obj = tracked_objects[id] = TrackedObject(
self.name, self.config.model,
self.config.model.colormap,
self.camera_config, self.camera_config,
self.frame_cache, self.frame_cache,
current_detections[id], current_detections[id],