store trcked points in class

Author: arch

funscript_editor/algorithms/funscriptgenerator.py

@@ -98,6 +98,7 @@ def __init__(self,
         self.params = params
         self.funscripts = funscripts
         self.video_info = FFmpegStream.get_video_info(self.params.video_path)
+        self.tracking_points = {}
         self.score = {
                 'x': [],
                 'y': [],
@@ -252,17 +253,9 @@ def scale_score(self, status: str, metric : str = 'y') -> None:
         min_frame = np.argmin(np.array(self.score[metric])) + self.params.start_frame
         max_frame = np.argmax(np.array(self.score[metric])) + self.params.start_frame
 
-        if False:
-            cap = cv2.VideoCapture(self.params.video_path)
-            cap.set(cv2.CAP_PROP_POS_FRAMES, min_frame)
-            success_min, imgMin = cap.read()
-            cap.set(cv2.CAP_PROP_POS_FRAMES, max_frame)
-            success_max, imgMax = cap.read()
-            cap.release()
-        else:
-            success_min, success_max = True, True
-            imgMin = FFmpegStream.get_frame(self.params.video_path, min_frame)
-            imgMax = FFmpegStream.get_frame(self.params.video_path, max_frame)
+        success_min, success_max = True, True
+        imgMin = FFmpegStream.get_frame(self.params.video_path, min_frame)
+        imgMax = FFmpegStream.get_frame(self.params.video_path, max_frame)
 
         if success_min and success_max:
             if self.is_vr_video():
@@ -283,6 +276,20 @@ def scale_score(self, status: str, metric : str = 'y') -> None:
             imgMin = cv2.resize(imgMin, None, fx=scale, fy=scale)
             imgMax = cv2.resize(imgMax, None, fx=scale, fy=scale)
 
+            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)
+
+            # TODO: draw points to image
+            # NOTE: Code below does not work because image do not use the same projection
+
+            # for points in min_tracking_points:
+            #     imgMin = OpenCV_GUI.draw_point_to_image(imgMin, points)
+
+            # for points in max_tracking_points:
+            #     imgMax = OpenCV_GUI.draw_point_to_image(imgMax, points)
+
+            # 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,
@@ -654,10 +661,53 @@ def tracking(self) -> str:
         self.logger.info(status)
         self.logger.info('Interpolate tracking boxes')
         interpolated_bboxes = self.interpolate_bboxes(bboxes)
+        self.tracking_points = self.determine_tracking_points(interpolated_bboxes)
         self.calculate_score(interpolated_bboxes)
         return status
 
 
+    def determine_tracking_points(self, interpolated_bboxes: dict) -> dict:
+        """ Determine the final tracking points
+
+        Args:
+            interpolate_bboxes (dict): interpolate bboxes from all trackers
+
+        Returns:
+            dict: final tracking points
+        """
+        result = {}
+        for tracker_type in interpolated_bboxes.keys():
+            result[tracker_type] = {}
+            for tracker_number in interpolated_bboxes[tracker_type].keys():
+                result[tracker_type][tracker_number] = [self.get_center(item) for item in interpolated_bboxes[tracker_type][tracker_number]]
+
+        return result
+
+
+    def get_tracking_points_by_frame_number(self, relative_frame_number: int) -> list:
+        """ Get tracking points by frame number
+
+        Args:
+            relative_frame_number (int): relative frame number
+
+        Returns:
+            list: all tracking points
+        """
+        result = []
+        for tracker_type in self.tracking_points.keys():
+            for idx, tracker_number in enumerate([k for k in self.tracking_points[tracker_type].keys()]):
+                if len(result) < idx + 1:
+                    if relative_frame_number < len(self.tracking_points[tracker_type][tracker_number]):
+                        result.append([self.tracking_points[tracker_type][tracker_number][relative_frame_number]])
+                    else:
+                        result.append([])
+                else:
+                    if relative_frame_number < len(self.tracking_points[tracker_type][tracker_number]):
+                        result[idx].append(self.tracking_points[tracker_type][tracker_number][relative_frame_number])
+
+        return result
+
+
     def finished(self, status: str, success :bool) -> None:
         """ Process necessary steps to complete the predicted funscript
 

funscript_editor/config/settings.yaml

@@ -20,10 +20,6 @@ tracker: 'CSRT'
 # Specify the wav file in the config directory to play when tracking finished (write 'off' to disable the sound notification)
 notification_sound: 'off'
 
-# Time in milliseconds at which the tracking is stopped if the selected feature is not found
-# THIS feature has no effekt since v0.1.2!
-tracking_lost_time: 0
-
 # Specify the scene detector algorithm. Tracking pause at detected scene changes. Available options:
 # - 'CSV': load scene csv file from [PySceneDetect](https://github.com/Breakthrough/PySceneDetect) if available
 # - 'CONTENT': Detects cuts using changes in colour and intensity between frames.

funscript_editor/ui/opencvui.py

@@ -228,7 +228,7 @@ def draw_box_to_image(image: np.ndarray, bbox, color: tuple = (255, 0, 255)) ->
 
         Args:
             image (np.ndarray): image
-            bbox (tuple): tracking box with (x,y,w,h)
+            bbox (tuple or list): tracking box with (x,y,w,h) or list of tracking boxes
             color (tuple): RGB color values for the box
 
         Returns:
@@ -259,6 +259,33 @@ def draw_box_to_image(image: np.ndarray, bbox, color: tuple = (255, 0, 255)) ->
         return image
 
 
+    @staticmethod
+    def draw_point_to_image(image: np.ndarray, point, color: tuple = (255, 0, 255)) -> np.ndarray:
+        """ Draw an point to given image
+
+        Args:
+            image (np.ndarray): image
+            point (tuple or list): points (x,y) or list of points
+            color (tuple): RGB color values for the box
+
+        Returns:
+            np.ndarray: opencv image
+        """
+        if not isinstance(point, list):
+            point = [point]
+
+        for p in point:
+            cv2.circle(
+                    image,
+                    (p[0], p[1]),
+                    4,
+                    color,
+                    2
+                )
+
+        return image
+
+
     def get_preview_fps(self) -> float:
         """ Get current processing FPS