From 38d398c9676d78e3e88abff31c0717023bebe71b Mon Sep 17 00:00:00 2001 From: Nicolas Mowen Date: Sat, 28 Sep 2024 07:49:04 -0600 Subject: [PATCH] Dynamic attributes config (#14035) * Add config for attribute map and generate all labels from the map * Update docs * Formatting * Use the dynamic label map * Fix check * Fix docs typo --- docs/docs/configuration/reference.md | 10 +++++++++ frigate/config.py | 3 +-- frigate/const.py | 9 +++----- frigate/detectors/detector_config.py | 33 +++++++++++++++++++++++++++- frigate/object_processing.py | 7 ++++-- frigate/review/maintainer.py | 7 +++--- frigate/video.py | 6 ++--- 7 files changed, 57 insertions(+), 18 deletions(-) diff --git a/docs/docs/configuration/reference.md b/docs/docs/configuration/reference.md index 16b4a331b..90bf22fad 100644 --- a/docs/docs/configuration/reference.md +++ b/docs/docs/configuration/reference.md @@ -138,6 +138,16 @@ model: # Optional: Label name modifications. These are merged into the standard labelmap. labelmap: 2: vehicle + # Optional: Map of object labels to their attribute labels (default: depends on model) + attributes_map: + person: + - amazon + - face + car: + - amazon + - fedex + - license_plate + - ups # Optional: Audio Events Configuration # NOTE: Can be overridden at the camera level diff --git a/frigate/config.py b/frigate/config.py index 81c5fa6e4..7705f4ee1 100644 --- a/frigate/config.py +++ b/frigate/config.py @@ -25,7 +25,6 @@ from ruamel.yaml import YAML from typing_extensions import Self from frigate.const import ( - ALL_ATTRIBUTE_LABELS, AUDIO_MIN_CONFIDENCE, CACHE_DIR, CACHE_SEGMENT_FORMAT, @@ -1566,7 +1565,7 @@ class FrigateConfig(FrigateBaseModel): self.notifications.enabled_in_config = self.notifications.enabled # set default min_score for object attributes - for attribute in ALL_ATTRIBUTE_LABELS: + for attribute in self.model.all_attributes: if not self.objects.filters.get(attribute): self.objects.filters[attribute] = FilterConfig(min_score=0.7) elif self.objects.filters[attribute].min_score == 0.5: diff --git a/frigate/const.py b/frigate/const.py index 890fbb3ca..b37ca662e 100644 --- a/frigate/const.py +++ b/frigate/const.py @@ -15,13 +15,10 @@ PLUS_API_HOST = "https://api.frigate.video" # Attribute & Object constants -ATTRIBUTE_LABEL_MAP = { - "person": ["face", "amazon"], - "car": ["ups", "fedex", "amazon", "license_plate"], +DEFAULT_ATTRIBUTE_LABEL_MAP = { + "person": ["amazon", "face"], + "car": ["amazon", "fedex", "license_plate", "ups"], } -ALL_ATTRIBUTE_LABELS = [ - item for sublist in ATTRIBUTE_LABEL_MAP.values() for item in sublist -] LABEL_CONSOLIDATION_MAP = { "car": 0.8, "face": 0.5, diff --git a/frigate/detectors/detector_config.py b/frigate/detectors/detector_config.py index d8c841f42..bc0a0ff11 100644 --- a/frigate/detectors/detector_config.py +++ b/frigate/detectors/detector_config.py @@ -9,6 +9,7 @@ import requests from pydantic import BaseModel, ConfigDict, Field from pydantic.fields import PrivateAttr +from frigate.const import DEFAULT_ATTRIBUTE_LABEL_MAP from frigate.plus import PlusApi from frigate.util.builtin import generate_color_palette, load_labels @@ -42,6 +43,10 @@ class ModelConfig(BaseModel): labelmap: Dict[int, str] = Field( default_factory=dict, title="Labelmap customization." ) + attributes_map: Dict[str, list[str]] = Field( + default=DEFAULT_ATTRIBUTE_LABEL_MAP, + title="Map of object labels to their attribute labels.", + ) input_tensor: InputTensorEnum = Field( default=InputTensorEnum.nhwc, title="Model Input Tensor Shape" ) @@ -53,6 +58,7 @@ class ModelConfig(BaseModel): ) _merged_labelmap: Optional[Dict[int, str]] = PrivateAttr() _colormap: Dict[int, Tuple[int, int, int]] = PrivateAttr() + _all_attributes: list[str] = PrivateAttr() _model_hash: str = PrivateAttr() @property @@ -63,6 +69,10 @@ class ModelConfig(BaseModel): def colormap(self) -> Dict[int, Tuple[int, int, int]]: return self._colormap + @property + def all_attributes(self) -> list[str]: + return self._all_attributes + @property def model_hash(self) -> str: return self._model_hash @@ -76,6 +86,14 @@ class ModelConfig(BaseModel): } self._colormap = {} + # generate list of attribute labels + unique_attributes = set() + + for attributes in self.attributes_map.values(): + unique_attributes.update(attributes) + + self._all_attributes = list(unique_attributes) + def check_and_load_plus_model( self, plus_api: PlusApi, detector: str = None ) -> None: @@ -100,7 +118,7 @@ class ModelConfig(BaseModel): json.dump(model_info, f) else: with open(model_info_path, "r") as f: - model_info = json.load(f) + model_info: dict[str, any] = json.load(f) if detector and detector not in model_info["supportedDetectors"]: raise ValueError(f"Model does not support detector type of {detector}") @@ -110,6 +128,19 @@ class ModelConfig(BaseModel): self.input_tensor = model_info["inputShape"] self.input_pixel_format = model_info["pixelFormat"] self.model_type = model_info["type"] + + # generate list of attribute labels + self.attributes_map = { + **model_info.get("attributes", DEFAULT_ATTRIBUTE_LABEL_MAP), + **self.attributes_map, + } + unique_attributes = set() + + for attributes in self.attributes_map.values(): + unique_attributes.update(attributes) + + self._all_attributes = list(unique_attributes) + self._merged_labelmap = { **{int(key): val for key, val in model_info["labelMap"].items()}, **self.labelmap, diff --git a/frigate/object_processing.py b/frigate/object_processing.py index e73186f46..19cb5ae2b 100644 --- a/frigate/object_processing.py +++ b/frigate/object_processing.py @@ -25,7 +25,7 @@ from frigate.config import ( SnapshotsConfig, ZoomingModeEnum, ) -from frigate.const import ALL_ATTRIBUTE_LABELS, CLIPS_DIR, UPDATE_CAMERA_ACTIVITY +from frigate.const import CLIPS_DIR, UPDATE_CAMERA_ACTIVITY from frigate.events.types import EventStateEnum, EventTypeEnum from frigate.ptz.autotrack import PtzAutoTrackerThread from frigate.util.image import ( @@ -752,7 +752,10 @@ class CameraState: sub_label = None if obj.obj_data.get("sub_label"): - if obj.obj_data.get("sub_label")[0] in ALL_ATTRIBUTE_LABELS: + if ( + obj.obj_data.get("sub_label")[0] + in self.config.model.all_attributes + ): label = obj.obj_data["sub_label"][0] else: label = f"{object_type}-verified" diff --git a/frigate/review/maintainer.py b/frigate/review/maintainer.py index bea62256d..51ec349f8 100644 --- a/frigate/review/maintainer.py +++ b/frigate/review/maintainer.py @@ -20,7 +20,6 @@ from frigate.comms.detections_updater import DetectionSubscriber, DetectionTypeE from frigate.comms.inter_process import InterProcessRequestor from frigate.config import CameraConfig, FrigateConfig from frigate.const import ( - ALL_ATTRIBUTE_LABELS, CLEAR_ONGOING_REVIEW_SEGMENTS, CLIPS_DIR, UPSERT_REVIEW_SEGMENT, @@ -153,6 +152,8 @@ class ReviewSegmentMaintainer(threading.Thread): self.active_review_segments: dict[str, Optional[PendingReviewSegment]] = {} self.frame_manager = SharedMemoryFrameManager() + logger.error(f"All attributes are {config.model.all_attributes}") + # create communication for review segments self.requestor = InterProcessRequestor() self.config_subscriber = ConfigSubscriber("config/record/") @@ -253,7 +254,7 @@ class ReviewSegmentMaintainer(threading.Thread): for object in active_objects: if not object["sub_label"]: segment.detections[object["id"]] = object["label"] - elif object["sub_label"][0] in ALL_ATTRIBUTE_LABELS: + elif object["sub_label"][0] in self.config.model.all_attributes: segment.detections[object["id"]] = object["sub_label"][0] else: segment.detections[object["id"]] = f'{object["label"]}-verified' @@ -347,7 +348,7 @@ class ReviewSegmentMaintainer(threading.Thread): for object in active_objects: if not object["sub_label"]: detections[object["id"]] = object["label"] - elif object["sub_label"][0] in ALL_ATTRIBUTE_LABELS: + elif object["sub_label"][0] in self.config.model.all_attributes: detections[object["id"]] = object["sub_label"][0] else: detections[object["id"]] = f'{object["label"]}-verified' diff --git a/frigate/video.py b/frigate/video.py index 33e71a454..485764d2e 100755 --- a/frigate/video.py +++ b/frigate/video.py @@ -16,8 +16,6 @@ from frigate.comms.config_updater import ConfigSubscriber from frigate.comms.inter_process import InterProcessRequestor from frigate.config import CameraConfig, DetectConfig, ModelConfig from frigate.const import ( - ALL_ATTRIBUTE_LABELS, - ATTRIBUTE_LABEL_MAP, CACHE_DIR, CACHE_SEGMENT_FORMAT, REQUEST_REGION_GRID, @@ -727,7 +725,7 @@ def process_frames( tracked_detections = [ d for d in consolidated_detections - if d[0] not in ALL_ATTRIBUTE_LABELS + if d[0] not in model_config.all_attributes ] # now that we have refined our detections, we need to track objects object_tracker.match_and_update(frame_time, tracked_detections) @@ -737,7 +735,7 @@ def process_frames( # group the attribute detections based on what label they apply to attribute_detections = {} - for label, attribute_labels in ATTRIBUTE_LABEL_MAP.items(): + for label, attribute_labels in model_config.attributes_map.items(): attribute_detections[label] = [ d for d in consolidated_detections if d[0] in attribute_labels ]