diff --git a/frigate/test/test_obects.py b/frigate/test/test_obects.py index 8fe831980e..ee0162ec01 100644 --- a/frigate/test/test_obects.py +++ b/frigate/test/test_obects.py @@ -1,6 +1,22 @@ +import random import unittest +import numpy as np + from frigate.track.tracked_object import TrackedObjectAttribute +from frigate.util.object import average_boxes + + +class TestBoxStatistics(unittest.TestCase): + def test_average_boxes_matches_numpy(self) -> None: + rng = random.Random(0) + for _ in range(5000): + boxes = [ + [rng.randint(0, 4000) for _ in range(4)] + for _ in range(rng.randint(1, 10)) + ] + expected = [float(np.mean([b[i] for b in boxes])) for i in range(4)] + self.assertEqual(average_boxes(boxes), expected) class TestAttribute(unittest.TestCase): diff --git a/frigate/util/object.py b/frigate/util/object.py index 311fccf477..b4ee72dc07 100644 --- a/frigate/util/object.py +++ b/frigate/util/object.py @@ -339,18 +339,13 @@ def reduce_boxes(boxes, iou_threshold=0.0): def average_boxes(boxes: list[list[int, int, int, int]]) -> list[int, int, int, int]: """Return a box that is the average of a list of boxes.""" - x_mins = [] - y_mins = [] - x_max = [] - y_max = [] - - for box in boxes: - x_mins.append(box[0]) - y_mins.append(box[1]) - x_max.append(box[2]) - y_max.append(box[3]) - - return [np.mean(x_mins), np.mean(y_mins), np.mean(x_max), np.mean(y_max)] + n = len(boxes) + return [ + sum(box[0] for box in boxes) / n, + sum(box[1] for box in boxes) / n, + sum(box[2] for box in boxes) / n, + sum(box[3] for box in boxes) / n, + ] def median_of_boxes(boxes: list[list[int, int, int, int]]) -> list[int, int, int, int]: