improve supervised tracking option

Author: arch

funscript_editor/algorithms/funscriptgenerator.py

@@ -35,9 +35,9 @@ class FunscriptGeneratorParameter:
     """ Funscript Generator Parameter Dataclass with default values """
     video_path: str
     track_men: bool
+    supervised_tracking: bool
     metric: str
     projection: str
-    supervised_tracking: bool = True
 
     # Settings
     start_frame: int = 0 # default is video start (input: set current video position)
@@ -393,6 +393,8 @@ def scale_score(self, status: str, metric : str = 'y') -> None:
                 elif 'ou' in self.params.projection.split('_'):
                     imgMin = imgMin[:int(imgMin.shape[0]/2), :]
                     imgMax = imgMax[:int(imgMax.shape[0]/2), :]
+                else:
+                    self.logger.warning("Unknown VR Projection Type: %s", self.params.projection)
 
             if PROJECTION[self.params.projection]['parameter']['width'] > 0:
                 scale = PROJECTION[self.params.projection]['parameter']['width'] / float(1.75*imgMax.shape[1])
@@ -671,7 +673,10 @@ def tracking(self) -> str:
         bbox_woman = self.get_bbox(first_frame, "Select Woman Feature")
         preview_frame = self.draw_box(first_frame, bbox_woman, color=(255,0,255))
         if self.params.supervised_tracking:
-            tracking_area_woman = self.get_bbox(preview_frame, "Select the Supervised Tracking Area for the Woman Feature")
+            while True:
+                tracking_area_woman = self.get_bbox(preview_frame, "Select the Supervised Tracking Area for the Woman Feature")
+                if StaticVideoTracker.is_bbox_in_tracking_area(bbox_woman, tracking_area_woman): break
+                self.logger.error("Invalid supervised tracking area selected")
             preview_frame = self.draw_box(preview_frame, tracking_area_woman, color=(0,255,0))
             tracker_woman = StaticVideoTracker(first_frame, bbox_woman, supervised_tracking_area = tracking_area_woman)
         else:
@@ -682,7 +687,10 @@ def tracking(self) -> str:
             bbox_men = self.get_bbox(preview_frame, "Select Men Feature")
             preview_frame = self.draw_box(preview_frame, bbox_men, color=(255,0,255))
             if self.params.supervised_tracking:
-                tracking_area_men = self.get_bbox(preview_frame, "Select the Supervised Tracking Area for the Men Feature")
+                while True:
+                    tracking_area_men = self.get_bbox(preview_frame, "Select the Supervised Tracking Area for the Men Feature")
+                    if StaticVideoTracker.is_bbox_in_tracking_area(bbox_men, tracking_area_men): break
+                    self.logger.error("Invalid supervised tracking area selected")
                 tracker_men = StaticVideoTracker(first_frame, bbox_men, supervised_tracking_area = tracking_area_men)
             else:
                 tracker_men = StaticVideoTracker(first_frame, bbox_men)
@@ -720,15 +728,15 @@ def tracking(self) -> str:
                 # Process data from last step while the next tracking points get predicted.
                 # This should improve the whole processing speed, because the tracker run in a seperate thread
                 bboxes['Woman'][frame_num-1] = bbox_woman
-                last_frame = self.draw_box(last_frame, bboxes['Woman'][frame_num-1])
+                last_frame = self.draw_box(last_frame, bboxes['Woman'][frame_num-1], color=(0,255,0))
                 if self.params.supervised_tracking:
                     last_frame = self.draw_box(last_frame, tracking_area_woman, color=(0,255,0))
 
                 if self.params.track_men:
                     bboxes['Men'][frame_num-1] = bbox_men
-                    last_frame = self.draw_box(last_frame, bboxes['Men'][frame_num-1])
+                    last_frame = self.draw_box(last_frame, bboxes['Men'][frame_num-1], color=(255,0,255))
                     if self.params.supervised_tracking:
-                        last_frame = self.draw_box(last_frame, tracking_area_men, color=(0,255,0))
+                        last_frame = self.draw_box(last_frame, tracking_area_men, color=(255,0,255))
 
                 last_frame = self.draw_fps(last_frame)
                 cv2.putText(last_frame, "Press 'q' if the tracking point shifts or a video cut occured",

funscript_editor/algorithms/videotracker.py

@@ -21,6 +21,7 @@ class StaticVideoTracker:
         first_frame (np.ndarray): open cv image representing the start frame
         tracking_bbox (tuple): tuple with (x,y,w,h) of the init tracking box
         limit_searchspace (dict) : only insert the specified region around the init box
+        supervised_tracking_area (tuple, optional): tuple with (x,y,w,h) of the supervised tracking area
         queue_size (int): in (work) and out (result) queue size
     """
 
@@ -46,6 +47,37 @@ def __init__(self,
     __logger = logging.getLogger(__name__)
 
 
+    @staticmethod
+    def is_bbox_in_tracking_area(bbox: tuple, supervised_tracking_area: tuple = None) -> bool:
+        """ Check if tracking box is inside the supervised tracking area
+
+        Args:
+            bbox (tuple): tuple with (x,y,w,h) of the tracking box
+            supervised_tracking_area (tuple, optional): tuple with (x,y,w,h) of the supervised tracking area
+
+        Returns:
+            bool: True if tracking box is in supervised_tracking_area else False
+        """
+        if bbox is None:
+            return False
+        if supervised_tracking_area is None:
+            return True
+        if supervised_tracking_area[2] <= 1:
+            return False
+        if supervised_tracking_area[3] <= 1:
+            return False
+        if bbox[0] < supervised_tracking_area[0]:
+            return False
+        if bbox[1] < supervised_tracking_area[1]:
+            return False
+        if bbox[0] + bbox[2] > supervised_tracking_area[0] + supervised_tracking_area[2]:
+            return False
+        if bbox[1] + bbox[3] > supervised_tracking_area[1] + supervised_tracking_area[3]:
+            return False
+
+        return True
+
+
     def stop(self) -> None:
         """ Stop the tracker thread """
         self.stopped = True
@@ -117,15 +149,8 @@ def run(self) -> None:
                 if success:
                     status = "OK"
                     bbox = (int(bbox[0] + x0), int(bbox[1] + y0), int(bbox[2]), int(bbox[3]))
-                    if self.supervised_tracking_area is not None:
-                        if bbox[0] < self.supervised_tracking_area[0]:
-                            status = "Feature outside the specified area"
-                        elif bbox[1] < self.supervised_tracking_area[1]:
-                            status = "Feature outside the specified area"
-                        elif bbox[0] + bbox[2] > self.supervised_tracking_area[0] + self.supervised_tracking_area[2]:
-                            status = "Feature outside the specified area"
-                        elif bbox[1] + bbox[3] > self.supervised_tracking_area[1] + self.supervised_tracking_area[3]:
-                            status = "Feature outside the specified area"
+                    if not StaticVideoTracker.is_bbox_in_tracking_area(bbox, self.supervised_tracking_area):
+                        status = "Feature outside the specified area"
 
                 self.queue_out.put((status, bbox))
 

funscript_editor/ui/funscript_generator_window.py

@@ -111,7 +111,8 @@ def run(self) -> None:
         self.funscript_generator = FunscriptGeneratorThread(
                 FunscriptGeneratorParameter(
                     video_path = self.video_file,
-                    track_men = self.settings['trackingMethod'] != 'Woman',
+                    track_men = 'Men' in self.settings['trackingMethod'],
+                    supervised_tracking = 'Supervised' in self.settings['trackingMethod'],
                     metric = self.settings['trackingMetric'],
                     projection = self.settings['videoType'],
                     start_frame = self.start_frame,

funscript_editor/ui/settings_dialog.py

@@ -28,10 +28,12 @@ def __init__(self, settings: dict, include_vr: bool = True):
     #: apply settings event
     applySettings = QtCore.pyqtSignal()
 
+
     def show(self):
-        """ Show wttings dialog """
+        """ Show settings dialog """
         self.form.show()
 
+
     def __setup_ui_bindings(self):
         self.ui.okButton.clicked.connect(self.__apply)
         self.ui.videoTypeComboBox.currentTextChanged.connect(
@@ -43,27 +45,34 @@ def __setup_ui_bindings(self):
         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() \
                 if 'vr' not in key.lower() or self.include_vr])
+        self.ui.trackingMethodComboBox.addItems(['Unsupervised Woman', 'Unsupervised Woman + Men', 'Supervised Woman', 'Supervised Woman + Men']) # set before tracking metric
         self.ui.trackingMetricComboBox.addItems(['y (up-down)', 'x (left-right)', 'euclideanDistance', 'roll (rotation)'])
-        # self.ui.trackingMethodComboBox.addItems(['Unsupervised Woman', 'Unsupervised Woman + Men', 'Supervised Woman', 'Supervised Woman + Men'])
-        self.ui.trackingMethodComboBox.addItems(['Woman', 'Woman + Men'])
+
 
     def __set_tracking_metric(self, value):
         value = value.split('(')[0].strip()
+        current_tracking_method_items = [self.ui.trackingMethodComboBox.itemText(i) for i in range(self.ui.trackingMethodComboBox.count())]
+
         if value in ['x', 'y']:
-            self.ui.trackingMethodComboBox.setEnabled(True)
+            if 'Unsupervised Woman' not in current_tracking_method_items:
+                self.ui.trackingMethodComboBox.addItems(['Unsupervised Woman', 'Supervised Woman'])
         else:
-            self.ui.trackingMethodComboBox.setCurrentText('Woman + Men')
-            self.ui.trackingMethodComboBox.setEnabled(False)
+            if 'Unsupervised Woman' in current_tracking_method_items:
+                self.ui.trackingMethodComboBox.clear()
+                self.ui.trackingMethodComboBox.addItems(['Unsupervised Woman + Men', 'Supervised Woman + Men'])
 
         self.__set_setting('trackingMetric', value)
 
+
     def __apply(self):
         self.form.hide()
         self.applySettings.emit()
 
+
     def __set_setting(self, key, value):
         self.settings[key] = value