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 = {
"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 = {
"car": 0.8,

View File

@ -59,6 +59,7 @@ class ModelConfig(BaseModel):
_merged_labelmap: Optional[Dict[int, str]] = PrivateAttr()
_colormap: Dict[int, Tuple[int, int, int]] = PrivateAttr()
_all_attributes: list[str] = PrivateAttr()
_all_attribute_logos: list[str] = PrivateAttr()
_model_hash: str = PrivateAttr()
@property
@ -73,6 +74,10 @@ class ModelConfig(BaseModel):
def all_attributes(self) -> list[str]:
return self._all_attributes
@property
def all_attribute_logos(self) -> list[str]:
return self._all_attribute_logos
@property
def model_hash(self) -> str:
return self._model_hash
@ -93,6 +98,9 @@ class ModelConfig(BaseModel):
unique_attributes.update(attributes)
self._all_attributes = list(unique_attributes)
self._all_attribute_logos = list(
unique_attributes - set(["face", "license_plate"])
)
def check_and_load_plus_model(
self, plus_api: PlusApi, detector: str = None
@ -140,6 +148,9 @@ class ModelConfig(BaseModel):
unique_attributes.update(attributes)
self._all_attributes = list(unique_attributes)
self._all_attribute_logos = list(
unique_attributes - set(["face", "license_plate"])
)
self._merged_labelmap = {
**{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 (
CameraConfig,
FrigateConfig,
ModelConfig,
MqttConfig,
RecordConfig,
SnapshotsConfig,
@ -109,8 +110,7 @@ def is_better_thumbnail(label, current_thumb, new_obj, frame_shape) -> bool:
class TrackedObject:
def __init__(
self,
camera,
colormap,
model_config: ModelConfig,
camera_config: CameraConfig,
frame_cache,
obj_data: dict[str, any],
@ -120,8 +120,8 @@ class TrackedObject:
del obj_data["score_history"]
self.obj_data = obj_data
self.camera = camera
self.colormap = colormap
self.colormap = model_config.colormap
self.logos = model_config.all_attribute_logos
self.camera_config = camera_config
self.frame_cache = frame_cache
self.zone_presence: dict[str, int] = {}
@ -245,9 +245,7 @@ class TrackedObject:
# populate the sub_label for object with highest scoring logo
if self.obj_data["label"] in ["car", "package", "person"]:
recognized_logos = {
k: self.attributes[k]
for k in ["ups", "fedex", "amazon"]
if k in self.attributes
k: self.attributes[k] for k in self.logos if k in self.attributes
}
if len(recognized_logos) > 0:
max_logo = max(recognized_logos, key=recognized_logos.get)
@ -291,7 +289,7 @@ class TrackedObject:
def to_dict(self, include_thumbnail: bool = False):
event = {
"id": self.obj_data["id"],
"camera": self.camera,
"camera": self.camera_config.name,
"frame_time": self.obj_data["frame_time"],
"snapshot": self.thumbnail_data,
"label": self.obj_data["label"],
@ -707,8 +705,7 @@ class CameraState:
for id in new_ids:
new_obj = tracked_objects[id] = TrackedObject(
self.name,
self.config.model.colormap,
self.config.model,
self.camera_config,
self.frame_cache,
current_detections[id],