draw center tracking box point

Author: arch

funscript_editor/algorithms/funscriptgenerator.py

@@ -407,20 +407,21 @@ def init_trackers(self, ffmpeg_stream: FFmpegStream) -> tuple:
         for tracker_number in range(self.params.number_of_trackers):
             bbox_woman = self.ui.bbox_selector(
                     preview_frame,
-                    "Select {} Feature #{}".format(self.get_target_name(0), tracker_number+1)
+                    "Select {} Feature #{}".format(self.get_target_name(0), tracker_number+1),
+                    add_center = True
                 )
 
             preview_frame = self.ui.draw_box_to_image(
-                    preview_frame, 
-                    bbox_woman, 
+                    preview_frame,
+                    bbox_woman,
                     color=(255,0,255)
                 )
 
             if self.params.supervised_tracking:
                 while True:
                     tracking_areas_woman[tracker_number] = self.ui.bbox_selector(
                             preview_frame,
-                            "Select the Supervised Tracking Area for the {} Feature #{}".format(self.get_target_name(0), tracker_number+1)
+                            "Select the Supervised Tracking Area for the {} Feature #{}".format(self.get_target_name(0), tracker_number+1),
                         )
 
                     if StaticVideoTracker.is_bbox_in_tracking_area(bbox_woman, tracking_areas_woman[tracker_number]):
@@ -456,7 +457,8 @@ def init_trackers(self, ffmpeg_stream: FFmpegStream) -> tuple:
             if self.params.track_men:
                 bbox_men = self.ui.bbox_selector(
                         preview_frame,
-                        "Select {} Feature #{}".format(self.get_target_name(1), tracker_number+1)
+                        "Select {} Feature #{}".format(self.get_target_name(1), tracker_number+1),
+                        add_center = True
                     )
                 preview_frame = self.ui.draw_box_to_image(preview_frame, bbox_men, color=(255,0,255))
                 if self.params.supervised_tracking:

funscript_editor/algorithms/videotracker.py

@@ -205,6 +205,26 @@ def get_center(box: tuple) -> tuple:
         return ( round(box[0] + box[2]/2), round(box[1] + box[3]/2) )
 
 
+    @staticmethod
+    def create_tracking_box(box: tuple):
+        """ Create an tracking box (x, y, h, w, x_center, y_center)
+
+        Args:
+            box (tuple): the predicted bounding box
+
+        Returns:
+            tuble: tracking box
+        """
+        if box is None:
+            return None
+
+        if len(box) < 4:
+            return None
+
+        center = StaticVideoTracker.get_center(box)
+        return (box[0], box[1], box[2], box[3], center[0], center[1])
+
+
     def run(self) -> None:
         """ The Video Tracker Thread Function """
         self.__setup_tracker()
@@ -241,6 +261,7 @@ def run(self) -> None:
                 frame = self.queue_in.get()
                 frame_roi = frame[y0:y1, x0:x1]
                 success, bbox = self.tracker.update(frame_roi)
+                bbox = self.create_tracking_box(bbox)
                 self.last_detected_tracking_box = self.current_detected_tracking_box
                 self.current_detected_tracking_box = copy.deepcopy(bbox)
                 self.tracking_counter += 1
@@ -249,7 +270,7 @@ def run(self) -> None:
                     bbox = None
                 else:
                     status = StaticVideoTracker.Status.OK
-                    bbox = (int(bbox[0] + x0), int(bbox[1] + y0), int(bbox[2]), int(bbox[3]))
+                    bbox = (int(bbox[0] + x0), int(bbox[1] + y0), int(bbox[2]), int(bbox[3]), int(bbox[4] + x0), int(bbox[5] + y0))
                     if self.params.tracking_plausibility_check:
                         if not self.__is_plausible(bbox):
                             status = StaticVideoTracker.Status.IMPLAUSIBLE

funscript_editor/ui/opencvui.py

@@ -200,7 +200,7 @@ def draw_box(self, bbox, color: tuple = (255, 0, 255)) -> None:
             bbox = [bbox]
 
         for box in bbox:
-            if box and len(box) == 4:
+            if box and len(box) >= 4:
                 cv2.rectangle(
                         self.preview_image,
                         (box[0], box[1]),
@@ -209,6 +209,15 @@ def draw_box(self, bbox, color: tuple = (255, 0, 255)) -> None:
                         3,
                         1
                     )
+                if len(box) >= 6:
+                    cv2.circle(
+                            self.preview_image,
+                            (box[4], box[5]),
+                            5,
+                            color,
+                            2
+                        )
+
 
     @staticmethod
     def draw_box_to_image(image: np.ndarray, bbox, color: tuple = (255, 0, 255)) -> np.ndarray:
@@ -226,7 +235,7 @@ def draw_box_to_image(image: np.ndarray, bbox, color: tuple = (255, 0, 255)) ->
             bbox = [bbox]
 
         for box in bbox:
-            if box and len(box) == 4:
+            if box and len(box) >= 4:
                 cv2.rectangle(
                         image,
                         (box[0], box[1]),
@@ -235,6 +244,14 @@ def draw_box_to_image(image: np.ndarray, bbox, color: tuple = (255, 0, 255)) ->
                         3,
                         1
                     )
+                if len(box) >= 6:
+                    cv2.circle(
+                            image,
+                            (box[4], box[5]),
+                            5,
+                            color,
+                            2
+                        )
 
         return image
 
@@ -473,13 +490,26 @@ def min_max_selector(self,
                 else (trackbarValueMax, trackbarValueMin)
 
 
+    @staticmethod
+    def get_center(box: tuple) -> tuple:
+        """ Get the cencter point of an box
 
-    def bbox_selector(self, image: np.ndarray, txt: str) -> tuple:
+        Args:
+            box (tuple): the predicted bounding box
+
+        Returns:
+            tuple (x,y) of the current point
+        """
+        return ( round(box[0] + box[2]/2), round(box[1] + box[3]/2) )
+
+
+    def bbox_selector(self, image: np.ndarray, txt: str, add_center: bool = False) -> tuple:
         """ Window to get an bounding box from user input
 
         Args:
             image (np.ndarray): opencv image e.g. the first frame to determine the bounding box
             txt (str): additional text to display on the selection window
+            add_center (bool): add center cordinates to the box
 
         Returns:
             tuple: user input bounding box tuple (x,y,w,h)
@@ -542,6 +572,10 @@ def bbox_selector(self, image: np.ndarray, txt: str) -> tuple:
                     round(bbox[3]/self.params.zoom_factor)
                 )
 
+        if add_center:
+            center = self.get_center(bbox)
+            bbox = (bbox[0], bbox[1], bbox[2], bbox[3], center[0], center[1])
+
         self.logger.info("User Input: %s", str(bbox))
         return bbox