From 730ba38bba17bd8c0a48997008b452cd62434fd6 Mon Sep 17 00:00:00 2001
From: Nick Mowen <nickmowen213@gmail.com>
Date: Wed, 18 Oct 2023 07:12:17 -0600
Subject: [PATCH] Add tests

---
 frigate/test/test_video.py | 34 ++++++++++++++++++++++++++++++++++
 frigate/util/object.py     |  4 ++--
 2 files changed, 36 insertions(+), 2 deletions(-)

diff --git a/frigate/test/test_video.py b/frigate/test/test_video.py
index f2667bb62..9fd46a877 100644
--- a/frigate/test/test_video.py
+++ b/frigate/test/test_video.py
@@ -10,6 +10,7 @@ from frigate.util.object import (
     get_cluster_boundary,
     get_cluster_candidates,
     get_cluster_region,
+    get_region_from_grid,
 )
 
 
@@ -190,3 +191,36 @@ class TestObjectBoundingBoxes(unittest.TestCase):
 
         assert intersection(box_a, box_b) == None
         assert intersection(box_b, box_c) == (899, 128, 985, 151)
+
+
+class TestRegionGrid(unittest.TestCase):
+    def setUp(self) -> None:
+        pass
+
+    def test_region_in_range(self):
+        """Test that region is kept at minimal size when within std dev."""
+        frame_shape = (720, 1280)
+        box = [450, 450, 550, 550]
+        region_grid = [
+            [],
+            [],
+            [],
+            [{}, {}, {}, {}, {}, {"sizes": [0.25], "mean": 0.26, "std_dev": 0.01}],
+        ]
+
+        region = get_region_from_grid(frame_shape, box, 320, region_grid)
+        assert region[2] - region[0] == 320
+
+    def test_region_out_of_range(self):
+        """Test that region is upsized when outside of std dev."""
+        frame_shape = (720, 1280)
+        box = [450, 450, 550, 550]
+        region_grid = [
+            [],
+            [],
+            [],
+            [{}, {}, {}, {}, {}, {"sizes": [0.5], "mean": 0.5, "std_dev": 0.1}],
+        ]
+
+        region = get_region_from_grid(frame_shape, box, 320, region_grid)
+        assert region[2] - region[0] > 320
diff --git a/frigate/util/object.py b/frigate/util/object.py
index 02103c469..3fa98df59 100644
--- a/frigate/util/object.py
+++ b/frigate/util/object.py
@@ -166,8 +166,8 @@ def get_region_from_grid(
         box[0] + (min(frame_shape[1], box[2]) - box[0]) / 2,
         box[1] + (min(frame_shape[0], box[3]) - box[1]) / 2,
     )
-    grid_x = int(centroid[0] / frame_shape[1] * len(region_grid))
-    grid_y = int(centroid[1] / frame_shape[0] * len(region_grid))
+    grid_x = int(centroid[0] / frame_shape[1] * GRID_SIZE)
+    grid_y = int(centroid[1] / frame_shape[0] * GRID_SIZE)
 
     cell = region_grid[grid_x][grid_y]