add ui to select a line

Author: arch

funscript_editor/algorithms/funscriptgenerator.py

@@ -270,6 +270,13 @@ def scale_score(self, status: str, metric : str = 'y') -> None:
             imgMin = FFmpegStream.get_projection(imgMin, self.projection_config)
             imgMax = FFmpegStream.get_projection(imgMax, self.projection_config)
 
+            if metric == 'roll':
+                max_distance_frame_num = np.argmax(np.array([abs(x) for x in self.score['distance']])) + self.params.start_frame
+                max_distance_frame = FFmpegStream.get_frame(self.params.video_path, max_distance_frame_num)
+                max_distance_frame = FFmpegStream.get_projection(max_distance_frame, self.projection_config)
+                center_line = self.ui.line_selector(max_distance_frame, "draw line on center of dick")
+                print('center line', center_line)
+
             min_tracking_points = self.get_tracking_points_by_frame_number(min_frame - self.params.start_frame)
             max_tracking_points = self.get_tracking_points_by_frame_number(max_frame - self.params.start_frame)
 
@@ -281,6 +288,7 @@ def scale_score(self, status: str, metric : str = 'y') -> None:
 
             # print('min_tracking_points', min_tracking_points, 'max_tracking_points', max_tracking_points)
 
+
             (desired_min, desired_max) = self.ui.min_max_selector(
                     image_min = imgMin,
                     image_max = imgMax,

funscript_editor/ui/opencvui.py

@@ -99,6 +99,37 @@ def was_key_pressed(self, key: str) -> bool:
         return False
 
 
+class DrawSingleLineWidget(object):
+    def __init__(self, background_img, window_name, preview_scaling):
+        self.original_image = background_img
+        self.clone = self.original_image.copy()
+        self.window_name = window_name
+        self.preview_scaling = preview_scaling
+
+        cv2.namedWindow(window_name)
+        cv2.setMouseCallback(window_name, self.extract_coordinates)
+
+        self.start_coordinate = None
+        self.end_coordinate = None
+
+    def extract_coordinates(self, event, x, y, flags, parameters):
+        x = round(x/self.preview_scaling)
+        y = round(y/self.preview_scaling)
+
+        if event == cv2.EVENT_LBUTTONDOWN:
+            self.start_coordinate = (x,y)
+
+        elif event == cv2.EVENT_LBUTTONUP:
+            self.end_coordinate = (x,y)
+            self.clone = self.original_image.copy()
+            cv2.line(self.clone, self.start_coordinate, self.end_coordinate, (36,255,12), 2)
+
+    def show_image(self):
+        return self.clone
+
+    def get_result(self):
+        return [ self.start_coordinate, self.end_coordinate ]
+
 
 class OpenCV_GUI(KeypressHandler):
     """ High Level OpenCV GUI.
@@ -473,6 +504,33 @@ def play_beep():
                 self.logger.warning("Notification sound file not found (%s)", self.params.notification_sound_file)
 
 
+    def line_selector(self,
+            image: np.ndarray,
+            txt: str):
+        """ Line Selector Widget
+
+        Args:
+            image (np.ndarray): background image
+            txt (str): text to display
+
+        Returns:
+            list: start and endpoint of line
+        """
+
+        line_widget = DrawSingleLineWidget(image, self.window_name, self.monitor_preview_scaling)
+        while True:
+            self.set_background_image(line_widget.show_image(), copy_image=False)
+            self.print_text(txt)
+            self.print_text("Press 'space' to continue")
+            ret = self.show(5)
+
+            if self.was_any_accept_key_pressed() or any(ret == x for x in [ord(' '), 13]):
+                result = line_widget.get_result()
+                if result[0] is not None and result[1] is not None:
+                    if abs(result[0][0] - result[1][0]) + abs(result[0][1] - result[1][1]) > 10:
+                        return result
+
+
     def min_max_selector(self,
             image_min: np.ndarray,
             image_max: np.ndarray,