eliminate duplicated code

Author: arch

funscript_editor/algorithms/funscriptgenerator.py

@@ -67,7 +67,7 @@ class FunscriptGeneratorParameter:
     right_threshold: float = float(HYPERPARAMETER['right_threshold'])
 
 
-class FunscriptGenerator(QtCore.QThread):
+class FunscriptGeneratorThread(QtCore.QThread):
     """ Funscript Generator Thread
 
     Args:
@@ -105,9 +105,6 @@ def __init__(self,
     #: completed event with reference to the funscript with the predicted actions, status message and success flag
     funscriptCompleted = QtCore.pyqtSignal(object, str, bool)
 
-    #: processing event with current processed frame number
-    processStatus = QtCore.pyqtSignal(int)
-
     logger = logging.getLogger(__name__)
 
 
@@ -687,7 +684,6 @@ def tracking(self) -> str:
 
             trackerWoman.update(frame)
             if self.params.track_men: trackerMen.update(frame)
-            self.processStatus.emit(frame_num)
 
             if last_frame is not None:
                 # Process data from last step while the next tracking points get predicted.

funscript_editor/api.py

@@ -5,7 +5,7 @@
 
 from funscript_editor.utils.logging import setup_logging
 from funscript_editor.ui.funscript_editor_window import FunscriptEditorWindow
-from funscript_editor.ui.minimal import MinimalFunscriptGenerator
+from funscript_editor.ui.funscript_generator_window import FunscriptGeneratorWindow
 from funscript_editor.utils.config import VERSION
 
 from PyQt5 import QtCore, QtGui, QtWidgets
@@ -40,6 +40,6 @@ def generate_funscript(
     logging.info("Args: video_file=%s, start_time=%s, end_time=%s, output_file=%s", \
             str(video_file), str(start_time), str(end_time), str(output_file))
     app = QtWidgets.QApplication(sys.argv)
-    generator = MinimalFunscriptGenerator(video_file, start_time, end_time, output_file)
+    generator = FunscriptGeneratorWindow(video_file, start_time, end_time, output_file)
     generator.run()
     sys.exit(app.exec_())

funscript_editor/config/settings.yaml

@@ -23,7 +23,7 @@ preview_scaling: 0.6
 # Set the video type. Available Options:
 # - 'flat': 2D Videos
 # - 'vr_he_sbs': 3D VR Side-By-Side Video
-projection: 'vr_he_sbs'
+projection: 'flat'
 
 # Specify the tracker algorithm. Available options are 'MIL', 'KCF', 'CSRT'.
 tracker: 'CSRT'

funscript_editor/ui/funscript_editor_window.py

@@ -21,10 +21,11 @@
 from funscript_editor.ui.video_player import VideoPlayer
 from funscript_editor.ui.funscript_visualizer import FunscriptVisualizer
 from funscript_editor.data.funscript import Funscript
-from funscript_editor.algorithms.funscriptgenerator import FunscriptGenerator, FunscriptGeneratorParameter
+from funscript_editor.algorithms.funscriptgenerator import FunscriptGeneratorThread, FunscriptGeneratorParameter
 from funscript_editor.utils.config import UI_CONFIG, VERSION
 from funscript_editor.definitions import APP_DOCUMENTATION_DIR, CODE_DOCUMENTATION_DIR
 from funscript_editor.data.ffmpegstream import FFmpegStream
+from funscript_editor.ui.funscript_generator_window import FunscriptGeneratorWindow
 
 class FunscriptEditorWindow(QtWidgets.QMainWindow):
     """ Funscript Editor window """
@@ -295,38 +296,28 @@ def __new_funscript(self):
         self.video_player.set_funscript(self.funscript)
         self.funscript_visualizer.set_funscript(self.funscript)
 
-    def __generator_status_changed(self, current):
-        self.statusBar().showMessage("{} ({})".format(current, datetime.now().strftime("%H:%M:%S")))
-
     def __generate_funscript(self):
         if self.funscript is None: return
         if self.video_player is None: return
         if self.video_player.get_video_file is None: return
-        start_frame = self.video_player.get_current_frame
+        start_time = self.video_player.get_current_timestamp_in_millis
         next_action = self.funscript.get_next_action(self.video_player.get_current_timestamp_in_millis+100)
         if next_action['at'] > self.video_player.get_current_timestamp_in_millis+100:
-            end_frame = self.video_player.millisec_to_frame(next_action['at'])
+            end_time = next_action['at']
         else:
-            end_frame = -1
+            end_time = -1
 
-        self.__logger.info("Stop at Frame {}".format(end_frame))
+        self.__logger.info("Stop at {}".format(end_time))
+
+        self.funscript_generator_window = FunscriptGeneratorWindow(
+                self.video_player.get_video_file,
+                start_time,
+                end_time,
+                self.funscript
+            )
+        self.funscript_generator_window.funscriptCompleted.connect(self.__funscript_generated)
+        self.funscript_generator_window.run()
 
-        reply = QtWidgets.QMessageBox.question(None, 'Track Men', 'Do you want to track the Men? ',
-                QtWidgets.QMessageBox.Yes | QtWidgets.QMessageBox.No, QtWidgets.QMessageBox.No)
-        trackMen = True if reply == QtWidgets.QMessageBox.Yes else False
-
-        self.video_player.set_indicate_bussy(True)
-        self.funscript_generator = FunscriptGenerator(
-                FunscriptGeneratorParameter(
-                    video_path = self.video_player.get_video_file,
-                    start_frame = start_frame,
-                    end_frame = end_frame,
-                    track_men = trackMen
-                ),
-                self.funscript)
-        self.funscript_generator.funscriptCompleted.connect(self.__funscript_generated)
-        self.funscript_generator.processStatus.connect(self.__generator_status_changed)
-        self.funscript_generator.start()
 
     def __funscript_generated(self, funscript, status, success):
         self.video_player.set_funscript(self.funscript)

funscript_editor/ui/minimal.py …_editor/ui/funscript_generator_window.pyfunscript_editor/ui/minimal.py renamed to funscript_editor/ui/funscript_generator_window.py funscript_editor/ui/minimal.py renamed to funscript_editor/ui/funscript_generator_window.py

@@ -4,14 +4,13 @@
 import cv2
 
 from funscript_editor.utils.logging import setup_logging
-from funscript_editor.ui.funscript_editor_window import FunscriptEditorWindow
-from funscript_editor.algorithms.funscriptgenerator import FunscriptGenerator, FunscriptGeneratorParameter
+from funscript_editor.algorithms.funscriptgenerator import FunscriptGeneratorThread, FunscriptGeneratorParameter
 from funscript_editor.data.funscript import Funscript
 
 from PyQt5 import QtCore, QtGui, QtWidgets
 
 
-class MinimalFunscriptGenerator(QtWidgets.QMainWindow):
+class FunscriptGeneratorWindow(QtWidgets.QMainWindow):
     """ Class to Generate a funscript with minimal UI
 
     Note:
@@ -21,22 +20,24 @@ class MinimalFunscriptGenerator(QtWidgets.QMainWindow):
         video_file (str): path to video file
         start_time (float): start position in video (timestamp in milliseconds)
         end_time (float): end position in video (timestamp in milliseconds) use -1.0 for video end.
-        output_file (str): csv output file path
+        output_file (str, Funscript): csv output file path (Optional you can pass a funscript object where to store the result)
     """
 
     def __init__(self,
             video_file: str,
             start_time: float,
             end_time: float,
             output_file: str):
-        super(MinimalFunscriptGenerator, self).__init__()
+        super(FunscriptGeneratorWindow, self).__init__()
 
-        if os.path.isdir(output_file):
-            self.__show_message("The output TempFile path must be a file not a folder", error=True)
-            sys.exit()
+        if not isinstance(output_file, Funscript):
+            output_file = os.path.abspath(output_file)
+            if os.path.isdir(output_file):
+                self.__show_message("The output TempFile path must be a file not a folder", error=True)
+                sys.exit()
 
-        if os.path.exists(output_file):
-            os.remove(output_file)
+            if os.path.exists(output_file):
+                os.remove(output_file)
 
         if video_file is None or video_file == "":
             self.__show_message("Video file was not specified! " \
@@ -53,7 +54,7 @@ def __init__(self,
         cap.release()
 
         self.funscript = Funscript(fps)
-        self.output_file = os.path.abspath(output_file)
+        self.output_file = output_file
 
         if False:
             reply = QtWidgets.QMessageBox.question(None, 'Generate Funscript ', \
@@ -72,7 +73,7 @@ def __init__(self,
 
         self.__logger.info("Set End Time to Frame Number %d", end_frame)
 
-        self.funscript_generator = FunscriptGenerator(
+        self.funscript_generator = FunscriptGeneratorThread(
                 FunscriptGeneratorParameter(
                     video_path = video_file,
                     start_frame = start_frame,
@@ -81,11 +82,12 @@ def __init__(self,
                 ),
                 self.funscript)
         self.funscript_generator.funscriptCompleted.connect(self.__funscript_generated)
-        self.funscript_generator.processStatus.connect(self.__status_changed)
 
 
     __logger = logging.getLogger(__name__)
 
+    funscriptCompleted = QtCore.pyqtSignal(object, str, bool)
+
 
     def __show_message(self, message :str, error: bool = False) -> None:
         if error: self.__logger.error(message)
@@ -98,17 +100,18 @@ def __show_message(self, message :str, error: bool = False) -> None:
 
 
     def __funscript_generated(self, funscript, msg, success) -> None:
-        os.makedirs(os.path.dirname(self.output_file), exist_ok=True)
-        with open(self.output_file, 'w') as f:
-            f.write('at;pos\n')
+        if isinstance(self.output_file, Funscript):
             for item in funscript.get_actions():
-                f.write('{at};{pos}\n'.format(at=item['at'], pos=item['pos']))
-        if not success: self.__show_message(msg, error=True)
-        sys.exit()
-
-
-    def __status_changed(self, frame) -> None:
-        pass
+                self.output_file.add_action(item['pos'], item['at'])
+            self.funscriptCompleted.emit(self.output_file, msg, success)
+        else:
+            os.makedirs(os.path.dirname(self.output_file), exist_ok=True)
+            with open(self.output_file, 'w') as f:
+                f.write('at;pos\n')
+                for item in funscript.get_actions():
+                    f.write('{at};{pos}\n'.format(at=item['at'], pos=item['pos']))
+            if not success: self.__show_message(msg, error=True)
+            sys.exit()
 
 
     def run(self) -> None: