This commit is contained in:
Blake Blackshear 2023-06-11 07:13:55 -05:00
parent 5631f45311
commit 2fbfff0394
2 changed files with 60 additions and 63 deletions

View File

@ -8,7 +8,7 @@ from norfair.drawing.drawer import Drawer
from frigate.video import ( from frigate.video import (
get_cluster_boundary, get_cluster_boundary,
get_cluster_candidates, get_cluster_candidates,
get_cluster_regions, get_cluster_region,
) )
@ -95,9 +95,10 @@ class TestConfig(unittest.TestCase):
self.frame_shape, self.min_region_size, boxes self.frame_shape, self.min_region_size, boxes
) )
regions = get_cluster_regions( regions = [
self.frame_shape, self.min_region_size, cluster_candidates, boxes get_cluster_region(self.frame_shape, self.min_region_size, candidate, boxes)
) for candidate in cluster_candidates
]
# save_clusters_image("regions", boxes, cluster_candidates, regions) # save_clusters_image("regions", boxes, cluster_candidates, regions)
assert len(regions) == 2 assert len(regions) == 2
@ -109,11 +110,12 @@ class TestConfig(unittest.TestCase):
self.frame_shape, self.min_region_size, boxes self.frame_shape, self.min_region_size, boxes
) )
regions = get_cluster_regions( regions = [
self.frame_shape, self.min_region_size, cluster_candidates, boxes get_cluster_region(self.frame_shape, self.min_region_size, candidate, boxes)
) for candidate in cluster_candidates
]
# save_clusters_image("too_small", boxes, cluster_candidates, regions) save_clusters_image("too_small", boxes, cluster_candidates, regions)
assert len(cluster_candidates) == 2 assert len(cluster_candidates) == 2
assert len(regions) == 2 assert len(regions) == 2
@ -125,9 +127,10 @@ class TestConfig(unittest.TestCase):
self.frame_shape, self.min_region_size, boxes self.frame_shape, self.min_region_size, boxes
) )
regions = get_cluster_regions( regions = [
self.frame_shape, self.min_region_size, cluster_candidates, boxes get_cluster_region(self.frame_shape, self.min_region_size, candidate, boxes)
) for candidate in cluster_candidates
]
# save_clusters_image("redundant", boxes, cluster_candidates, regions) # save_clusters_image("redundant", boxes, cluster_candidates, regions)
@ -148,9 +151,10 @@ class TestConfig(unittest.TestCase):
self.frame_shape, self.min_region_size, boxes self.frame_shape, self.min_region_size, boxes
) )
regions = get_cluster_regions( regions = [
self.frame_shape, self.min_region_size, cluster_candidates, boxes get_cluster_region(self.frame_shape, self.min_region_size, candidate, boxes)
) for candidate in cluster_candidates
]
# save_clusters_image("combine", boxes, cluster_candidates, regions) # save_clusters_image("combine", boxes, cluster_candidates, regions)
assert len(regions) == 1 assert len(regions) == 1
@ -165,9 +169,10 @@ class TestConfig(unittest.TestCase):
self.frame_shape, self.min_region_size, boxes self.frame_shape, self.min_region_size, boxes
) )
regions = get_cluster_regions( regions = [
self.frame_shape, self.min_region_size, cluster_candidates, boxes get_cluster_region(self.frame_shape, self.min_region_size, candidate, boxes)
) for candidate in cluster_candidates
]
# save_clusters_image("dont_combine", boxes, cluster_candidates, regions) # save_clusters_image("dont_combine", boxes, cluster_candidates, regions)
assert len(regions) == 2 assert len(regions) == 2

View File

@ -623,16 +623,27 @@ def get_cluster_candidates(frame_shape, min_region, boxes):
if compare_index in used_boxes: if compare_index in used_boxes:
continue continue
# get the region if you were to add this box to the cluster # if the box is not inside the potential cluster area, cluster them
cluster_regions = get_cluster_regions( if not box_inside(cluster_boundary, compare_box):
frame_shape, min_region, [cluster + [compare_index]], boxes
)
# if adding the box would result in multiple regions, dont cluster
if len(cluster_regions) > 1:
continue continue
# if the box is inside the potential cluster area, cluster them # get the region if you were to add this box to the cluster
if box_inside(cluster_boundary, compare_box): potential_cluster = cluster + [compare_index]
cluster_region = get_cluster_region(
frame_shape, min_region, potential_cluster, boxes
)
# if region could be smaller and either box would be too small
# for the resulting region, dont cluster
if (cluster_region[2] - cluster_region[0]) > min_region:
should_cluster = True
for b in potential_cluster:
box = boxes[b]
# boxes should be more than 5% of the area of the region
if area(box) / area(cluster_region) < 0.05:
should_cluster = False
break
if should_cluster:
cluster.append(compare_index) cluster.append(compare_index)
used_boxes.append(compare_index) used_boxes.append(compare_index)
cluster_candidates.append(cluster) cluster_candidates.append(cluster)
@ -642,41 +653,19 @@ def get_cluster_candidates(frame_shape, min_region, boxes):
return [list(tup) for tup in unique] return [list(tup) for tup in unique]
def get_cluster_regions(frame_shape, min_region, clusters, boxes): def get_cluster_region(frame_shape, min_region, cluster, boxes):
regions = [] min_x = frame_shape[1]
for c in clusters: min_y = frame_shape[0]
min_x = frame_shape[1] max_x = 0
min_y = frame_shape[0] max_y = 0
max_x = 0 for b in cluster:
max_y = 0 min_x = min(boxes[b][0], min_x)
for b in c: min_y = min(boxes[b][1], min_y)
min_x = min(boxes[b][0], min_x) max_x = max(boxes[b][2], max_x)
min_y = min(boxes[b][1], min_y) max_y = max(boxes[b][3], max_y)
max_x = max(boxes[b][2], max_x) return calculate_region(
max_y = max(boxes[b][3], max_y) frame_shape, min_x, min_y, max_x, max_y, min_region, multiplier=1.2
region = calculate_region( )
frame_shape, min_x, min_y, max_x, max_y, min_region, multiplier=1.2
)
regions.append(region)
# TODO: move this to a dedicated check function so it doesnt run again
# now check each box in the region to ensure it's not too small
# if it is, create a dedicated region for it
if len(c) > 1 and (region[2] - region[0]) > min_region:
for b in c:
box = boxes[b]
if area(box) / area(region) < 0.05:
regions.append(
calculate_region(
frame_shape,
box[0],
box[1],
box[2],
box[3],
min_region,
multiplier=1.35,
)
)
return regions
def process_frames( def process_frames(
@ -771,9 +760,12 @@ def process_frames(
frame_shape, region_min_size, combined_boxes frame_shape, region_min_size, combined_boxes
) )
regions = get_cluster_regions( regions = [
frame_shape, region_min_size, cluster_candidates, combined_boxes get_cluster_region(
) frame_shape, region_min_size, candidate, combined_boxes
)
for candidate in cluster_candidates
]
# if starting up, get the next startup scan region # if starting up, get the next startup scan region
if startup_scan_counter < 9: if startup_scan_counter < 9: