add testcode for pitch

Author: arch

funscript_editor/algorithms/funscriptgenerator.py

@@ -1,13 +1,12 @@
 """ Top level process to generate the funscript actions by tracking selected features in the video """
 
 import cv2
-import json
 import copy
 import time
+import math
 import logging
 
 from screeninfo import get_monitors
-from threading import Thread
 from queue import Queue
 from pynput.keyboard import Key, Listener
 from dataclasses import dataclass
@@ -18,7 +17,7 @@
 
 from funscript_editor.algorithms.kalmanfilter import KalmanFilter2D
 from funscript_editor.algorithms.videotracker import StaticVideoTracker
-from funscript_editor.data.ffmpegstream import FFmpegStream, VideoInfo
+from funscript_editor.data.ffmpegstream import FFmpegStream
 from funscript_editor.data.funscript import Funscript
 from funscript_editor.utils.config import HYPERPARAMETER, SETTINGS, PROJECTION
 from funscript_editor.utils.logging import get_logfiles_paths
@@ -95,7 +94,8 @@ def __init__(self,
         self.score = {
                 'x': [],
                 'y': [],
-                'euclideanDistance': []
+                'euclideanDistance': [],
+                'pitch': []
             }
         self.bboxes = {
                 'Men': [],
@@ -263,8 +263,8 @@ def min_max_selector(self,
         Returns:
             tuple: with selected (min: flaot, max float)
         """
-        cv2.createTrackbar("Min", self.window_name, lower_limit, upper_limit, lambda x: None)
-        cv2.createTrackbar("Max", self.window_name, upper_limit, upper_limit, lambda x: None)
+        cv2.createTrackbar("Min", self.window_name, lower_limit, upper_limit, lambda _: None)
+        cv2.createTrackbar("Max", self.window_name, upper_limit, upper_limit, lambda _: None)
         image = np.concatenate((image_min, image_max), axis=1)
 
         if info != "":
@@ -317,13 +317,33 @@ def calculate_score(self) -> None:
             self.score['euclideanDistance'] = [np.sqrt(np.sum((np.array(m) - np.array(w)) ** 2, axis=0)) \
                     for w, m in zip(woman_center, men_center)]
 
+            for i in range( min(( len(men_center), len(woman_center) )) ):
+                x = self.bboxes['Woman'][i][0] - self.bboxes['Men'][i][0]
+                y = self.bboxes['Men'][i][1] - self.bboxes['Woman'][i][1]
+                if x >= 0 and y >= 0:
+                    self.score['pitch'].append(np.arctan(np.array(y / max((10e-3, x)))))
+                elif x >= 0 and y < 0:
+                    self.score['pitch'].append(-1.0*np.arctan(np.array(y / max((10e-3, x)))))
+                elif x < 0 and y < 0:
+                    self.score['pitch'].append(math.pi + -1.0*np.arctan(np.array(y / x)))
+                elif x < 0 and y >= 0:
+                    self.score['pitch'].append(math.pi + np.arctan(np.array(y / x)))
+                else:
+                    # this should never happen
+                    self.logger.error('Calculate score not implement for x=%d, y=%d', x, y)
+
+            # invert
+            self.score['pitch'] = [-1.0*item for item in self.score['pitch']]
+
+
         else:
             self.score['x'] = [w[0] - min([x[0] for x in self.bboxes['Woman']]) for w in self.bboxes['Woman']]
             self.score['y'] = [max([x[1] for x in self.bboxes['Woman']]) - w[1] for w in self.bboxes['Woman']]
 
         self.score['x'] = sp.scale_signal(self.score['x'], 0, 100)
         self.score['y'] = sp.scale_signal(self.score['y'], 0, 100)
         self.score['euclideanDistance'] = sp.scale_signal(self.score['euclideanDistance'], 0, 100)
+        self.score['pitch'] = sp.scale_signal(self.score['pitch'], 0, 100)
 
 
     def scale_score(self, status: str, metric : str = 'y') -> None:
@@ -339,8 +359,6 @@ def scale_score(self, status: str, metric : str = 'y') -> None:
         if len(self.score['y']) < 2: return
 
         cap = cv2.VideoCapture(self.params.video_path)
-        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
-        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
 
         if metric == 'euclideanDistance':
             min_frame = np.argmin(np.array(self.score['euclideanDistance'])) + self.params.start_frame
@@ -898,11 +916,11 @@ def determin_change_points(self) -> dict:
         """
         self.logger.info("Determine change points")
         if self.params.metric == 'euclideanDistance':
-            idx_dict = sp.get_local_max_and_min_idx(self.score['euclideanDistance'], self.video_info.fps)
+            idx_dict = sp.get_local_max_and_min_idx(self.score['euclideanDistance'], round(self.video_info.fps))
         elif self.params.metric == 'x':
-            idx_dict = sp.get_local_max_and_min_idx(self.score['x'], self.video_info.fps)
+            idx_dict = sp.get_local_max_and_min_idx(self.score['x'], round(self.video_info.fps))
         else:
-            idx_dict = sp.get_local_max_and_min_idx(self.score['y'], self.video_info.fps)
+            idx_dict = sp.get_local_max_and_min_idx(self.score['y'], round(self.video_info.fps))
         return idx_dict
 
 
@@ -976,7 +994,7 @@ def run(self) -> None:
                 self.logger.warning("Raw output is enabled!")
 
             # NOTE: score['y'] and score['x'] should have the same number size so it should be enouth to check one score length
-            with Listener(on_press=self.on_key_press) as listener:
+            with Listener(on_press=self.on_key_press) as _:
                 status = self.tracking()
 
                 if self.params.use_kalman_filter:

funscript_editor/data/ffmpegstream.py

@@ -96,7 +96,7 @@ def get_frame(
         cap = cv2.VideoCapture(video_path)
         if frame_number > 0:
             cap.set(cv2.CAP_PROP_POS_FRAMES, frame_number)
-        success, frame = cap.read()
+        _, frame = cap.read()
         cap.release()
         return frame
 
@@ -172,8 +172,10 @@ def get_projection(
                 )
 
         pipe.terminate()
-        pipe.stdin.close()
-        pipe.stdout.close()
+        try: pipe.stdin.close()
+        except: pass
+        try: pipe.stdout.close()
+        except: pass
 
         return projection
 
@@ -322,4 +324,5 @@ def run(self) -> None:
         self.stopped = True
         self.logger.info('Close FFmpeg Stream')
         pipe.terminate()
-        pipe.stdout.close()
+        try: pipe.stdout.close()
+        except: pass

funscript_editor/data/funscript.py

@@ -23,16 +23,12 @@
     - at : time to be at position in milliseconds
 """
 
-import copy
 import os
 import cv2
 import json
-import logging
 
 import numpy as np
 
-from PyQt5 import QtGui, QtCore ,QtWidgets
-
 
 class Funscript:
     """
@@ -57,9 +53,6 @@ def __init__(self, fps, version='1.0', inverted=False, moment_range=90):
         self.changed = False
 
 
-    __logger = logging.getLogger(__name__)
-
-
     def get_fps(self) -> float:
         """ Get Video FPS