add simple shift point function

Author: arch

funscript_editor/algorithms/funscriptgenerator.py

@@ -15,7 +15,7 @@
 from PyQt5 import QtGui, QtCore, QtWidgets
 from matplotlib.figure import Figure
 from funscript_editor.definitions import VIDEO_SCALING_CONFIG_FILE
-from funscript_editor.utils.config import HYPERPARAMETER
+from funscript_editor.utils.config import HYPERPARAMETER, SETTINGS
 from datetime import datetime
 
 import funscript_editor.algorithms.signalprocessing as sp
@@ -29,9 +29,11 @@ class FunscriptGeneratorParameter:
     video_path: str # no default value
     start_frame: int = 0 # default is video start
     skip_frames: int = HYPERPARAMETER['skip_frames']
-    max_playback_fps: int = HYPERPARAMETER['max_playback_fps']
-    direction: str = HYPERPARAMETER['tracking_direction']
-    use_zoom: bool = HYPERPARAMETER['use_zoom']
+    max_playback_fps: int = SETTINGS['max_playback_fps']
+    direction: str = SETTINGS['tracking_direction']
+    use_zoom: bool = SETTINGS['use_zoom']
+    shift_bottom_points :int = HYPERPARAMETER['shift_bottom_points']
+    shift_top_points :int = HYPERPARAMETER['shift_top_points']
     track_men: bool = True
 
 
@@ -324,7 +326,7 @@ def plot_y_score(self, name: str, idx_list: list, dpi : int = 300) -> None:
 
         Args:
             name (str): file name for the figure
-            idx_list (list): list with funscript action points
+            idx_list (list): list with all frame numbers with funscript action points
             dpi (int): picture output dpi
         """
         if len(self.score_y) < 2: return
@@ -432,6 +434,9 @@ def tracking(self) -> str:
         if first_frame is None:
             return
 
+        if self.params.skip_frames < 0:
+            self.params.skip_frames = 0
+
         bboxWoman = self.get_bbox(first_frame, "Select Woman Feature")
         trackerWoman = StaticVideoTracker(first_frame, bboxWoman)
         self.bboxes['Woman'].append(bboxWoman)
@@ -441,7 +446,7 @@ def tracking(self) -> str:
             trackerMen = StaticVideoTracker(first_frame, bboxMen)
             self.bboxes['Men'].append(bboxMen)
 
-        if self.params.max_playback_fps > 2:
+        if self.params.max_playback_fps > (self.params.skip_frames+1):
             cycle_time_in_ms = (float(1000) / float(self.params.max_playback_fps)) * (self.params.skip_frames+1)
         else:
             cycle_time_in_ms = 0
@@ -562,6 +567,23 @@ def finished(self, status: str, success :bool) -> None:
         self.stopped = True
 
 
+    def apply_shift(self, frame_number, position: str) -> int:
+        """ Apply shift to predicted frame positions
+
+        Args:
+            position (str): is max or min
+        """
+        if position in ['max', 'top'] and self.params.direction != 'x':
+            if frame_number >= -1*self.params.shift_top_points and frame_number + self.params.shift_top_points < len(self.score_y):
+                return self.params.start_frame + frame_number + self.params.shift_top_points
+
+        if position in ['min', 'bottom'] and self.params.direction != 'x':
+            if frame_number >= -1*self.params.shift_bottom_points and frame_number + self.params.shift_bottom_points < len(self.score_y):
+                return self.params.start_frame + frame_number + self.params.shift_bottom_points
+
+        return self.params.start_frame + frame_number
+
+
     def run(self) -> None:
         """ The Funscript Generator Thread Function """
         # NOTE: score_y and score_x should have the same number of elements so it should be enouth to check one score length
@@ -578,18 +600,28 @@ def run(self) -> None:
             return
 
         if self.params.direction != 'x':
-            idx_list = sp.get_local_max_and_min_idx(self.score_y, self.fps)
+            idx_dict = sp.get_local_max_and_min_idx(self.score_y, self.fps)
         else:
-            idx_list = sp.get_local_max_and_min_idx(self.score_x, self.fps)
+            idx_dict = sp.get_local_max_and_min_idx(self.score_x, self.fps)
+
+        idx_list = [x for k in ['min', 'max'] for x in idx_dict[k]]
+        idx_list.sort()
 
         if False:
-            if self.params.direction != 'x': self.plot_y_score('debug_001.png', idx_list)
-            self.plot_scores('debug_002.png')
+            self.plot_scores('debug_001.png')
+            if self.params.direction != 'x':
+                self.plot_y_score('debug_002.png', idx_list)
+
+        for idx in idx_dict['min']:
+            self.funscript.add_action(
+                    self.score_y[idx],
+                    self.frame_to_millisec(self.apply_shift(idx, 'min'), self.fps)
+                )
 
-        for idx in idx_list:
+        for idx in idx_dict['max']:
             self.funscript.add_action(
                     self.score_y[idx],
-                    self.frame_to_millisec(idx+self.params.start_frame, self.fps)
+                    self.frame_to_millisec(self.apply_shift(idx, 'max'), self.fps)
                 )
 
         self.finished(status, True)

funscript_editor/algorithms/signalprocessing.py

@@ -62,30 +62,39 @@ def moving_average(x :list, w: int) -> list:
             +[avg[-1] for _ in range(len(avg)+w,len(x))]
 
 
-def get_local_max_and_min_idx(score :list, fps: int) -> list:
+def get_local_max_and_min_idx(score :list, fps: int, shift_min :int = 0, shift_max :int = 0) -> dict:
     """ Get the local max and min positions in given signal
 
     Args:
         score (list): list with float or int signal values
         fps (int): rounded fps of the video
+        shift_min (int): shift the local min indexes by given value (default is 0)
+        shift_max (int): shift the local max indexes by given value (default is 0)
 
     Returns:
-        list: list with all local max and min indexes
+        dict: dict with 2 lists with all local max and min indexes ({'min':[], 'max':[]})
     """
     avg = moving_average(score, w=round(fps * HYPERPARAMETER['avg_sec_for_local_min_max_extraction']))
-    tmp_min_idx, tmp_max_idx, idx_list = -1, -1, []
+    result = {'min': [], 'max': []}
+    tmp_min_idx, tmp_max_idx = -1, -1
     for pos in range(len(score)):
         if score[pos] < avg[pos]:
             if tmp_min_idx < 0: tmp_min_idx = pos
             elif score[tmp_min_idx] > score[pos]: tmp_min_idx = pos
         elif tmp_min_idx >= 0:
-            idx_list.append(tmp_min_idx)
+            if tmp_min_idx >= -1*shift_min and tmp_min_idx + shift_min < len(score):
+                result['min'].append(tmp_min_idx + shift_min)
+            else:
+                result['min'].append(tmp_min_idx)
             tmp_min_idx = -1
 
         if score[pos] > avg[pos]:
             if tmp_max_idx < 0: tmp_max_idx = pos
             elif score[tmp_max_idx] < score[pos]: tmp_max_idx = pos
         elif tmp_max_idx >= 0:
-            idx_list.append(tmp_max_idx)
+            if tmp_max_idx >= -1*shift_max and tmp_max_idx + shift_max < len(score):
+                result['max'].append(tmp_max_idx + shift_max)
+            else:
+                result['max'].append(tmp_max_idx)
             tmp_max_idx = -1
-    return idx_list
+    return result

funscript_editor/config/hyperparameter.yaml

@@ -1,6 +1,13 @@
+# interpolate given frames to increase the processing speed
 skip_frames: 1
+
+# specify the window size for the calculation of the reference value for the local min and max search.
 avg_sec_for_local_min_max_extraction: 1.9
+
+# specify the minimum required frames for the tracking.
 min_frames: 120
-use_zoom: False
-tracking_direction: 'y'
-max_playback_fps: 0
+
+# shift predicted points by given frame number.
+# positive values delay the position and negative values result in an earlier position.
+shift_top_points: 0
+shift_bottom_points: 0

funscript_editor/config/settings.yaml

@@ -0,0 +1,3 @@
+use_zoom: False
+tracking_direction: 'y'
+max_playback_fps: 0

funscript_editor/definitions.py

@@ -38,6 +38,7 @@
 WINDOWS_LOG_CONFIG_FILE = os.path.join(CONFIG_DIR, 'logging_windows.yaml')
 LINUX_LOG_CONFIG_FILE = os.path.join(CONFIG_DIR, 'logging_linux.yaml')
 HYPERPARAMETER_CONFIG_FILE = os.path.join(CONFIG_DIR, 'hyperparameter.yaml')
+SETTINGS_CONFIG_FILE = os.path.join(CONFIG_DIR, 'settings.yaml')
 
 
 ##########

funscript_editor/ui/video_player.py

@@ -291,7 +291,7 @@ def increase_speed(self) -> None:
 
     def decrease_speed(self) -> None:
         """ Decrease the playback speed """
-        self.player.speed -= 0.2
+        if self.player.speed > 0.21: self.player.speed -= 0.2
         if 0.9 < self.player.speed < 1.1: self.player.speed = 1.0
         self.show_message('Speed: ' + str(self.player.speed))
 

funscript_editor/utils/config.py

@@ -3,7 +3,8 @@
 import os
 import yaml
 
-from funscript_editor.definitions import ROOT_DIR, UI_CONFIG_FILE, HYPERPARAMETER_CONFIG_FILE
+from funscript_editor.definitions import ROOT_DIR, UI_CONFIG_FILE, \
+        HYPERPARAMETER_CONFIG_FILE, SETTINGS_CONFIG_FILE
 
 def read_yaml_config(config_file: str) -> dict:
     """ Parse a yaml config file
@@ -37,3 +38,6 @@ def read_version() -> str:
 
 #: hyperparameter for the algorithms
 HYPERPARAMETER = read_yaml_config(HYPERPARAMETER_CONFIG_FILE)
+
+#: settings
+SETTINGS = read_yaml_config(SETTINGS_CONFIG_FILE)