add option for muliaxis json output

Author: arch

funscript_editor/__main__.py

@@ -9,8 +9,9 @@ def main():
     """ CLI Main Function """
     parser = argparse.ArgumentParser()
     parser.add_argument("--generator", action = 'store_true', help = "Run only the generator")
+    parser.add_argument("--multiaxis", action = 'store_true', help = "Show options for multiaxis output")
     parser.add_argument("-i", "--input", type = str, help = "Video File")
-    parser.add_argument("-o", "--output", type = str, default = "/tmp/funscript_actions.csv", help = "Output Path")
+    parser.add_argument("-o", "--output", type = str, default = "/tmp/funscript_actions.json", help = "Output Path")
     parser.add_argument("-s", "--start", type = float, default = 0.0, help = "Start Time in Milliseconds")
     parser.add_argument("-e", "--end", type = float, default = -1.0, help = "End/Stop Time in Milliseconds")
     args = parser.parse_args()
@@ -19,4 +20,4 @@ def main():
         os.environ['PATH'] = os.getcwd() + os.sep + os.environ['PATH']
 
     if not args.generator: show_editor()
-    else: generate_funscript(args.input, args.start, args.end, args.output)
+    else: generate_funscript(args.input, args.start, args.end, args.output, args.multiaxis)

funscript_editor/api.py

@@ -26,20 +26,22 @@ def generate_funscript(
         video_file: str,
         start_time: float,
         end_time :float,
-        output_file: str) -> None:
+        output_file: str,
+        include_multiaxis_options: bool = False) -> None:
     """ Generate a funscript with minimal UI
 
     Args:
         video_file (str): path to video file
         start_time (float): start time in milliseconds
         end_time (float): end time in milliseconds (set -1.0 to use video end)
         output_file (str): path for the output file
+        include_multiaxis_options (bool): include options for multiaxis output
     """
     setup_logging()
     logging.info("Python Funscript Generator %s", VERSION)
     logging.info("Startup Path: %s", str(os.getcwd()))
     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)
-    _ = FunscriptGeneratorWindow(video_file, start_time, end_time, output_file)
+    _ = FunscriptGeneratorWindow(video_file, start_time, end_time, output_file, include_multiaxis_options)
     sys.exit(app.exec_())

funscript_editor/config/hyperparameter.yaml

@@ -29,10 +29,10 @@ signal:
   high_second_derivative_points_threshold: 1.2
 
   # filter length to detect a direction change
-  direction_change_filter_len: 4
+  direction_change_filter_len: 3
 
   # threshold value in milliseconds to merge additional points
-  additional_points_merge_time_threshold_in_ms: 100
+  additional_points_merge_time_threshold_in_ms: 70
 
   # threshold value to merge additional points
   additional_points_merge_distance_threshold: 10.0

funscript_editor/ui/funscript_generator_window.py

@@ -1,6 +1,7 @@
 import sys
 import funscript_editor.utils.logging as logging
 import os
+import json
 import time
 import platform
 import cv2
@@ -15,7 +16,7 @@
 
 from PyQt5 import QtCore, QtGui, QtWidgets
 
-USE_OPTICALFLOW = False
+USE_OPTICALFLOW = False # Enable some hardcoded optical flow testcode
 if USE_OPTICALFLOW:
     from funscript_editor.algorithms.opticalflow import OpticalFlowFunscriptGeneratorThread, OpticalFlowFunscriptGeneratorParameter
 
@@ -44,6 +45,9 @@ def __init__(self,
         if os.path.exists(definitions.ICON_PATH):
             self.setWindowIcon(QtGui.QIcon(definitions.ICON_PATH))
 
+        if include_multiaxis:
+            self.__logger.info("Enable multiaxis output")
+
         if not isinstance(output_file, Funscript):
             output_file = os.path.abspath(output_file)
             if os.path.isdir(output_file):
@@ -108,15 +112,34 @@ def __funscript_generated(self, funscripts, msg, success) -> None:
         first_metric = [x for x in funscripts.keys()][0]
 
         if isinstance(self.output_file, Funscript):
+            if len(funscripts) > 1:
+                self.__logger.warning("Multiaxis output for build-in UI is not implemented")
             for item in funscripts[first_metric].get_actions():
                 self.output_file.add_action(item['pos'], item['at'], SETTINGS['raw_output'])
             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 funscripts[first_metric].get_actions():
-                    f.write('{at};{pos}\n'.format(at=item['at'], pos=item['pos']))
+            if self.output_file.lower().endswith('.json'):
+                funscript_json_output = {
+                        'version': 1,
+                        'actions': {}
+                    }
+
+                for key in funscripts.keys():
+                    funscript_json_output['actions'][key] = []
+                    for item in funscripts[key].get_actions():
+                        funscript_json_output['actions'][key].append(item)
+
+                with open(self.output_file, 'w') as f:
+                    json.dump(funscript_json_output, f)
+            else:
+                # dump to CSV
+                if len(funscripts) > 1:
+                    self.__logger.warning("Multiaxis output for csv is  not implemented")
+                with open(self.output_file, 'w') as f:
+                    f.write('at;pos\n')
+                    for item in funscripts[first_metric].get_actions():
+                        f.write('{at};{pos}\n'.format(at=item['at'], pos=item['pos']))
 
             self.__logger.info("Save result to %s", self.output_file)
             if not success: self.__show_message(msg, error=True)