Add sub labels to data

This commit is contained in:
Nicolas Mowen 2024-02-20 16:19:37 -07:00
parent b6b8c470be
commit 63ff3c980f

View File

@ -42,6 +42,7 @@ class PendingReviewSegment:
severity: SeverityEnum, severity: SeverityEnum,
detections: set[str] = set(), detections: set[str] = set(),
objects: set[str] = set(), objects: set[str] = set(),
sub_labels: set[str] = set(),
zones: set[str] = set(), zones: set[str] = set(),
audio: set[str] = set(), audio: set[str] = set(),
motion: list[int] = [], motion: list[int] = [],
@ -53,6 +54,7 @@ class PendingReviewSegment:
self.severity = severity self.severity = severity
self.detections = detections self.detections = detections
self.objects = objects self.objects = objects
self.sub_labels = sub_labels
self.zones = zones self.zones = zones
self.audio = audio self.audio = audio
self.sig_motion_areas = motion self.sig_motion_areas = motion
@ -109,6 +111,7 @@ class PendingReviewSegment:
ReviewSegment.data: { ReviewSegment.data: {
"detections": list(self.detections), "detections": list(self.detections),
"objects": list(self.objects), "objects": list(self.objects),
"sub_labels": list(self.sub_labels),
"zones": list(self.zones), "zones": list(self.zones),
"audio": list(self.audio), "audio": list(self.audio),
"significant_motion_areas": self.sig_motion_areas, "significant_motion_areas": self.sig_motion_areas,
@ -168,6 +171,9 @@ class ReviewSegmentMaintainer(threading.Thread):
segment.detections.add(object["id"]) segment.detections.add(object["id"])
segment.objects.add(object["label"]) segment.objects.add(object["label"])
if object["sub_label"]:
segment.sub_labels.add(object["sub_label"])
# if object is alert label and has qualified for recording # if object is alert label and has qualified for recording
# mark this review as alert # mark this review as alert
if ( if (
@ -207,6 +213,7 @@ class ReviewSegmentMaintainer(threading.Thread):
has_sig_object = False has_sig_object = False
detections: set = set() detections: set = set()
objects: set = set() objects: set = set()
sub_labels: set = (set(),)
zones: set = set() zones: set = set()
for object in active_objects: for object in active_objects:
@ -219,6 +226,10 @@ class ReviewSegmentMaintainer(threading.Thread):
detections.add(object["id"]) detections.add(object["id"])
objects.add(object["label"]) objects.add(object["label"])
if object["sub_label"]:
sub_labels.add(object["sub_label"])
zones.update(object["current_zones"]) zones.update(object["current_zones"])
self.active_review_segments[camera] = PendingReviewSegment( self.active_review_segments[camera] = PendingReviewSegment(
@ -226,8 +237,9 @@ class ReviewSegmentMaintainer(threading.Thread):
frame_time, frame_time,
SeverityEnum.alert if has_sig_object else SeverityEnum.detection, SeverityEnum.alert if has_sig_object else SeverityEnum.detection,
detections, detections,
objects, objects=objects,
zones, sub_labels=sub_labels,
zones=zones,
) )
elif len(motion) >= 20: elif len(motion) >= 20:
self.active_review_segments[camera] = PendingReviewSegment( self.active_review_segments[camera] = PendingReviewSegment(