add prototype for kalman filter

Author: arch

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

@@ -19,6 +19,7 @@ Config Files:
 
 #### `settings.yaml`
 
+- `use_kalman_filter`: Enable Kalman Filter
 - `use_zoom` (bool): Enable or disable an additional step to zoom in the Video before selecting a tracking feature for the Woman or Men.
 - `zoom_factor:` (float): Set the desired zoom value which will be used when the zoom function is activated.
 - `tracking_direction` (str): Specify the tracking direction. Allowed values are `'x'` and `'y'`.

funscript_editor/algorithms/funscriptgenerator.py

@@ -18,6 +18,7 @@
 from funscript_editor.utils.config import HYPERPARAMETER, SETTINGS, PROJECTION
 from datetime import datetime
 from funscript_editor.data.ffmpegstream import FFmpegStream, VideoInfo
+from funscript_editor.algorithms.kalmanfilter import KalmanFilter2D
 
 import funscript_editor.algorithms.signalprocessing as sp
 import numpy as np
@@ -44,6 +45,7 @@ class FunscriptGeneratorParameter:
     bottom_threshold: float = float(HYPERPARAMETER['bottom_threshold'])
     preview_scaling: float = float(SETTINGS['preview_scaling'])
     projection: str = str(SETTINGS['projection']).lower()
+    use_kalman_filter: bool = SETTINGS['use_kalman_filter']
 
 
 class FunscriptGenerator(QtCore.QThread):
@@ -403,13 +405,13 @@ def delete_last_tracking_predictions(self, num :int) -> None:
         Args:
             num (int): number of frames to remove from predicted boxes
         """
-        if len(self.bboxes['Woman']) <= num-1:
+        if len(self.bboxes['Woman']) <= num:
             self.bboxes['Woman'] = []
             self.bboxes['Men'] = []
         else:
-            for i in range(len(self.bboxes['Woman'])-1,len(self.bboxes['Woman'])-num,-1):
-                del self.bboxes['Woman'][i]
-                if self.params.track_men: del self.bboxes['Men'][i]
+            for _ in range(num):
+                del self.bboxes['Woman'][-1]
+                if self.params.track_men: del self.bboxes['Men'][-1]
 
 
     def preview_scaling(self, preview_image :np.ndarray) -> np.ndarray:
@@ -438,7 +440,7 @@ def get_vr_projection_config(self, image :np.ndarray) -> dict:
         Returns:
             dict: projection config
         """
-        config = PROJECTION[self.params.projection]
+        config = copy.deepcopy(PROJECTION[self.params.projection])
 
         self.determine_preview_scaling(config['parameter']['width'], config['parameter']['height'])
 
@@ -639,7 +641,7 @@ def tracking(self) -> str:
 
                 if self.was_key_pressed('q') or cv2.waitKey(1) == ord('q'):
                     status = 'Tracking stopped by user'
-                    self.delete_last_tracking_predictions(int(self.get_average_tracking_fps()+1)*3)
+                    self.delete_last_tracking_predictions(int((self.get_average_tracking_fps()+1)*2.2))
                     break
 
             (successWoman, bboxWoman) = trackerWoman.result()
@@ -663,6 +665,7 @@ def tracking(self) -> str:
 
         video.stop()
         self.logger.info(status)
+        self.logger.info('Calculate score')
         self.calculate_score()
         return status
 
@@ -762,12 +765,38 @@ def get_score_with_offset(self, idx_dict) -> list:
         return score
 
 
+    def apply_kalman_filter(self) -> None:
+        """ Apply Kalman Filter to the tracking prediction """
+        if len(self.bboxes['Woman']) < self.video_info.fps: return
+
+        # TODO: we should use the center of the tracking box not x0,y0 of the box
+
+        self.logger.info("Apply kalman filter")
+        kalman = KalmanFilter2D(self.video_info.fps)
+        kalman.init(self.bboxes['Woman'][0][0], self.bboxes['Woman'][0][1])
+        for idx, item in enumerate(self.bboxes['Woman']):
+            prediction = kalman.update(item[0], item[1])
+            self.bboxes['Woman'][idx] = (prediction[0], prediction[1], item[2], item[3])
+
+        if self.params.track_men:
+            kalman = KalmanFilter2D(self.video_info.fps)
+            kalman.init(self.bboxes['Men'][0][0], self.bboxes['Men'][0][1])
+            for idx, item in enumerate(self.bboxes['Men']):
+                prediction = kalman.update(item[0], item[1])
+                self.bboxes['Men'][idx] = (prediction[0], prediction[1], item[2], item[3])
+
+
     def run(self) -> None:
         """ The Funscript Generator Thread Function """
         # NOTE: score['y'] and score['x'] should have the same number size so it should be enouth to check one score length
         with Listener(on_press=self.on_key_press) as listener:
             status = self.tracking()
+
+            if self.params.use_kalman_filter:
+                self.apply_kalman_filter()
+
             if len(self.score['y']) >= HYPERPARAMETER['min_frames']:
+                self.logger.info("Scale score")
                 if self.params.direction != 'x':
                     self.scale_score(status, direction='y')
                 else:
@@ -777,6 +806,7 @@ def run(self) -> None:
             self.finished(status + ' -> Tracking time insufficient', False)
             return
 
+        self.logger.info("Determine local max and min")
         if self.params.direction != 'x':
             idx_dict = sp.get_local_max_and_min_idx(self.score['y'], self.video_info.fps)
         else:

funscript_editor/algorithms/kalmanfilter.py

@@ -0,0 +1,76 @@
+""" Kalman Filter """
+
+import cv2
+
+import numpy as np
+
+
+class KalmanFilter2D:
+    """ Kalman 2D Filter
+
+    Args:
+        fps (float): frames per second
+    """
+
+    def __init__(self, fps: float):
+        self.kalman = cv2.KalmanFilter(4, 2)
+        self.prediction = np.zeros((2, 1), np.float32)
+        self.prediction_counter = 0
+        self.fps = fps
+
+
+    def init(self, x: float, y: float) -> None:
+        """ Initialize the 2D Kalman Filter
+
+        Args:
+            x (float): x measurement
+            y (float): y measurement
+        """
+        # dt = 1.0 / self.fps
+        dt = 1.0
+
+        self.kalman.measurementMatrix = np.array(
+                [
+                    [1, 0, 0, 0],
+                    [0, 1, 0, 0]
+                ], np.float32)
+
+        self.kalman.transitionMatrix = np.array(
+                [
+                    [1, 0, dt,0 ],
+                    [0, 1, 0, dt],
+                    [0, 0, 1, 0 ],
+                    [0, 0, 0, 1 ]
+                ], np.float32)
+
+        self.kalman.processNoiseCov = np.array(
+                [
+                    [1, 0, 0, 0],
+                    [0, 1, 0, 0],
+                    [0, 0, 1, 0],
+                    [0, 0, 0, 1]
+                ], np.float32) * 0.03
+
+        self.kalman.correct(np.array([np.float32(x-1), np.float32(y-1)], np.float32))
+        self.prediction = self.kalman.predict()
+        self.prediction_counter += 1
+
+
+    def update(self, x: float, y: float) -> np.ndarray:
+        """ Update the 2D Kalman Filter
+
+        Args:
+            x (float): x measurement
+            y (float): y measurement
+
+        Returns:
+            list: prediction [x', y', vx', vy']
+        """
+        self.kalman.correct(np.array([np.float32(x), np.float32(y)], np.float32))
+        self.prediction = self.kalman.predict()
+        self.prediction_counter += 1
+        # TODO howo to init/fit the OpenCV Kalman Filter?
+        if self.prediction_counter < self.fps:
+            return [round(x), round(y), 0.0, 0.0]
+        else:
+            return [round(item[0]) if idx < 2 else float(item[0]) for idx, item in enumerate(self.prediction)]

funscript_editor/config/hyperparameter.yaml

@@ -11,7 +11,7 @@ avg_sec_for_local_min_max_extraction: 1.9
 
 # Specify the minimum required frames for the tracking. Wee need this parameter to
 # ensure there is at leas two strokes in the tracking result.
-min_frames: 120
+min_frames: 100
 
 # Shift predicted top points by given frame number. Positive values delay the position
 # and negative values result in an earlier position.

funscript_editor/config/settings.yaml

@@ -1,5 +1,8 @@
 # Python Funscript Editor settings
 
+# Enable Kalman Filter
+use_kalman_filter: False
+
 # Enable or disable an additional step to zoom in the Video before selecting a tracking
 # feature for the Woman or Men.
 use_zoom: False

funscript_editor/data/ffmpegstream.py

@@ -155,9 +155,9 @@ def get_projection(
                     (config['parameter']['height'], config['parameter']['width'], 3)
                 )
 
+        pipe.terminate()
         pipe.stdin.close()
         pipe.stdout.close()
-        pipe.terminate()
 
         return projection
 
@@ -218,7 +218,6 @@ def millisec_to_timestamp(millis :int)->str:
     def stop(self) -> None:
         """ Stop FFmpeg video stream """
         self.stopped = True
-        self.thread.join()
 
 
     def read(self) -> np.ndarray:
@@ -305,5 +304,5 @@ def run(self) -> None:
 
         self.stopped = True
         self.logger.info('Close FFmpeg Stream')
-        pipe.stdout.close()
         pipe.terminate()
+        pipe.stdout.close()