Update logos handling (#14396)

* Add attribute for logos

* Clean up tracked object to pass model data

* Update default attributes map
This commit is contained in:
Nicolas Mowen 2024-10-16 15:22:34 -06:00 committed by GitHub
parent e836523bc3
commit 5f77408956
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 33 additions and 11 deletions

View File

@ -17,7 +17,21 @@ PLUS_API_HOST = "https://api.frigate.video"
DEFAULT_ATTRIBUTE_LABEL_MAP = { DEFAULT_ATTRIBUTE_LABEL_MAP = {
"person": ["amazon", "face"], "person": ["amazon", "face"],
"car": ["amazon", "fedex", "license_plate", "ups"], "car": [
"amazon",
"an_post",
"dhl",
"dpd",
"fedex",
"gls",
"license_plate",
"nzpost",
"postnl",
"postnord",
"purolator",
"ups",
"usps",
],
} }
LABEL_CONSOLIDATION_MAP = { LABEL_CONSOLIDATION_MAP = {
"car": 0.8, "car": 0.8,

View File

@ -59,6 +59,7 @@ class ModelConfig(BaseModel):
_merged_labelmap: Optional[Dict[int, str]] = PrivateAttr() _merged_labelmap: Optional[Dict[int, str]] = PrivateAttr()
_colormap: Dict[int, Tuple[int, int, int]] = PrivateAttr() _colormap: Dict[int, Tuple[int, int, int]] = PrivateAttr()
_all_attributes: list[str] = PrivateAttr() _all_attributes: list[str] = PrivateAttr()
_all_attribute_logos: list[str] = PrivateAttr()
_model_hash: str = PrivateAttr() _model_hash: str = PrivateAttr()
@property @property
@ -73,6 +74,10 @@ class ModelConfig(BaseModel):
def all_attributes(self) -> list[str]: def all_attributes(self) -> list[str]:
return self._all_attributes return self._all_attributes
@property
def all_attribute_logos(self) -> list[str]:
return self._all_attribute_logos
@property @property
def model_hash(self) -> str: def model_hash(self) -> str:
return self._model_hash return self._model_hash
@ -93,6 +98,9 @@ class ModelConfig(BaseModel):
unique_attributes.update(attributes) unique_attributes.update(attributes)
self._all_attributes = list(unique_attributes) self._all_attributes = list(unique_attributes)
self._all_attribute_logos = list(
unique_attributes - set(["face", "license_plate"])
)
def check_and_load_plus_model( def check_and_load_plus_model(
self, plus_api: PlusApi, detector: str = None self, plus_api: PlusApi, detector: str = None
@ -140,6 +148,9 @@ class ModelConfig(BaseModel):
unique_attributes.update(attributes) unique_attributes.update(attributes)
self._all_attributes = list(unique_attributes) self._all_attributes = list(unique_attributes)
self._all_attribute_logos = list(
unique_attributes - set(["face", "license_plate"])
)
self._merged_labelmap = { self._merged_labelmap = {
**{int(key): val for key, val in model_info["labelMap"].items()}, **{int(key): val for key, val in model_info["labelMap"].items()},

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],