always use ffmpeg to extract frames

Author: arch

funscript_editor/algorithms/funscriptgenerator.py

@@ -254,12 +254,17 @@ 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
 
-        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()
+        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)
 
         if success_min and success_max:
             if self.is_vr_video():
@@ -534,6 +539,7 @@ def tracking(self) -> str:
         delete_last_predictions = 0
         bbox_woman = [None for _ in range(self.params.number_of_trackers)]
         bbox_men = [None for _ in range(self.params.number_of_trackers)]
+        self.ui.clear_keypress_queue()
         try:
             while video.isOpen():
                 cycle_start = time.time()

funscript_editor/algorithms/videotracker.py

@@ -1,8 +1,8 @@
 """ Video Tracker """
 
-import os
 import cv2
 import time
+import copy
 import funscript_editor.utils.logging as logging
 import platform
 
@@ -65,6 +65,8 @@ def __init__(self,
         self.tracking_counter = 0
         self.cluster_center = [0, 0]
         self.plausibility_thresholds = [0, 0]
+        self.current_detected_tracking_box = tracking_bbox
+        self.lats_detected_tracking_box = tracking_bbox
         self.last_valid_tracking_box = tracking_bbox
         self.supervised_tracking_is_exit_condition = supervised_tracking_is_exit_condition
 
@@ -190,6 +192,19 @@ def __is_plausible(self, box) -> bool:
         return True
 
 
+    @staticmethod
+    def get_center(box: tuple) -> tuple:
+        """ Get the cencter point of an box
+
+        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 run(self) -> None:
         """ The Video Tracker Thread Function """
         self.__setup_tracker()
@@ -226,6 +241,8 @@ def run(self) -> None:
                 frame = self.queue_in.get()
                 frame_roi = frame[y0:y1, x0:x1]
                 success, bbox = self.tracker.update(frame_roi)
+                self.last_detected_tracking_box = self.current_detected_tracking_box
+                self.current_detected_tracking_box = copy.deepcopy(bbox)
                 self.tracking_counter += 1
                 status = StaticVideoTracker.Status.TRACKING_LOST
                 if not success or bbox is None:

funscript_editor/data/ffmpegstream.py

@@ -111,27 +111,6 @@ def get_video_info(
         return video_info
 
 
-    @staticmethod
-    def get_frame(
-            video_path :str,
-            frame_number :int) -> np.ndarray:
-        """ Get Video frame
-
-        Args:
-            video_path (str): path to video
-            frame_number (int): frame number to extract from video
-
-        Returns:
-            np.ndarray: opencv image
-        """
-        cap = cv2.VideoCapture(video_path)
-        if frame_number > 0:
-            cap.set(cv2.CAP_PROP_POS_FRAMES, frame_number)
-        _, frame = cap.read()
-        cap.release()
-        return frame
-
-
     @staticmethod
     def get_ffmpeg_command() -> str:
         """ Get FFmpeg binary command string
@@ -314,6 +293,65 @@ def more(self) -> bool:
         return self.frame_buffer.qsize() > 0
 
 
+    @staticmethod
+    def get_frame(
+            video_path :str,
+            frame_number :int) -> np.ndarray:
+        """ Get Video frame
+
+        Args:
+            video_path (str): path to video
+            frame_number (int): frame number to extract from video
+
+        Returns:
+            np.ndarray: opencv image
+        """
+        if False:
+            cap = cv2.VideoCapture(video_path)
+            if frame_number > 0:
+                cap.set(cv2.CAP_PROP_POS_FRAMES, frame_number)
+            _, frame = cap.read()
+            cap.release()
+            return frame
+        else:
+            video_info = FFmpegStream.get_video_info(video_path)
+            seek = FFmpegStream.frame_to_timestamp(frame_number, video_info.fps)
+
+            command = [
+                    FFmpegStream.get_ffmpeg_command(),
+                    '-hide_banner',
+                    '-loglevel', 'warning',
+                    '-ss', str(seek),
+                    '-hwaccel', 'auto',
+                    '-i', video_path,
+                    '-f', 'image2pipe',
+                    '-pix_fmt', 'bgr24',
+                    '-vsync', '0',
+                    '-vcodec', 'rawvideo',
+                    '-an',
+                    '-sn',
+                    '-'
+                ]
+
+            pipe = sp.Popen(
+                    command,
+                    stdout = sp.PIPE,
+                    stderr = sp.PIPE,
+                    bufsize= 3 * video_info.height * video_info.width
+                )
+
+            data = pipe.stdout.read(video_info.width * video_info.height * 3)
+            pipe.terminate()
+            try: pipe.stdout.close()
+            except: pass
+            try: pipe.stderr.close()
+            except: pass
+
+            return np.frombuffer(data, dtype='uint8').reshape(
+                    (video_info.height, video_info.width, 3)
+                )
+
+
     def run(self) -> None:
         """ Function to read transformed frames from ffmpeg video stream into a queue """