add option to set tracker algorithm

Author: arch

docs/app/docs/user-guide/config.md

@@ -26,6 +26,7 @@ Config Files:
 - `max_playback_fps` (int): Limit the max player speed in the tracking preview window (0 = disable limit)
 - `preview_scaling` (float): Set the preview image scaling factor. With a value of `1.0`, the window should fill the height or width of the screen depending on the aspect ratio of the video.
 - `projection`: (str): Set the video type. All available options can be obtained from the `projection.yaml` file. Each root keys represent an option. Available Options: `'flat'` = 2D Videos, `'vr_he_sbs'` = 3D VR Side-By-Side Video
+- `tracker`: (str) Specify the tracker algorithm. Available options are `'MIL'`, `'KCF'`, `'CSRT'`.
 
 #### `hyperparameter.yaml`
 

funscript_editor/algorithms/videotracker.py

@@ -6,6 +6,7 @@
 
 from threading import Thread
 from queue import Queue
+from funscript_editor.utils.config import SETTINGS
 
 import numpy as np
 
@@ -66,9 +67,23 @@ def result(self) -> tuple:
         return self.queue_out.get()
 
 
+    def __setup_tracker(self) -> None:
+        """ Setup the tracker specified in the config """
+        if SETTINGS['tracker'].upper() == 'MIL':
+            self.__logger.info("Start MIL Tracker")
+            self.tracker = cv2.TrackerMIL_create()
+        elif SETTINGS['tracker'].upper() == 'KCF':
+            self.__logger.info("Start KCF Tracker")
+            self.tracker = cv2.TrackerKCF_create()
+        else:
+            # falback is CSRT tracker
+            self.__logger.info("Start CSRT Tracker")
+            self.tracker = cv2.TrackerCSRT_create()
+
+
     def run(self) -> None:
         """ The Video Tracker Thread Function """
-        self.tracker = cv2.TrackerCSRT_create() # NOTE: you can change this to your favorite tracker
+        self.__setup_tracker()
         frame_heigt, frame_width = self.first_frame.shape[:2]
 
         dh, dw = int(frame_heigt*self.limit_searchspace['h']), int(frame_width*self.limit_searchspace['w'])

funscript_editor/config/settings.yaml

@@ -27,3 +27,6 @@ preview_scaling: 0.6
 # - 'flat': 2D Videos
 # - 'vr_he_sbs': 3D VR Side-By-Side Video
 projection: 'vr_he_sbs'
+
+# Specify the tracker algorithm. Available options are 'MIL', 'KCF', 'CSRT'.
+tracker: 'CSRT'

funscript_editor/data/ffmpegstream.py

@@ -279,6 +279,7 @@ def run(self) -> None:
                 '-'
             ]
 
+        self.logger.info("Open FFmpeg Stream")
         pipe = sp.Popen(
                 command,
                 stdout = sp.PIPE,