update data.py

utils/data.py

@@ -135,34 +135,48 @@ def extract_joints(
 
 def preprocess(
     source: str,
-    data_args: dict,
     keypoints_detector,
-    device: str = 'cpu',
-) -> torch.Tensor:
+    normalization: bool = True,
+    random_choose: bool = True,
+    window_size: int = 120,
+) -> np.ndarray:
     '''
     Preprocess the video.
-
     Parameters
     ----------
     source : str
         The path to the video.
-
+    keypoints_detector : mediapipe.solutions.holistic.Holistic
+        The keypoints detector.
+    normalization : bool, default=True
+        Whether to normalize the data.
+    random_choose : bool, default=True
+        Whether to randomly sample the data.
+    window_size : int, default=120
+        The window size.
     Returns
     -------
-    dict
-        The model inputs.
+    np.ndarray
+        The processed inputs for model.
     '''
-    print('Extracting joints from pose...')
-    inputs = extract_joints_new(source=source, keypoints_detector=keypoints_detector)
+    inputs = extract_joints(source=source, keypoints_detector=keypoints_detector)
+
     T = inputs.shape[1]
-    print('Sampling video...')
-    if data_args['random_choose']:
-        inputs = random_sample_np(inputs, data_args['window_size'])
+    # ori_data = inputs
+    # for t in range(T - 1):
+    #     inputs[:, t, :, :] = ori_data[:, t + 1, :, :] - ori_data[:, t, :, :]
+    # inputs[:, T - 1, :, :] = 0
+
+    if random_choose:
+        inputs = random_sample_np(inputs, window_size)
     else:
-        inputs = uniform_sample_np(inputs, data_args['window_size'])
+        inputs = uniform_sample_np(inputs, window_size)
+
+    # if normalization:
+    #     assert inputs.shape[0] == 3
+    #     inputs[0, :, :, :] = inputs[0, :, :, :] - inputs[0, :, 0, 0].mean(axis=0)
+    #     inputs[1, :, :, :] = inputs[1, :, :, :] - inputs[1, :, 0, 0].mean(axis=0)
 
-    print('Normalizing video...')
-    print(inputs.shape, inputs)
     return np.squeeze(inputs).transpose(1, 2, 0).astype(np.float32)