From 50f113277b10056c4cb784952644f1d2f0ac95c0 Mon Sep 17 00:00:00 2001
From: Josh Hawkins <32435876+hawkeye217@users.noreply.github.com>
Date: Tue, 11 Jul 2023 12:17:58 -0500
Subject: [PATCH] reset motion estimator after returning to preset

---
 frigate/app.py                   |  1 +
 frigate/ptz/autotrack.py         | 29 ++++++++++++++++++++++-------
 frigate/track/norfair_tracker.py | 12 +++++++++++-
 frigate/types.py                 |  1 +
 4 files changed, 35 insertions(+), 8 deletions(-)

diff --git a/frigate/app.py b/frigate/app.py
index ecbf7c2e0..87a55aca2 100644
--- a/frigate/app.py
+++ b/frigate/app.py
@@ -172,6 +172,7 @@ class FrigateApp:
                     self.config.cameras[camera_name].onvif.autotracking.enabled,
                 ),
                 "ptz_stopped": mp.Event(),
+                "ptz_reset": mp.Event(),
                 "ptz_start_time": mp.Value("d", 0.0),  # type: ignore[typeddict-item]
                 # issue https://github.com/python/typeshed/issues/8799
                 # from mypy 0.981 onwards
diff --git a/frigate/ptz/autotrack.py b/frigate/ptz/autotrack.py
index 9421e36f9..721b15268 100644
--- a/frigate/ptz/autotrack.py
+++ b/frigate/ptz/autotrack.py
@@ -36,22 +36,33 @@ def ptz_moving_at_frame_time(frame_time, ptz_start_time, ptz_stop_time):
 
 
 class PtzMotionEstimator:
-    def __init__(self, config: CameraConfig, ptz_metrics: PTZMetricsTypes) -> None:
+    def __init__(
+        self, config: CameraConfig, ptz_metrics: dict[str, PTZMetricsTypes]
+    ) -> None:
         self.frame_manager = SharedMemoryFrameManager()
-        # homography is nice (zooming) but slow, translation is pan/tilt only but fast.
-        self.norfair_motion_estimator = MotionEstimator(
-            transformations_getter=TranslationTransformationGetter(),
-            min_distance=30,
-            max_points=900,
-        )
+        self.norfair_motion_estimator = None
         self.camera_config = config
         self.coord_transformations = None
         self.ptz_metrics = ptz_metrics
         self.ptz_start_time = self.ptz_metrics["ptz_start_time"]
         self.ptz_stop_time = self.ptz_metrics["ptz_stop_time"]
+
+        self.ptz_metrics["ptz_reset"].set()
         logger.debug(f"Motion estimator init for cam: {config.name}")
 
     def motion_estimator(self, detections, frame_time, camera_name):
+        # If we've just started up or returned to our preset, reset motion estimator for new tracking session
+        if self.ptz_metrics["ptz_reset"].is_set():
+            self.ptz_metrics["ptz_reset"].clear()
+            logger.debug("Motion estimator reset")
+            # homography is nice (zooming) but slow, translation is pan/tilt only but fast.
+            self.norfair_motion_estimator = MotionEstimator(
+                transformations_getter=TranslationTransformationGetter(),
+                min_distance=30,
+                max_points=900,
+            )
+            self.coord_transformations = None
+
         if ptz_moving_at_frame_time(
             frame_time, self.ptz_start_time.value, self.ptz_stop_time.value
         ):
@@ -242,6 +253,9 @@ class PtzAutoTracker:
         camera_config = self.config.cameras[camera]
 
         if camera_config.onvif.autotracking.enabled:
+            if not self.autotracker_init[camera]:
+                self._autotracker_setup(self.config.cameras[camera], camera)
+
             # either this is a brand new object that's on our camera, has our label, entered the zone, is not a false positive,
             # and is not initially motionless - or one we're already tracking, which assumes all those things are already true
             if (
@@ -404,4 +418,5 @@ class PtzAutoTracker:
                 camera,
                 autotracker_config.return_preset.lower(),
             )
+            self.ptz_metrics[camera]["ptz_reset"].set()
             self.tracked_object_previous[camera] = None
diff --git a/frigate/track/norfair_tracker.py b/frigate/track/norfair_tracker.py
index 5e5d764e3..2e9e5034f 100644
--- a/frigate/track/norfair_tracker.py
+++ b/frigate/track/norfair_tracker.py
@@ -67,7 +67,9 @@ class NorfairTracker(ObjectTracker):
         self.max_disappeared = config.detect.max_disappeared
         self.camera_config = config
         self.detect_config = config.detect
+        self.ptz_metrics = ptz_metrics
         self.ptz_autotracker_enabled = ptz_metrics["ptz_autotracker_enabled"]
+        self.ptz_motion_estimator = {}
         self.camera_name = config.name
         self.track_id_map = {}
         # TODO: could also initialize a tracker per object class if there
@@ -79,7 +81,9 @@ class NorfairTracker(ObjectTracker):
             hit_counter_max=self.max_disappeared,
         )
         if self.ptz_autotracker_enabled.value:
-            self.ptz_motion_estimator = PtzMotionEstimator(config, ptz_metrics)
+            self.ptz_motion_estimator = PtzMotionEstimator(
+                self.camera_config, self.ptz_metrics
+            )
 
     def register(self, track_id, obj):
         rand_id = "".join(random.choices(string.ascii_lowercase + string.digits, k=6))
@@ -244,6 +248,12 @@ class NorfairTracker(ObjectTracker):
         coord_transformations = None
 
         if self.ptz_autotracker_enabled.value:
+            # we must have been enabled by mqtt, so set up the estimator
+            if not self.ptz_motion_estimator:
+                self.ptz_motion_estimator = PtzMotionEstimator(
+                    self.camera_config, self.ptz_metrics
+                )
+
             coord_transformations = self.ptz_motion_estimator.motion_estimator(
                 detections, frame_time, self.camera_name
             )
diff --git a/frigate/types.py b/frigate/types.py
index 3940ff056..05181d714 100644
--- a/frigate/types.py
+++ b/frigate/types.py
@@ -29,6 +29,7 @@ class CameraMetricsTypes(TypedDict):
 class PTZMetricsTypes(TypedDict):
     ptz_autotracker_enabled: Synchronized
     ptz_stopped: Event
+    ptz_reset: Event
     ptz_start_time: Synchronized
     ptz_stop_time: Synchronized