forked from Github/frigate
Label attributes (#6829)
* pass attribute labels as attributes * add label attrs to events and snapshots * incorporate area of license_plate and face into snapshot selection * populate sublabels for cars with logos
This commit is contained in:
Blake Blackshear
@@ -24,6 +24,7 @@ from frigate.const import CLIPS_DIR
|
||||
from frigate.events.maintainer import EventTypeEnum
|
||||
from frigate.util import (
|
||||
SharedMemoryFrameManager,
|
||||
area,
|
||||
calculate_region,
|
||||
draw_box_with_label,
|
||||
draw_timestamp,
|
||||
@@ -42,11 +43,45 @@ def on_edge(box, frame_shape):
|
||||
return True
|
||||
|
||||
|
||||
def is_better_thumbnail(current_thumb, new_obj, frame_shape) -> bool:
|
||||
def has_better_attr(current_thumb, new_obj, attr_label) -> bool:
|
||||
max_new_attr = max(
|
||||
[0]
|
||||
+ [area(a["box"]) for a in new_obj["attributes"] if a["label"] == attr_label]
|
||||
)
|
||||
max_current_attr = max(
|
||||
[0]
|
||||
+ [
|
||||
area(a["box"])
|
||||
for a in current_thumb["attributes"]
|
||||
if a["label"] == attr_label
|
||||
]
|
||||
)
|
||||
|
||||
# if the thumb has a higher scoring attr
|
||||
return max_new_attr > max_current_attr
|
||||
|
||||
|
||||
def is_better_thumbnail(label, current_thumb, new_obj, frame_shape) -> bool:
|
||||
# larger is better
|
||||
# cutoff images are less ideal, but they should also be smaller?
|
||||
# better scores are obviously better too
|
||||
|
||||
# check face on person
|
||||
if label == "person":
|
||||
if has_better_attr(current_thumb, new_obj, "face"):
|
||||
return True
|
||||
# if the current thumb has a face attr, dont update unless it gets better
|
||||
if any([a["label"] == "face" for a in current_thumb["attributes"]]):
|
||||
return False
|
||||
|
||||
# check license_plate on car
|
||||
if label == "car":
|
||||
if has_better_attr(current_thumb, new_obj, "license_plate"):
|
||||
return True
|
||||
# if the current thumb has a license_plate attr, dont update unless it gets better
|
||||
if any([a["label"] == "license_plate" for a in current_thumb["attributes"]]):
|
||||
return False
|
||||
|
||||
# if the new_thumb is on an edge, and the current thumb is not
|
||||
if on_edge(new_obj["box"], frame_shape) and not on_edge(
|
||||
current_thumb["box"], frame_shape
|
||||
@@ -76,6 +111,7 @@ class TrackedObject:
|
||||
self.zone_presence = {}
|
||||
self.current_zones = []
|
||||
self.entered_zones = []
|
||||
self.attributes = set()
|
||||
self.false_positive = True
|
||||
self.has_clip = False
|
||||
self.has_snapshot = False
|
||||
@@ -125,7 +161,10 @@ class TrackedObject:
|
||||
if not self.false_positive:
|
||||
# determine if this frame is a better thumbnail
|
||||
if self.thumbnail_data is None or is_better_thumbnail(
|
||||
self.thumbnail_data, obj_data, self.camera_config.frame_shape
|
||||
self.obj_data["label"],
|
||||
self.thumbnail_data,
|
||||
obj_data,
|
||||
self.camera_config.frame_shape,
|
||||
):
|
||||
self.thumbnail_data = {
|
||||
"frame_time": obj_data["frame_time"],
|
||||
@@ -133,6 +172,7 @@ class TrackedObject:
|
||||
"area": obj_data["area"],
|
||||
"region": obj_data["region"],
|
||||
"score": obj_data["score"],
|
||||
"attributes": obj_data["attributes"],
|
||||
}
|
||||
thumb_update = True
|
||||
|
||||
@@ -164,6 +204,19 @@ class TrackedObject:
|
||||
if 0 < zone_score < zone.inertia:
|
||||
self.zone_presence[name] = zone_score - 1
|
||||
|
||||
# maintain attributes
|
||||
for attr in obj_data["attributes"]:
|
||||
self.attributes.add(attr["label"])
|
||||
|
||||
# populate the sub_label for car with first logo if it exists
|
||||
if self.obj_data["label"] == "car" and "sub_label" not in self.obj_data:
|
||||
recognized_logos = self.attributes.intersection(
|
||||
set(["ups", "fedex", "amazon"])
|
||||
)
|
||||
if len(recognized_logos) > 0:
|
||||
self.obj_data["sub_label"] = recognized_logos.pop()
|
||||
|
||||
# check for significant change
|
||||
if not self.false_positive:
|
||||
# if the zones changed, signal an update
|
||||
if set(self.current_zones) != set(current_zones):
|
||||
@@ -214,6 +267,8 @@ class TrackedObject:
|
||||
"entered_zones": self.entered_zones.copy(),
|
||||
"has_clip": self.has_clip,
|
||||
"has_snapshot": self.has_snapshot,
|
||||
"attributes": list(self.attributes),
|
||||
"current_attributes": self.obj_data["attributes"],
|
||||
}
|
||||
|
||||
if include_thumbnail:
|
||||
@@ -294,6 +349,21 @@ class TrackedObject:
|
||||
color=color,
|
||||
)
|
||||
|
||||
# draw any attributes
|
||||
for attribute in self.thumbnail_data["attributes"]:
|
||||
box = attribute["box"]
|
||||
draw_box_with_label(
|
||||
best_frame,
|
||||
box[0],
|
||||
box[1],
|
||||
box[2],
|
||||
box[3],
|
||||
attribute["label"],
|
||||
f"{attribute['score']:.0%}",
|
||||
thickness=thickness,
|
||||
color=color,
|
||||
)
|
||||
|
||||
if crop:
|
||||
box = self.thumbnail_data["box"]
|
||||
box_size = 300
|
||||
@@ -421,6 +491,21 @@ class CameraState:
|
||||
color=color,
|
||||
)
|
||||
|
||||
# draw any attributes
|
||||
for attribute in obj["current_attributes"]:
|
||||
box = attribute["box"]
|
||||
draw_box_with_label(
|
||||
frame_copy,
|
||||
box[0],
|
||||
box[1],
|
||||
box[2],
|
||||
box[3],
|
||||
attribute["label"],
|
||||
f"{attribute['score']:.0%}",
|
||||
thickness=thickness,
|
||||
color=color,
|
||||
)
|
||||
|
||||
if draw_options.get("regions"):
|
||||
for region in regions:
|
||||
cv2.rectangle(
|
||||
@@ -553,6 +638,7 @@ class CameraState:
|
||||
# or the current object is older than desired, use the new object
|
||||
if (
|
||||
is_better_thumbnail(
|
||||
object_type,
|
||||
current_best.thumbnail_data,
|
||||
obj.thumbnail_data,
|
||||
self.camera_config.frame_shape,
|
||||
|
||||
Reference in New Issue
Block a user