refactor distance to metric

Author: arch

funscript_editor/algorithms/funscriptgenerator.py

@@ -35,7 +35,7 @@ class FunscriptGeneratorParameter:
     # No default values
     video_path: str # no default value
     track_men: bool
-    direction: str
+    metric: str
     projection: str
 
     # Settings
@@ -95,7 +95,7 @@ def __init__(self,
         self.score = {
                 'x': [],
                 'y': [],
-                'd': []
+                'euclideanDistance': []
             }
         self.bboxes = {
                 'Men': [],
@@ -309,7 +309,7 @@ def calculate_score(self) -> None:
         if self.params.track_men:
             self.score['x'] = [w[0] - m[0] for w, m in zip(self.bboxes['Woman'], self.bboxes['Men'])]
             self.score['y'] = [m[1] - w[1] for w, m in zip(self.bboxes['Woman'], self.bboxes['Men'])]
-            self.score['d'] = [np.sqrt(np.sum((np.array(m[:2]) - np.array(w[:2])) ** 2, axis=0)) \
+            self.score['euclideanDistance'] = [np.sqrt(np.sum((np.array(m[:2]) - np.array(w[:2])) ** 2, axis=0)) \
                     for w, m in zip(self.bboxes['Woman'], self.bboxes['Men'])]
         else:
             self.score['x'] = [w[0] - min([x[0] for x in self.bboxes['Woman']]) for w in self.bboxes['Woman']]
@@ -318,29 +318,29 @@ def calculate_score(self) -> None:
 
         self.score['x'] = sp.scale_signal(self.score['x'], 0, 100)
         self.score['y'] = sp.scale_signal(self.score['y'], 0, 100)
-        self.score['d'] = sp.scale_signal(self.score['d'], 0, 100)
+        self.score['euclideanDistance'] = sp.scale_signal(self.score['euclideanDistance'], 0, 100)
 
 
-    def scale_score(self, status: str, direction : str = 'y') -> None:
+    def scale_score(self, status: str, metric : str = 'y') -> None:
         """ Scale the score to desired stroke high
 
         Note:
             We determine the lowerst and highes positions in the score and request the real position from user.
 
         Args:
             status (str): a status/info message to display in the window
-            direction (str): scale the 'y' or 'x' score
+            metric (str): scale the 'y' or 'x' score
         """
         if len(self.score['y']) < 2: return
 
         cap = cv2.VideoCapture(self.params.video_path)
         width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
         height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
 
-        if direction == 'd':
-            min_frame = np.argmin(np.array(self.score['d'])) + self.params.start_frame
-            max_frame = np.argmax(np.array(self.score['d'])) + self.params.start_frame
-        elif direction == 'x':
+        if metric == 'euclideanDistance':
+            min_frame = np.argmin(np.array(self.score['euclideanDistance'])) + self.params.start_frame
+            max_frame = np.argmax(np.array(self.score['euclideanDistance'])) + self.params.start_frame
+        elif metric == 'x':
             min_frame = np.argmin(np.array(self.score['x'])) + self.params.start_frame
             max_frame = np.argmax(np.array(self.score['x'])) + self.params.start_frame
         else:
@@ -370,16 +370,16 @@ def scale_score(self, status: str, direction : str = 'y') -> None:
             imgMin = cv2.resize(imgMin, None, fx=scale, fy=scale)
             imgMax = cv2.resize(imgMax, None, fx=scale, fy=scale)
 
-            if direction == 'y':
+            if metric == 'y':
                 title_min = "Bottom"
-            elif direction == 'x':
+            elif metric == 'x':
                 title_min = "Left"
             else:
                 title_min = "Minimum"
 
-            if direction == 'y':
+            if metric == 'y':
                 title_max = "Top"
-            elif direction == 'x':
+            elif metric == 'x':
                 title_max = "Right"
             else:
                 title_max = "Maximum"
@@ -396,9 +396,9 @@ def scale_score(self, status: str, direction : str = 'y') -> None:
             desired_min = 0
             desired_max = 99
 
-        if direction == 'd':
-            self.score['d'] = sp.scale_signal(self.score['d'], desired_min, desired_max)
-        elif direction == 'x':
+        if metric == 'euclideanDistance':
+            self.score['euclideanDistance'] = sp.scale_signal(self.score['euclideanDistance'], desired_min, desired_max)
+        elif metric == 'x':
             self.score['x'] = sp.scale_signal(self.score['x'], desired_min, desired_max)
         else:
             self.score['y'] = sp.scale_signal(self.score['y'], desired_min, desired_max)
@@ -796,16 +796,16 @@ def apply_shift(self, frame_number, position: str) -> int:
         Args:
             position (str): is max or min
         """
-        if self.params.direction == 'd':
+        if self.params.metric == 'euclideanDistance':
             shift_a = self.params.shift_top_points
-        elif self.params.direction == 'x':
+        elif self.params.metric == 'x':
             shift_a = self.params.shift_right_points
         else:
             shift_a = self.params.shift_top_points
 
-        if self.params.direction == 'd':
+        if self.params.metric == 'euclideanDistance':
             shift_b = self.params.shift_bottom_points
-        elif self.params.direction == 'x':
+        elif self.params.metric == 'x':
             shift_b = self.params.shift_left_points
         else:
             shift_b = self.params.shift_bottom_points
@@ -832,23 +832,23 @@ def get_score_with_offset(self, idx_dict) -> list:
         Returns:
             list: score with offset
         """
-        if self.params.direction == 'd':
+        if self.params.metric == 'euclideanDistance':
             offset_a = self.params.top_points_offset
-        elif self.params.direction == 'x':
+        elif self.params.metric == 'x':
             offset_a = self.params.right_points_offset
         else:
             offset_a = self.params.top_points_offset
 
-        if self.params.direction == 'd':
+        if self.params.metric == 'euclideanDistance':
             offset_b = self.params.bottom_points_offset
-        elif self.params.direction == 'x':
+        elif self.params.metric == 'x':
             offset_b = self.params.left_points_offset
         else:
             offset_b = self.params.bottom_points_offset
 
-        if self.params.direction == 'd':
-            score = copy.deepcopy(self.score['d'])
-        elif self.params.direction == 'x':
+        if self.params.metric == 'euclideanDistance':
+            score = copy.deepcopy(self.score['euclideanDistance'])
+        elif self.params.metric == 'x':
             score = copy.deepcopy(self.score['x'])
         else:
             score = copy.deepcopy(self.score['y'])
@@ -892,9 +892,9 @@ def determin_change_points(self) -> dict:
             dict: all local max and min points in score {'min':[idx1, idx2, ...], 'max':[idx1, idx2, ...]}
         """
         self.logger.info("Determine change points")
-        if self.params.direction == 'd':
-            idx_dict = sp.get_local_max_and_min_idx(self.score['d'], self.video_info.fps)
-        elif self.params.direction == 'x':
+        if self.params.metric == 'euclideanDistance':
+            idx_dict = sp.get_local_max_and_min_idx(self.score['euclideanDistance'], self.video_info.fps)
+        elif self.params.metric == 'x':
             idx_dict = sp.get_local_max_and_min_idx(self.score['x'], self.video_info.fps)
         else:
             idx_dict = sp.get_local_max_and_min_idx(self.score['y'], self.video_info.fps)
@@ -909,9 +909,9 @@ def create_funscript(self, idx_dict: dict) -> None:
                              {'min':[idx1, idx2, ...], 'max':[idx1, idx2, ...]}
         """
         if self.params.raw_output:
-            if self.params.direction == 'd':
-                output_score = copy.deepcopy(self.score['d'])
-            elif self.params.direction == 'x':
+            if self.params.metric == 'euclideanDistance':
+                output_score = copy.deepcopy(self.score['euclideanDistance'])
+            elif self.params.metric == 'x':
                 output_score = copy.deepcopy(self.score['x'])
             else:
                 output_score = copy.deepcopy(self.score['y'])
@@ -925,16 +925,16 @@ def create_funscript(self, idx_dict: dict) -> None:
         else:
             output_score = self.get_score_with_offset(idx_dict)
 
-            if self.params.direction == 'd':
+            if self.params.metric == 'euclideanDistance':
                 threshold_a = self.params.bottom_threshold
-            elif self.params.direction == 'x':
+            elif self.params.metric == 'x':
                 threshold_a = self.params.left_threshold
             else:
                 threshold_a = self.params.bottom_threshold
 
-            if self.params.direction == 'd':
+            if self.params.metric == 'euclideanDistance':
                 threshold_b = self.params.top_threshold
-            elif self.params.direction == 'x':
+            elif self.params.metric == 'x':
                 threshold_b = self.params.right_threshold
             else:
                 threshold_b = self.params.top_threshold
@@ -959,9 +959,8 @@ def create_funscript(self, idx_dict: dict) -> None:
     def run(self) -> None:
         """ The Funscript Generator Thread Function """
         try:
-            if self.params.direction == 'd' and not self.params.track_men:
-                self.finished("Tracking with 'd' and no men tracker is not implemented!", False)
-                return
+            if self.params.metric not in ['x', 'y']:
+                self.params.track_men = True # we need 2 tracking points
 
             if self.video_info.fps < 31.0 and self.params.skip_frames > 0:
                 self.logger.warning("The Video has less than 30 frames per seconds and you have set skip_frames to %d "\
@@ -980,7 +979,7 @@ def run(self) -> None:
 
                 if len(self.score['y']) >= HYPERPARAMETER['min_frames']:
                     self.logger.info("Scale score")
-                    self.scale_score(status, direction=self.params.direction)
+                    self.scale_score(status, metric=self.params.metric)
 
             if len(self.score['y']) < HYPERPARAMETER['min_frames']:
                 self.finished(status + ' -> Tracking time insufficient', False)

funscript_editor/ui/funscript_generator_window.py

@@ -108,7 +108,7 @@ def run(self) -> None:
                 FunscriptGeneratorParameter(
                     video_path = self.video_file,
                     track_men = self.settings['trackingMethod'] != 'Woman',
-                    direction = self.settings['trackingAxis'],
+                    metric = self.settings['trackingMetric'],
                     projection = self.settings['videoType'],
                     start_frame = self.start_frame,
                     end_frame = self.end_frame

funscript_editor/ui/settings_dialog.py

@@ -38,14 +38,23 @@ def __setup_ui_bindings(self):
                     list(filter(lambda x: PROJECTION[x]['name'] == value, PROJECTION.keys()))[0]
                 )
             )
-        self.ui.trackingAxisComboBox.currentTextChanged.connect(lambda value: self.__set_setting('trackingAxis', value))
+        self.ui.trackingMetricComboBox.currentTextChanged.connect(self.__set_tracking_metric)
         self.ui.trackingMethodComboBox.currentTextChanged.connect(lambda value: self.__set_setting('trackingMethod', value))
 
     def __setup_combo_boxes(self):
         self.ui.videoTypeComboBox.addItems([PROJECTION[key]['name'] for key in PROJECTION.keys()])
-        self.ui.trackingAxisComboBox.addItems(['y', 'x'])
+        self.ui.trackingMetricComboBox.addItems(['y', 'x', 'euclideanDistance'])
         self.ui.trackingMethodComboBox.addItems(['Woman', 'Woman + Men'])
 
+    def __set_tracking_metric(self, value):
+        if value in ['x', 'y']:
+            self.ui.trackingMethodComboBox.setEnabled(True)
+        else:
+            self.ui.trackingMethodComboBox.setCurrentText('Woman + Men')
+            self.ui.trackingMethodComboBox.setEnabled(False)
+
+        self.__set_setting('trackingMetric', value)
+
     def __apply(self):
         self.form.hide()
         self.applySettings.emit()

funscript_editor/ui/settings_view.py

@@ -41,14 +41,14 @@ def setupUi(self, Form):
         self.label_3 = QtWidgets.QLabel(Form)
         self.label_3.setObjectName("label_3")
         self.formLayout.setWidget(3, QtWidgets.QFormLayout.LabelRole, self.label_3)
-        self.trackingAxisComboBox = QtWidgets.QComboBox(Form)
+        self.trackingMetricComboBox = QtWidgets.QComboBox(Form)
         sizePolicy = QtWidgets.QSizePolicy(QtWidgets.QSizePolicy.MinimumExpanding, QtWidgets.QSizePolicy.Fixed)
         sizePolicy.setHorizontalStretch(0)
         sizePolicy.setVerticalStretch(0)
-        sizePolicy.setHeightForWidth(self.trackingAxisComboBox.sizePolicy().hasHeightForWidth())
-        self.trackingAxisComboBox.setSizePolicy(sizePolicy)
-        self.trackingAxisComboBox.setObjectName("trackingAxisComboBox")
-        self.formLayout.setWidget(2, QtWidgets.QFormLayout.FieldRole, self.trackingAxisComboBox)
+        sizePolicy.setHeightForWidth(self.trackingMetricComboBox.sizePolicy().hasHeightForWidth())
+        self.trackingMetricComboBox.setSizePolicy(sizePolicy)
+        self.trackingMetricComboBox.setObjectName("trackingMetricComboBox")
+        self.formLayout.setWidget(2, QtWidgets.QFormLayout.FieldRole, self.trackingMetricComboBox)
         self.trackingMethodComboBox = QtWidgets.QComboBox(Form)
         sizePolicy = QtWidgets.QSizePolicy(QtWidgets.QSizePolicy.MinimumExpanding, QtWidgets.QSizePolicy.Fixed)
         sizePolicy.setHorizontalStretch(0)
@@ -83,7 +83,7 @@ def retranslateUi(self, Form):
         _translate = QtCore.QCoreApplication.translate
         Form.setWindowTitle(_translate("Form", "Form"))
         self.label.setText(_translate("Form", "Video Type:"))
-        self.label_2.setText(_translate("Form", "Tracking Axis:"))
+        self.label_2.setText(_translate("Form", "Tracking Metric:"))
         self.label_3.setText(_translate("Form", "Tracking Method:"))
         self.okButton.setText(_translate("Form", "OK"))
         self.label_4.setText(_translate("Form", "  # Settings"))

funscript_editor/ui/settings_view.ui

@@ -64,7 +64,7 @@
      <item row="2" column="0">
       <widget class="QLabel" name="label_2">
        <property name="text">
-        <string>Tracking Axis:</string>
+        <string>Tracking Metric:</string>
        </property>
       </widget>
      </item>
@@ -76,7 +76,7 @@
       </widget>
      </item>
      <item row="2" column="1">
-      <widget class="QComboBox" name="trackingAxisComboBox">
+      <widget class="QComboBox" name="trackingMetricComboBox">
        <property name="sizePolicy">
         <sizepolicy hsizetype="MinimumExpanding" vsizetype="Fixed">
          <horstretch>0</horstretch>