blakeblackshear.frigate/frigate/test/test_video.py
Nicolas Mowen ec4d79eafc
Fix bug in intersection logic (#6780)
* Fix bug in intersection logic

* Fix isort

* Remove unrelated test

* Formatting

* Fix type in test
2023-06-12 17:38:22 -05:00

193 lines
5.7 KiB
Python

import unittest
import cv2
import numpy as np
from norfair.drawing.color import Palette
from norfair.drawing.drawer import Drawer
from frigate.util import intersection
from frigate.video import (
get_cluster_boundary,
get_cluster_candidates,
get_cluster_region,
)
def draw_box(frame, box, color=(255, 0, 0), thickness=2):
cv2.rectangle(
frame,
(box[0], box[1]),
(box[2], box[3]),
color,
thickness,
)
def save_clusters_image(name, boxes, candidates, regions=[]):
canvas = np.zeros((1000, 2000, 3), np.uint8)
for cluster in candidates:
color = Palette.choose_color(np.random.rand())
for b in cluster:
box = boxes[b]
draw_box(canvas, box, color, 2)
# bottom right
text_anchor = (
box[2],
box[3],
)
canvas = Drawer.text(
canvas,
str(b),
position=text_anchor,
size=None,
color=(255, 255, 255),
thickness=None,
)
for r in regions:
draw_box(canvas, r, (0, 255, 0), 2)
cv2.imwrite(
f"debug/frames/{name}.jpg",
canvas,
)
def save_cluster_boundary_image(name, boxes, bounding_boxes):
canvas = np.zeros((1000, 2000, 3), np.uint8)
color = Palette.choose_color(np.random.rand())
for box in boxes:
draw_box(canvas, box, color, 2)
for bound in bounding_boxes:
draw_box(canvas, bound, (0, 255, 0), 2)
cv2.imwrite(
f"debug/frames/{name}.jpg",
canvas,
)
class TestRegion(unittest.TestCase):
def setUp(self):
self.frame_shape = (1000, 2000)
self.min_region_size = 160
def test_cluster_candidates(self):
boxes = [(100, 100, 200, 200), (202, 150, 252, 200), (900, 900, 950, 950)]
cluster_candidates = get_cluster_candidates(
self.frame_shape, self.min_region_size, boxes
)
# save_clusters_image("cluster_candidates", boxes, cluster_candidates)
assert len(cluster_candidates) == 2
def test_cluster_boundary(self):
boxes = [(100, 100, 200, 200), (215, 215, 325, 325)]
boundary_boxes = [
get_cluster_boundary(box, self.min_region_size) for box in boxes
]
# save_cluster_boundary_image("bound", boxes, boundary_boxes)
assert len(boundary_boxes) == 2
def test_cluster_regions(self):
boxes = [(100, 100, 200, 200), (202, 150, 252, 200), (900, 900, 950, 950)]
cluster_candidates = get_cluster_candidates(
self.frame_shape, self.min_region_size, boxes
)
regions = [
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)
assert len(regions) == 2
def test_box_too_small_for_cluster(self):
boxes = [(100, 100, 600, 600), (655, 100, 700, 145)]
cluster_candidates = get_cluster_candidates(
self.frame_shape, self.min_region_size, boxes
)
regions = [
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)
assert len(cluster_candidates) == 2
assert len(regions) == 2
def test_redundant_clusters(self):
boxes = [(100, 100, 200, 200), (305, 305, 415, 415)]
cluster_candidates = get_cluster_candidates(
self.frame_shape, self.min_region_size, boxes
)
regions = [
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)
assert len(cluster_candidates) == 2
assert all([len(c) == 1 for c in cluster_candidates])
assert len(regions) == 2
def test_combine_boxes(self):
boxes = [
(460, 0, 561, 144),
(565, 0, 586, 71),
]
# boundary_boxes = [get_cluster_boundary(box) for box in boxes]
# save_cluster_boundary_image("combine_bound", boxes, boundary_boxes)
cluster_candidates = get_cluster_candidates(
self.frame_shape, self.min_region_size, boxes
)
regions = [
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)
assert len(regions) == 1
def test_dont_combine_boxes(self):
boxes = [(460, 0, 532, 129), (586, 0, 606, 46)]
# boundary_boxes = [get_cluster_boundary(box) for box in boxes]
# save_cluster_boundary_image("dont_combine_bound", boxes, boundary_boxes)
cluster_candidates = get_cluster_candidates(
self.frame_shape, self.min_region_size, boxes
)
regions = [
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)
assert len(regions) == 2
class TestObjectBoundingBoxes(unittest.TestCase):
def setUp(self) -> None:
pass
def test_box_intersection(self):
box_a = [2012, 191, 2031, 205]
box_b = [887, 92, 985, 151]
box_c = [899, 128, 1080, 175]
assert intersection(box_a, box_b) == None
assert intersection(box_b, box_c) == (899, 128, 985, 151)