Natsha commited on
Commit
0514416
·
1 Parent(s): 765b3de

Added docs and fixed the math for the data augmentation.

Browse files
Files changed (7) hide show
  1. fbx_handler.md +2 -1
  2. fbx_handler.py +72 -74
  3. globals.py +17 -0
  4. labeler/data_setup.py +636 -136
  5. preprocess_files.py +9 -12
  6. requirements.txt +4 -1
  7. utils.py +5 -18
fbx_handler.md CHANGED
@@ -26,13 +26,14 @@ actors_train, markers_train, t_test, _, _ = container.get_split_transforms(mode=
26
  ```
27
 
28
  ## Testing workflow:
 
29
  ```python
30
  # Load file.
31
  container = FBXContainer(input_file)
32
  # Get splitted original data (no transforms applied).
33
  actors_test, markers_test, t_test, r_test_, s_test = container.get_split_transforms(mode='test')
34
  # Predict the new actors and classes...
35
- actors_pred, markers_pred = Labeler(container.transform_translations(t_test))
36
  # Merge the new labels with their original translations.
37
  merged = merge_tdc(actors_pred, markers_pred, t_test, r_test, s_test)
38
  # Convert the full cloud into a dict structured for easy keyframes.
 
26
  ```
27
 
28
  ## Testing workflow:
29
+
30
  ```python
31
  # Load file.
32
  container = FBXContainer(input_file)
33
  # Get splitted original data (no transforms applied).
34
  actors_test, markers_test, t_test, r_test_, s_test = container.get_split_transforms(mode='test')
35
  # Predict the new actors and classes...
36
+ actors_pred, markers_pred = Labeler(scale_translations(t_test))
37
  # Merge the new labels with their original translations.
38
  merged = merge_tdc(actors_pred, markers_pred, t_test, r_test, s_test)
39
  # Convert the full cloud into a dict structured for easy keyframes.
fbx_handler.py CHANGED
@@ -22,7 +22,6 @@ def center_axis(a: Union[List[float], np.array]) -> np.array:
22
  # Turn list into np array for optimized math.
23
  if not isinstance(a, np.ndarray):
24
  a = np.array(a)
25
-
26
  # Find the centroid by subtracting the lowest value from the highest value.
27
  _min = np.min(a)
28
  _max = np.max(a)
@@ -296,7 +295,7 @@ def get_keyed_frames_from_curve(curve: fbx.FbxAnimCurve, length: int = -1) -> Li
296
 
297
 
298
  def get_world_transforms(actor_idx: int, marker_idx: int, m: fbx.FbxNode,
299
- r: List[int], c: fbx.FbxAnimCurve, incl_keyed: int = 1) -> List[List[float]]:
300
  """
301
  For the given marker node, gets the world transform for each frame in r, and stores the translation, rotation
302
  and scaling values as a list of lists. Stores the actor and marker classes at the start of this list of lists.
@@ -308,7 +307,6 @@ def get_world_transforms(actor_idx: int, marker_idx: int, m: fbx.FbxNode,
308
  :param m: `fbx.FbxNode` to evaluate the world transform of at each frame.
309
  :param r: `List[int]` list of frame numbers to evaluate the world transform at.
310
  :param c: `fbx.FbxAnimCurve` node to read the keyframes from.
311
- :param incl_keyed: `bool` whether to include if there was a key on a given frame or not. 0 if not.
312
  :return:
313
  """
314
  # Create a list of zeros with the same length as r.
@@ -341,19 +339,7 @@ def get_world_transforms(actor_idx: int, marker_idx: int, m: fbx.FbxNode,
341
  sy.append(wts[1])
342
  sz.append(wts[2])
343
 
344
- # If we don't need to include keyed frames, return the list of lists as is.
345
- if not incl_keyed:
346
- return [
347
- actors,
348
- markers,
349
- tx, ty, tz, zeros,
350
- rx, ry, rz, zeros,
351
- sx, sy, sz, ones
352
- ]
353
-
354
- # However, if we do need those keys, we first retrieve all the keyframed frame numbers from the curve.
355
- # Note: We do this after returning the previous results, because the following lines are very slow
356
- # and unnecessary for inference.
357
  keyed_frames = get_keyed_frames_from_curve(c)
358
  # Then we check if any of the frame numbers are in the keyed frames, which means it had a keyframe and should be 1.
359
  keyed_bools = [1 if f in keyed_frames else 0 for f in r]
@@ -365,7 +351,7 @@ def get_world_transforms(actor_idx: int, marker_idx: int, m: fbx.FbxNode,
365
  tx, ty, tz, zeros,
366
  rx, ry, rz, zeros,
367
  sx, sy, sz, ones,
368
- keyed_bools
369
  ]
370
 
371
 
@@ -382,14 +368,14 @@ def flatten_labeled_transforms(arr: np.array) -> np.array:
382
  """
383
  Flattens the given array so that it has the shape (n_actors * n_frames, 15, 73).
384
  :param arr: `np.array` to process.
385
- :return: `np.array` of shape (n_frames * n_actors, 15, 73).
386
  """
387
  # Transpose the array, so we get this order: (n_actors, n_frames, 15, 73).
388
  # That way, we can stack the actors after each other instead of the frames
389
  # (which would happen with the previous order).
390
  flattened = arr.transpose(1, 0, 2, 3)
391
  # Flatten the array, so we get a list of frames where with all actors stacked after each other.
392
- # Reshape to (n_frames * n_actors, 15, 73).
393
  return np.concatenate(flattened, axis=0)
394
 
395
 
@@ -403,10 +389,52 @@ def replace_zeros_with_inf(arr: np.array) -> np.array:
403
  # and set their transforms to np.inf.
404
  mask = arr[:, -1] == 0
405
  for i in range(arr.shape[0]):
406
- arr[i, 2:, mask[i]] = np.inf
407
  return arr
408
 
409
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
410
  class FBXContainerBase:
411
  def __init__(self, fbx_file: Path, debug: int = -1) -> None:
412
  """
@@ -430,7 +458,7 @@ class FBXContainerBase:
430
  """
431
  # Create an FBX manager and importer.
432
  self.manager = fbx.FbxManager.Create()
433
- importer = fbx.FbxImporter.Create(self.manager, '')
434
 
435
  # Import the FBX file.
436
  importer.Initialize(str(self.input_fbx))
@@ -720,16 +748,14 @@ class FBXContainer(FBXContainerBase):
720
  Calls the init functions for the labeled and unlabeled world transforms.
721
  :param r: Custom frame range to extract.
722
  """
723
- self.init_labeled_world_transforms(r=r, incl_keyed=1)
724
- self.init_unlabeled_world_transforms(r=r, incl_keyed=1)
725
 
726
- def init_labeled_world_transforms(self, r: Union[int, Tuple[int, int], Tuple[int, int, int]] = None,
727
- incl_keyed: int = 1) -> np.array:
728
  """
729
  For each actor, for each marker, stores a list for each element in the world transform for each frame
730
  in r. This can later be used to recreate the world transform matrix.
731
  :param r: Custom frame range to use.
732
- :param incl_keyed: `bool` whether to check if the marker was keyed at the frame.
733
  :return: `np.array` of shape (n_frames, 15, n_markers).
734
  """
735
  r = self.convert_r(r)
@@ -745,7 +771,7 @@ class FBXContainer(FBXContainerBase):
745
  # This requires the animation layer, so we can't do it within the function itself.
746
  curve = m.LclTranslation.GetCurve(self.anim_layer, 'X', True)
747
  # Get a list of each world transform element for all frames.
748
- marker_data = get_world_transforms(actor_idx + 1, marker_idx + 1, m, r, curve, incl_keyed)
749
  # Add the result to actor_data.
750
  actor_data.append(marker_data)
751
  self._print(f'Actor {actor_idx} marker {marker_idx} done', 1)
@@ -753,19 +779,17 @@ class FBXContainer(FBXContainerBase):
753
  labeled_data.append(actor_data)
754
 
755
  # Convert the list to a np array. This will have all frames at the last dimension because of this order:
756
- # Shape (n_actors, n_markers, 15, n_frames).
757
  wide_layout = np.array(labeled_data)
758
- # Transpose the array so that the order becomes (n_frames, n_actors, 15, n_markers).
759
  self.labeled_world_transforms = np.transpose(wide_layout, axes=(3, 0, 2, 1))
760
  return self.labeled_world_transforms
761
 
762
- def init_unlabeled_world_transforms(self, r: Union[int, Tuple[int, int], Tuple[int, int, int]] = None,
763
- incl_keyed: int = 1) -> np.array:
764
  """
765
  For all unlabeled markers, stores a list for each element in the world transform for each frame
766
  in r. This can later be used to recreate the world transform matrix.
767
  :param r: Custom frame range to use.
768
- :param incl_keyed: `bool` whether to check if the marker was keyed at the frame.
769
  :return: `np.array` of shape (n_frames, 15, n_unlabeled_markers).
770
  """
771
  r = self.convert_r(r)
@@ -777,15 +801,15 @@ class FBXContainer(FBXContainerBase):
777
  # This requires the animation layer, so we can't do it within the function itself.
778
  curve = ulm.LclTranslation.GetCurve(self.anim_layer, 'X', True)
779
  # Get a list of each world transform element for all frames.
780
- marker_data = get_world_transforms(0, 0, ulm, r, curve, incl_keyed=incl_keyed)
781
  # Add the result to marker_data.
782
  unlabeled_data.append(marker_data)
783
  self._print(f'Unlabeled marker {ulm.GetName()} done', 1)
784
 
785
  # Convert the list to a np array. This will have all frames at the last dimension because of this order:
786
- # Shape (n_unlabeled_markers, 15, n_frames).
787
  wide_layout = np.array(unlabeled_data)
788
- # Transpose the array so that the order becomes (n_frames, 15, n_unlabeled_markers).
789
  self.unlabeled_world_transforms = np.transpose(wide_layout, axes=(2, 1, 0))
790
  return self.unlabeled_world_transforms
791
 
@@ -825,21 +849,23 @@ class FBXContainer(FBXContainerBase):
825
  mask = mask_x1 & mask_x2 & mask_z1 & mask_z2
826
  return arr[mask]
827
 
828
- def extract_training_translations(self, r: Union[int, Tuple[int, int], Tuple[int, int, int]] = None) -> np.array:
 
829
  """
830
  Manipulates the existing labeled world transform array into one that is suitable for training.
831
  It does this through flattening the array to shape (n_frames, n_actors * 73, 15), then removing
832
  all clipping frames and finally transforms the frames to the right location and scale.
833
  :param r: Custom frame range to use if the labeled transforms are not stored yet.
 
834
  :return: Transformed labeled world transforms.
835
  """
836
  if self.labeled_world_transforms is None:
837
- self.init_labeled_world_transforms(r=r, incl_keyed=1)
838
 
839
  flattened = flatten_labeled_transforms(self.labeled_world_transforms)
840
  # Isolate the poses with all keyframes present by checking the last elements.
841
  # Start with the mask.
842
- # Returns shape of (n_frames * n_actors, 15, 73).
843
  mask = flattened[:, -1] == 1
844
  # We only need a filter for the first dimension, so use .all to check if all markers
845
  # have a keyframe. This results in shape (n_frames * n_actors,).
@@ -851,13 +877,7 @@ class FBXContainer(FBXContainerBase):
851
  # Remove any frames that cross the limits of the volume.
852
  flattened = self.remove_clipping_poses(flattened)
853
 
854
- for frame in range(flattened.shape[0]):
855
- # Center the X axis values.
856
- flattened[frame, 2] = center_axis(flattened[frame, 2])
857
- # Center the Z axis values.
858
- flattened[frame, 4] = center_axis(flattened[frame, 4])
859
-
860
- return self.transform_translations(flattened)
861
 
862
  def extract_inf_translations(self, r: Union[int, Tuple[int, int], Tuple[int, int, int]] = None,
863
  merged: bool = True) -> Union[np.array, Tuple[np.array, np.array]]:
@@ -872,18 +892,18 @@ class FBXContainer(FBXContainerBase):
872
  # If either of the arrays is None, we can initialize them with r.
873
  if self.labeled_world_transforms is None:
874
  # For inference, we don't need keyed frames, so incl_keyed is False.
875
- self.init_labeled_world_transforms(r=r, incl_keyed=1)
876
  if self.unlabeled_world_transforms is None:
877
  # Note: Unlabeled data is already flattened.
878
- self.init_unlabeled_world_transforms(r=r, incl_keyed=1)
879
 
880
- # Starting with (n_frames, n_actors, 15, 73).
881
  # Flatten the array, so we get a list of frames.
882
- # Returns shape (n_frames, 15, n_actors, 73).
883
  flat_labeled = self.labeled_world_transforms.transpose(0, 2, 1, 3)
884
 
885
  # Stack the elements in the last 2 dimension after each other.
886
- # Returns shape (n_frames, 15, n_actors * 73).
887
  ls = flat_labeled.shape
888
  flat_labeled = flat_labeled.reshape(ls[0], ls[1], -1)
889
  del ls
@@ -899,29 +919,6 @@ class FBXContainer(FBXContainerBase):
899
  else:
900
  return flat_labeled, self.unlabeled_world_transforms
901
 
902
- def transform_translations(self, arr: np.array) -> np.array:
903
- """
904
- Applies a scaling to the translation values in the given array.
905
- :param arr: `np.array` that can either be a timeline dense cloud or translation vectors.
906
- :return: Modified `np.array`.
907
- """
908
- # If the second dimension has 3 elements, it is a translation vector of shape (tx, ty, tz).
909
- # If it has 14 elements, it is a full marker row of shape (actor, marker, tx, ty, tz, tw, rx, ry, rz, rw, etc.).
910
- start = 0 if arr.shape[1] == 3 else 2
911
-
912
- # First multiply by self.scale, which turns centimeters to meters.
913
- # Then divide by volume dimensions, to normalize to the total area of the capture volume.
914
- arr[:, start + 0] *= self.scale / self.vol_x
915
- arr[:, start + 1] *= self.scale / self.vol_y
916
- arr[:, start + 2] *= self.scale / self.vol_z
917
-
918
- # Optional: Clip the translation values.
919
- # arr[:, start + 0] = np.clip(arr[:, start + 0], -0.5, 0.5)
920
- # arr[:, start + 1] = np.clip(arr[:, start + 1], -0.5, 0.5)
921
- # arr[:, start + 2] = np.clip(arr[:, start + 2], -0.5, 0.5)
922
-
923
- return arr
924
-
925
  def get_split_transforms(self, r: Union[int, Tuple[int, int], Tuple[int, int, int]] = None,
926
  mode: str = 'train') -> Tuple[np.array, np.array, np.array, np.array, np.array]:
927
  """
@@ -959,7 +956,7 @@ class FBXContainer(FBXContainerBase):
959
  Exports train data to an HDF5 file.
960
  :param output_file: `Path` to the file.
961
  :param r: Custom frame range to use.
962
- :return: `np.array` of shape (n_poses, 73, 14) of train data.
963
  """
964
  if output_file.suffix == '.h5':
965
  array_4d = self.extract_training_translations(r)
@@ -1164,6 +1161,7 @@ class FBXContainer(FBXContainerBase):
1164
  self._print(f'Replacing keys for actor {actor_idx}', 1)
1165
  self.replace_keyframes_per_actor(actor_idx, actor_dict)
1166
 
 
1167
  # if __name__ == '__main__':
1168
  # np.printoptions(precision=2, suppress=True)
1169
  # # container = FBXContainer(Path(r'G:\Firestorm\mocap-ai\data\fbx\dowg\TAKE_01+1_ALL_001.fbx'))
 
22
  # Turn list into np array for optimized math.
23
  if not isinstance(a, np.ndarray):
24
  a = np.array(a)
 
25
  # Find the centroid by subtracting the lowest value from the highest value.
26
  _min = np.min(a)
27
  _max = np.max(a)
 
295
 
296
 
297
  def get_world_transforms(actor_idx: int, marker_idx: int, m: fbx.FbxNode,
298
+ r: List[int], c: fbx.FbxAnimCurve) -> List[List[float]]:
299
  """
300
  For the given marker node, gets the world transform for each frame in r, and stores the translation, rotation
301
  and scaling values as a list of lists. Stores the actor and marker classes at the start of this list of lists.
 
307
  :param m: `fbx.FbxNode` to evaluate the world transform of at each frame.
308
  :param r: `List[int]` list of frame numbers to evaluate the world transform at.
309
  :param c: `fbx.FbxAnimCurve` node to read the keyframes from.
 
310
  :return:
311
  """
312
  # Create a list of zeros with the same length as r.
 
339
  sy.append(wts[1])
340
  sz.append(wts[2])
341
 
342
+ # Get the keyed values.
 
 
 
 
 
 
 
 
 
 
 
 
343
  keyed_frames = get_keyed_frames_from_curve(c)
344
  # Then we check if any of the frame numbers are in the keyed frames, which means it had a keyframe and should be 1.
345
  keyed_bools = [1 if f in keyed_frames else 0 for f in r]
 
351
  tx, ty, tz, zeros,
352
  rx, ry, rz, zeros,
353
  sx, sy, sz, ones,
354
+ r, keyed_bools
355
  ]
356
 
357
 
 
368
  """
369
  Flattens the given array so that it has the shape (n_actors * n_frames, 15, 73).
370
  :param arr: `np.array` to process.
371
+ :return: `np.array` of shape (n_actors * n_frames, 15, 73).
372
  """
373
  # Transpose the array, so we get this order: (n_actors, n_frames, 15, 73).
374
  # That way, we can stack the actors after each other instead of the frames
375
  # (which would happen with the previous order).
376
  flattened = arr.transpose(1, 0, 2, 3)
377
  # Flatten the array, so we get a list of frames where with all actors stacked after each other.
378
+ # Reshapes to (n_actors * n_frames, 15, 73).
379
  return np.concatenate(flattened, axis=0)
380
 
381
 
 
389
  # and set their transforms to np.inf.
390
  mask = arr[:, -1] == 0
391
  for i in range(arr.shape[0]):
392
+ arr[i, 2:-2, mask[i]] = np.inf
393
  return arr
394
 
395
 
396
+ def scale_translations(arr: np.array, scale: float = 0.01,
397
+ dims: Tuple[float, float, float] = (10., 10., 10.)) -> np.array:
398
+ """
399
+ Applies a scaling to the translation values in the given array.
400
+ :param arr: `np.array` that can either be a timeline dense cloud or translation vectors.
401
+ :param scale: `float` scaling factor.
402
+ :param dims: `tuple` of `float` values that determine the dimensions of the volume.
403
+ :return: Modified `np.array`.
404
+ """
405
+ # If the second dimension has 3 elements, it is a translation vector of shape (tx, ty, tz).
406
+ # If it has 15 elements, it is a full marker row of shape (actor, marker, tx, ty, tz, tw, rx, ry, rz, rw, etc.).
407
+ start = 0 if arr.shape[0] == 3 else 2
408
+
409
+ # First multiply by self.scale, which turns centimeters to meters.
410
+ # Then divide by volume dimensions, to normalize to the total area of the capture volume.
411
+ arr[:, start + 0] *= scale / dims[0]
412
+ arr[:, start + 1] *= scale / dims[1]
413
+ arr[:, start + 2] *= scale / dims[2]
414
+
415
+ return arr
416
+
417
+
418
+ def transform_translations(arr: np.array, move_to_center: bool = True,
419
+ scale: float = 0.01, dims: Tuple[float, float, float] = (10., 10., 10.)) -> np.array:
420
+ """
421
+ First moves the x and y values to their axis' center. Then scales all values to normalize them.
422
+ :param arr: `np.array` that can either be a timeline dense cloud or translation vectors.
423
+ :param move_to_center: Uses center_axis() to move the x and y translations to the center of their axes.
424
+ :param scale: `float` scaling factor.
425
+ :param dims: `tuple` of `float` values that determine the dimensions of the volume.
426
+ :return: Modified `np.array`.
427
+ """
428
+ if move_to_center:
429
+ for frame in range(arr.shape[0]):
430
+ # Center the X axis values.
431
+ arr[frame, 2] = center_axis(arr[frame, 2])
432
+ # Center the Z axis values.
433
+ arr[frame, 4] = center_axis(arr[frame, 4])
434
+
435
+ return scale_translations(arr, scale, dims)
436
+
437
+
438
  class FBXContainerBase:
439
  def __init__(self, fbx_file: Path, debug: int = -1) -> None:
440
  """
 
458
  """
459
  # Create an FBX manager and importer.
460
  self.manager = fbx.FbxManager.Create()
461
+ importer = fbx.FbxImporter.Create(self.manager, 'MyScene')
462
 
463
  # Import the FBX file.
464
  importer.Initialize(str(self.input_fbx))
 
748
  Calls the init functions for the labeled and unlabeled world transforms.
749
  :param r: Custom frame range to extract.
750
  """
751
+ self.init_labeled_world_transforms(r=r)
752
+ self.init_unlabeled_world_transforms(r=r)
753
 
754
+ def init_labeled_world_transforms(self, r: Union[int, Tuple[int, int], Tuple[int, int, int]] = None) -> np.array:
 
755
  """
756
  For each actor, for each marker, stores a list for each element in the world transform for each frame
757
  in r. This can later be used to recreate the world transform matrix.
758
  :param r: Custom frame range to use.
 
759
  :return: `np.array` of shape (n_frames, 15, n_markers).
760
  """
761
  r = self.convert_r(r)
 
771
  # This requires the animation layer, so we can't do it within the function itself.
772
  curve = m.LclTranslation.GetCurve(self.anim_layer, 'X', True)
773
  # Get a list of each world transform element for all frames.
774
+ marker_data = get_world_transforms(actor_idx + 1, marker_idx + 1, m, r, curve)
775
  # Add the result to actor_data.
776
  actor_data.append(marker_data)
777
  self._print(f'Actor {actor_idx} marker {marker_idx} done', 1)
 
779
  labeled_data.append(actor_data)
780
 
781
  # Convert the list to a np array. This will have all frames at the last dimension because of this order:
782
+ # Shape (n_actors, n_markers, 16, n_frames).
783
  wide_layout = np.array(labeled_data)
784
+ # Transpose the array so that the order becomes (n_frames, n_actors, 16, n_markers).
785
  self.labeled_world_transforms = np.transpose(wide_layout, axes=(3, 0, 2, 1))
786
  return self.labeled_world_transforms
787
 
788
+ def init_unlabeled_world_transforms(self, r: Union[int, Tuple[int, int], Tuple[int, int, int]] = None) -> np.array:
 
789
  """
790
  For all unlabeled markers, stores a list for each element in the world transform for each frame
791
  in r. This can later be used to recreate the world transform matrix.
792
  :param r: Custom frame range to use.
 
793
  :return: `np.array` of shape (n_frames, 15, n_unlabeled_markers).
794
  """
795
  r = self.convert_r(r)
 
801
  # This requires the animation layer, so we can't do it within the function itself.
802
  curve = ulm.LclTranslation.GetCurve(self.anim_layer, 'X', True)
803
  # Get a list of each world transform element for all frames.
804
+ marker_data = get_world_transforms(0, 0, ulm, r, curve)
805
  # Add the result to marker_data.
806
  unlabeled_data.append(marker_data)
807
  self._print(f'Unlabeled marker {ulm.GetName()} done', 1)
808
 
809
  # Convert the list to a np array. This will have all frames at the last dimension because of this order:
810
+ # Shape (n_unlabeled_markers, 16, n_frames).
811
  wide_layout = np.array(unlabeled_data)
812
+ # Transpose the array so that the order becomes (n_frames, 16, n_unlabeled_markers).
813
  self.unlabeled_world_transforms = np.transpose(wide_layout, axes=(2, 1, 0))
814
  return self.unlabeled_world_transforms
815
 
 
849
  mask = mask_x1 & mask_x2 & mask_z1 & mask_z2
850
  return arr[mask]
851
 
852
+ def extract_training_translations(self, r: Union[int, Tuple[int, int], Tuple[int, int, int]] = None,
853
+ move_to_center: bool = True) -> np.array:
854
  """
855
  Manipulates the existing labeled world transform array into one that is suitable for training.
856
  It does this through flattening the array to shape (n_frames, n_actors * 73, 15), then removing
857
  all clipping frames and finally transforms the frames to the right location and scale.
858
  :param r: Custom frame range to use if the labeled transforms are not stored yet.
859
+ :param move_to_center: If True, the x and y axes is moved to the center of the volume.
860
  :return: Transformed labeled world transforms.
861
  """
862
  if self.labeled_world_transforms is None:
863
+ self.init_labeled_world_transforms(r=r)
864
 
865
  flattened = flatten_labeled_transforms(self.labeled_world_transforms)
866
  # Isolate the poses with all keyframes present by checking the last elements.
867
  # Start with the mask.
868
+ # Returns shape of (n_frames * n_actors, 16, 73).
869
  mask = flattened[:, -1] == 1
870
  # We only need a filter for the first dimension, so use .all to check if all markers
871
  # have a keyframe. This results in shape (n_frames * n_actors,).
 
877
  # Remove any frames that cross the limits of the volume.
878
  flattened = self.remove_clipping_poses(flattened)
879
 
880
+ return transform_translations(flattened, move_to_center, self.scale, (self.vol_x, self.vol_y, self.vol_z))
 
 
 
 
 
 
881
 
882
  def extract_inf_translations(self, r: Union[int, Tuple[int, int], Tuple[int, int, int]] = None,
883
  merged: bool = True) -> Union[np.array, Tuple[np.array, np.array]]:
 
892
  # If either of the arrays is None, we can initialize them with r.
893
  if self.labeled_world_transforms is None:
894
  # For inference, we don't need keyed frames, so incl_keyed is False.
895
+ self.init_labeled_world_transforms(r=r)
896
  if self.unlabeled_world_transforms is None:
897
  # Note: Unlabeled data is already flattened.
898
+ self.init_unlabeled_world_transforms(r=r)
899
 
900
+ # Starting with (n_frames, n_actors, 16, 73).
901
  # Flatten the array, so we get a list of frames.
902
+ # Returns shape (n_frames, 16, n_actors, 73).
903
  flat_labeled = self.labeled_world_transforms.transpose(0, 2, 1, 3)
904
 
905
  # Stack the elements in the last 2 dimension after each other.
906
+ # Returns shape (n_frames, 16, n_actors * 73).
907
  ls = flat_labeled.shape
908
  flat_labeled = flat_labeled.reshape(ls[0], ls[1], -1)
909
  del ls
 
919
  else:
920
  return flat_labeled, self.unlabeled_world_transforms
921
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
922
  def get_split_transforms(self, r: Union[int, Tuple[int, int], Tuple[int, int, int]] = None,
923
  mode: str = 'train') -> Tuple[np.array, np.array, np.array, np.array, np.array]:
924
  """
 
956
  Exports train data to an HDF5 file.
957
  :param output_file: `Path` to the file.
958
  :param r: Custom frame range to use.
959
+ :return: `np.array` of shape (n_poses, 14, 73) of train data.
960
  """
961
  if output_file.suffix == '.h5':
962
  array_4d = self.extract_training_translations(r)
 
1161
  self._print(f'Replacing keys for actor {actor_idx}', 1)
1162
  self.replace_keyframes_per_actor(actor_idx, actor_dict)
1163
 
1164
+
1165
  # if __name__ == '__main__':
1166
  # np.printoptions(precision=2, suppress=True)
1167
  # # container = FBXContainer(Path(r'G:\Firestorm\mocap-ai\data\fbx\dowg\TAKE_01+1_ALL_001.fbx'))
globals.py CHANGED
@@ -13,3 +13,20 @@ def get_marker_names():
13
  'RIDX3', 'RIDX6', 'RMID0', 'RMID6', 'RRNG3', 'RRNG6', 'RPNK3', 'RPNK6', 'LFWT', 'MFWT',
14
  'RFWT', 'LBWT', 'MBWT', 'RBWT', 'LTHI', 'LKNE', 'LKNI', 'LSHN', 'LANK', 'LHEL', 'LMT5',
15
  'LMT1', 'LTOE', 'RTHI', 'RKNE', 'RKNI', 'RSHN', 'RANK', 'RHEL', 'RMT5', 'RMT1', 'RTOE')
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
13
  'RIDX3', 'RIDX6', 'RMID0', 'RMID6', 'RRNG3', 'RRNG6', 'RPNK3', 'RPNK6', 'LFWT', 'MFWT',
14
  'RFWT', 'LBWT', 'MBWT', 'RBWT', 'LTHI', 'LKNE', 'LKNI', 'LSHN', 'LANK', 'LHEL', 'LMT5',
15
  'LMT1', 'LTOE', 'RTHI', 'RKNE', 'RKNI', 'RSHN', 'RANK', 'RHEL', 'RMT5', 'RMT1', 'RTOE')
16
+
17
+
18
+ def get_joint_names():
19
+ return ('Hips', 'Spine', 'Spine1', 'Spine2', 'Spine3', 'Neck', 'Neck1', 'Head', 'HeadEnd',
20
+ 'RightShoulder', 'RightArm', 'RightForeArm', 'RightHand', 'RightHandMiddle1',
21
+ 'RightHandMiddle2', 'RightHandMiddle3', 'RightHandMiddle4', 'RightHandRing',
22
+ 'RightHandRing1', 'RightHandRing2', 'RightHandRing3', 'RightHandRing4', 'RightHandPinky',
23
+ 'RightHandPinky1', 'RightHandPinky2', 'RightHandPinky3', 'RightHandPinky4', 'RightHandIndex',
24
+ 'RightHandIndex1', 'RightHandIndex2', 'RightHandIndex3', 'RightHandIndex4', 'RightHandThumb1',
25
+ 'RightHandThumb2', 'RightHandThumb3', 'RightHandThumb4', 'LeftShoulder', 'LeftArm',
26
+ 'LeftForeArm', 'LeftHand', 'LeftHandMiddle1', 'LeftHandMiddle2', 'LeftHandMiddle3',
27
+ 'LeftHandMiddle4', 'LeftHandRing', 'LeftHandRing1', 'LeftHandRing2', 'LeftHandRing3',
28
+ 'LeftHandRing4', 'LeftHandPinky', 'LeftHandPinky1', 'LeftHandPinky2', 'LeftHandPinky3',
29
+ 'LeftHandPinky4', 'LeftHandIndex', 'LeftHandIndex1', 'LeftHandIndex2', 'LeftHandIndex3',
30
+ 'LeftHandIndex4', 'LeftHandThumb1', 'LeftHandThumb2', 'LeftHandThumb3', 'LeftHandThumb4',
31
+ 'RightUpLeg', 'RightLeg', 'RightFoot', 'RightToeBase', 'RightToeBaseEnd', 'LeftUpLeg',
32
+ 'LeftLeg', 'LeftFoot', 'LeftToeBase', 'LeftToeBaseEnd')
labeler/data_setup.py CHANGED
@@ -1,161 +1,661 @@
1
  from pathlib import Path
2
- from typing import Tuple
 
3
 
 
4
  import numpy as np
5
  import torch
6
- from torch.utils import data
7
- import math
8
-
9
-
10
- def apply_random_y_rotation(point_cloud_data: torch.Tensor) -> torch.Tensor:
11
- # Convert the random angle from degrees to radians
12
- angle = (torch.rand(1).item() * 2 - 1) * 180 * torch.tensor(math.pi / 180, device='cuda')
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
13
 
14
  # Create the rotation matrix for the y-axis
15
- rotation_matrix = torch.tensor([[torch.cos(angle), 0, torch.sin(angle)],
16
- [0, 1, 0],
17
- [-torch.sin(angle), 0, torch.cos(angle)]], device='cuda')
18
-
19
- # Apply the rotation to the point cloud data
20
- return torch.matmul(point_cloud_data, rotation_matrix.T)
21
-
22
-
23
- class PointCloudDataset(data.Dataset):
24
- def __init__(self, file: Path,
25
- n_samples=100,
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
26
  max_actors: int = 8,
27
- translation_factor=0.1,
28
- max_overlap: Tuple[float] = (0.2, 0.2, 0.2)):
29
- point_clouds_np = torch.tensor(np.load(str(file)), dtype=torch.float32, device='cuda')
30
- self.sparse_point_clouds = point_clouds_np
31
- self.n_samples = n_samples
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
32
  self.max_actors = max_actors
 
33
  self.translation_factor = translation_factor
34
- self.max_overlap = max_overlap
35
-
36
- # Generate a random permutation of indices.
37
- self.indices = torch.randperm(len(self.sparse_point_clouds))
38
 
39
- dataset = []
40
- for _ in range(n_samples):
41
- accumulated_cloud = []
42
- # TODO: Get a random number up to the max of actors.
43
- # TODO: Transform one row of the available rows, and check if it doesn't overlap.
44
- # TODO: Accumulate all actors into one point cloud and append that to dataset.
45
- # TODO: __getitem__() needs to get one of these point cloud rows.
46
- for i in range(max_actors):
47
-
48
- # Get a point cloud from the tensor using the shuffled index, shape (1, 1024).
49
- point_cloud = self.sparse_point_clouds[self.indices[index]]
50
-
51
- point_cloud_data = point_cloud[:, 2:5] # returns shape: (1024, 3)
52
-
53
- valid_transform = False
54
- while not valid_transform:
55
-
56
- point_cloud = point_cloud_data.clone()
57
- # Randomly translate the point cloud along the x and z axes
58
-
59
- self.apply_random_translation(point_cloud)
60
- # Apply random rotation around the y-axis
61
- rotated_point_cloud_data = apply_random_y_rotation(point_cloud)
62
-
63
- if not does_overlap(accumulated_cloud, point_cloud, self.max_overlap):
64
- accumulated_cloud.append(point_cloud)
65
- valid_transform = True
66
 
67
- def apply_random_translation(self, point_cloud: torch.Tensor) -> None:
68
- x_translation = (torch.rand(1).item() * 2 - 1) * self.translation_factor
69
- z_translation = (torch.rand(1).item() * 2 - 1) * self.translation_factor
70
- point_cloud[:, [0, 2]] += torch.tensor([x_translation, z_translation], device='cuda')
71
 
72
- def fill_point_cloud(self, point_cloud):
73
- target_num_points = 73 * self.max_actors
74
- current_num_points = point_cloud.shape[1]
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
75
 
76
- if current_num_points < target_num_points:
77
- num_points_to_add = target_num_points - current_num_points
78
- random_indices = torch.randint(0, current_num_points, (num_points_to_add,))
79
- additional_points = point_cloud[:, random_indices, :]
80
 
81
- filled_point_cloud = torch.cat((point_cloud, additional_points), dim=1)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
82
  else:
83
- filled_point_cloud = point_cloud
84
 
85
- return filled_point_cloud
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
86
 
87
  def __getitem__(self, index):
88
-
89
- point_cloud = np.vstack(accumulated_cloud)
90
- # Separate the labels from the point cloud data
91
- actor_labels = point_cloud[:, :, 0] # shape: (1024,)
92
- marker_labels = point_cloud[:, :, 1] # shape: (1024,)
93
-
94
- return actor_labels, marker_labels, rotated_point_cloud_data
95
 
96
  def __len__(self):
97
- return len(self.sparse_point_clouds)
98
-
99
-
100
- def does_overlap(accumulated_point_cloud, new_point_cloud, overlap_thresholds=(0.2, 0.2, 0.2)):
101
- def project_to_axis(point_cloud, axis):
102
- projected_points = point_cloud.clone()
103
- projected_points[:, axis] = 0
104
- return projected_points
105
-
106
- def get_bounding_box_2d(points):
107
- min_values, _ = torch.min(points, dim=0)
108
- max_values, _ = torch.max(points, dim=0)
 
 
 
 
 
 
 
 
109
  return min_values, max_values
110
 
111
- def check_surface_area_overlap(bb1_min, bb1_max, bb2_min, bb2_max, axis, overlap_threshold):
112
- bb1_area = (bb1_max[axis] - bb1_min[axis]) * (bb1_max[1] - bb1_min[1])
113
- bb2_area = (bb2_max[axis] - bb2_min[axis]) * (bb2_max[1] - bb2_min[1])
114
-
115
- overlap_min = torch.max(bb1_min, bb2_min)
116
- overlap_max = torch.min(bb1_max, bb2_max)
117
-
118
- overlap_area = (overlap_max[axis] - overlap_min[axis]) * (overlap_max[1] - overlap_min[1])
119
- overlap_area = torch.max(torch.tensor(0.0, device='cuda'), overlap_area) # Clamp to 0 if negative
120
-
121
- overlap_percentage = overlap_area / torch.min(bb1_area, bb2_area)
122
-
123
- return overlap_percentage >= overlap_threshold
124
-
125
- new_point_cloud_xz = project_to_axis(new_point_cloud, 1) # Project to xz-plane (remove y-axis values)
126
- new_point_cloud_min, new_point_cloud_max = get_bounding_box_2d(new_point_cloud_xz)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
127
 
128
  overlaps = []
129
 
130
- for pc in accumulated_point_cloud:
131
- for axis in range(len(overlap_thresholds)):
132
- pc_xz = project_to_axis(pc, axis) # Project to xz-plane (remove y-axis values)
133
- pc_min, pc_max = get_bounding_box_2d(pc_xz)
134
-
135
- if all(
136
- check_surface_area_overlap(
137
- new_point_cloud_min,
138
- new_point_cloud_max,
139
- pc_min,
140
- pc_max,
141
- axis,
142
- overlap_thresholds[axis],
143
- )
144
- for axis in range(len(overlap_thresholds))
145
- ):
146
- return True
147
-
148
- return False
149
-
150
-
151
- class NoOverlapDataLoader(data.DataLoader):
152
- def __init__(self, dataset: data.Dataset, max_overlap: Tuple[float] = (0.2, 0.2, 0.2), *args, **kwargs):
153
- super().__init__(dataset, *args, **kwargs)
154
- self.max_overlap = max_overlap
155
-
156
- def __iter__(self):
157
- accumulated_point_clouds = []
158
- for actor_labels, marker_labels, point_cloud_data in super().__iter__():
159
- if not does_overlap(accumulated_point_clouds, point_cloud_data, self.max_overlap):
160
- accumulated_point_clouds.append(point_cloud_data)
161
- yield actor_labels, marker_labels, point_cloud_data
 
1
  from pathlib import Path
2
+ from typing import Tuple, List, Union
3
+ from random import randint
4
 
5
+ import h5py
6
  import numpy as np
7
  import torch
8
+ from torch import Tensor
9
+ from torch.utils.data import Dataset
10
+ import matplotlib.pyplot as plt
11
+
12
+ import fbx_handler
13
+ import utils
14
+
15
+
16
+ def apply_y_rotation(point_cloud_data: Tensor, angle: float = None, device: str = 'cuda') -> Tensor:
17
+ """
18
+ Apply a random rotation to the point cloud.
19
+ :param point_cloud_data: `Tensor` of shape (3, 73) to modify.
20
+ :param angle: Angle as `float` in degrees to rotate the point cloud. If this is given, the rotation is not random.
21
+ :param device: `str` device on which to create the extra tensors.
22
+ :return: Modified `Tensor`.
23
+ """
24
+ # Convert the random angle from degrees to radians.
25
+ if angle is None:
26
+ # If no angle is given, use a random angle between -180 and 180.
27
+ angle = (torch.rand(1).item() * 2 - 1) * 180 * torch.tensor(torch.pi / 180, device=device)
28
+ else:
29
+ # If an angle is given, convert this angle instead.
30
+ angle *= torch.tensor(torch.pi / 180, device=device)
31
+
32
+ # Transpose the point_cloud_data from (3, 73) to (73, 3) so we can use torch.matmul.
33
+ point_cloud_data = point_cloud_data.transpose(1, 0)
34
 
35
  # Create the rotation matrix for the y-axis
36
+ rotation_matrix = torch.tensor([
37
+ [torch.cos(angle), 0, torch.sin(angle)],
38
+ [0, 1, 0],
39
+ [-torch.sin(angle), 0, torch.cos(angle)]], device=device)
40
+
41
+ # Apply the rotation to the point cloud data and reverse the transpose to get back to the original shape (3, 73).
42
+ return torch.matmul(point_cloud_data, rotation_matrix).transpose(1, 0)
43
+
44
+
45
+ def fill_1d_tensor_with_zeros(point_cloud: Tensor, pc_size: int = 1024, device: str = 'cuda') -> Tensor:
46
+ """
47
+ Fill a 1D tensor with zeros, so it is as long as pc_size.
48
+ :param point_cloud: `Tensor` of shape (73,) to add zeros to.
49
+ :param pc_size: `int` amount of points that need to be in the final tensor in total.
50
+ :param device: `str` device on which to create the extra tensors.
51
+ :return: `Tensor` of shape (pc_size,).
52
+ """
53
+ length = len(point_cloud)
54
+ if length < pc_size:
55
+ zeros = torch.zeros(pc_size - length, dtype=torch.int, device=device)
56
+ point_cloud = torch.cat((point_cloud, zeros), dim=0)
57
+
58
+ # Since we don't check if the length is longer than pc_size, always return the tensor with the pc_size slice.
59
+ return point_cloud[:pc_size]
60
+
61
+
62
+ def fill_frames_tensor(point_cloud: Tensor, pc_size: int = 1024, filler: int = -1, device: str = 'cuda') -> Tensor:
63
+ """
64
+ Fill a 1D tensor with ones, so it is as long as pc_size.
65
+ :param point_cloud: `Tensor` of shape (73,) to add `int` -1s to.
66
+ :param pc_size: `int` amount of points that need to be in the final tensor in total.
67
+ :param filler: `int` value to fill the remainder of the tensor with.
68
+ :param device: `str` device on which to create the extra tensors.
69
+ :return: `Tensor` of shape (pc_size,).
70
+ """
71
+ length = len(point_cloud)
72
+ if length < pc_size:
73
+ zeros = torch.full((pc_size - length,), filler, dtype=torch.int, device=device)
74
+ point_cloud = torch.cat((point_cloud, zeros), dim=0)
75
+
76
+ # Since we don't check if the length is longer than pc_size, always return the tensor with the pc_size slice.
77
+ return point_cloud[:pc_size]
78
+
79
+
80
+ def convert_max_overlap(max_overlap: Union[Tuple[float, float, float], float]) -> Tuple[float, float, float]:
81
+ """
82
+ Convert the argument max_overlap to a float tuple of length 3.
83
+ :param max_overlap: Either 3 floats or 1 float.
84
+ :return: If max_overlap is 3 floats, returns max_overlap unchanged.
85
+ If it is 1 `float`, returns a tuple of size 3 of that `float`.
86
+ """
87
+ if isinstance(max_overlap, float):
88
+ return max_overlap, max_overlap, max_overlap
89
+ if len(max_overlap) != 3:
90
+ raise ValueError(f'max_overlap must be a tuple of length 3, not {len(max_overlap)}.')
91
+ return max_overlap
92
+
93
+
94
+ def convert_n_samples(n_samples: Union[int, float], _max: int) -> int:
95
+ """
96
+ Convert the argument n_samples to an `int` that serves as a total samples amount.
97
+ :param n_samples: Either a `float` (representing a ratio) or an `int` (representing a number of samples).
98
+ :param _max: `int` that indicates the highest possible n_samples.
99
+ :return: An int that is never higher than _max.
100
+ """
101
+ # If n_samples is between 0-1, it is considered a ratio, and we calculate the amount of rows to use.
102
+ if isinstance(n_samples, float):
103
+ n_samples = int(n_samples * _max)
104
+ # If n_samples is negative, subtract the amount from the total amount of rows.
105
+ elif n_samples < 0:
106
+ n_samples = _max - n_samples
107
+ # If n_samples is 0, use all rows.
108
+ elif n_samples == 0 or n_samples > _max:
109
+ n_samples = _max
110
+
111
+ return n_samples
112
+
113
+
114
+ def plot_point_cloud(point_cloud: Tensor, scale: Union[int, float] = 50):
115
+ tensor = point_cloud.cpu().numpy()
116
+ # Extract x, y, and z coordinates from the tensor
117
+ x = tensor[:, 0]
118
+ y = tensor[:, 1]
119
+ z = tensor[:, 2]
120
+
121
+ # Create a 3D plot
122
+ fig = plt.figure()
123
+ ax = fig.add_subplot(111, projection='3d')
124
+
125
+ # Scatter plot
126
+ ax.scatter(x, y, z, s=scale)
127
+
128
+ # Set axis labels
129
+ ax.set_xlabel('X')
130
+ ax.set_ylabel('Y')
131
+ ax.set_zlabel('Z')
132
+
133
+ ax.set_xlim([-0.5, 0.5])
134
+ ax.set_ylim([-0.5, 0.5])
135
+ ax.set_zlim([-0.5, 0.5])
136
+
137
+ ax.zaxis._axinfo['juggled'] = (1, 1, 0)
138
+ ax.xaxis.pane.fill = False
139
+ ax.yaxis.pane.fill = False
140
+ ax.zaxis.pane.fill = False
141
+
142
+ # Show the plot
143
+ plt.show()
144
+
145
+
146
+ def compare_point_clouds(existing, title='plot'):
147
+ colors = plt.cm.jet(np.linspace(0, 1, len(existing)))
148
+
149
+ n_tensors = len(existing)
150
+ plt.figure(figsize=(10, 7))
151
+ for idx, tensor in enumerate(existing):
152
+ tensor = tensor.cpu().numpy()
153
+ # Extract the first and third elements
154
+ x_coords = tensor[0]
155
+ z_coords = tensor[2]
156
+
157
+ # Create a scatter plot
158
+ plt.scatter(x_coords, z_coords, c=colors[idx], label=f'Tensor {idx + 1}', s=1)
159
+
160
+ plt.show()
161
+
162
+
163
+ def fill_translation_cloud(translations: Tensor, n_points: int = 1024, augment=torch.rand,
164
+ apply_shuffle: bool = True, shuffle: Tensor = None, device: str = 'cuda') \
165
+ -> Tuple[Tensor, Tensor]:
166
+ """
167
+ Fill a translation tensor with filler data, so it is as long as pc_size.
168
+ :param translations: `Tensor` of shape (3, xxx).
169
+ :param n_points: `int` amount of total points that need to be in the output.
170
+ :param augment: Torch filler function to use for generating filler points, default `torch.rand`.
171
+ :param apply_shuffle: `bool` whether to shuffle the output.
172
+ :param shuffle: `Tensor` that contains a shuffled index order that needs to be used for shuffling.
173
+ This does nothing if apply_shuffle is False.
174
+ :param device: `str` device on which to create the extra tensors.
175
+ :return: Translation and shuffle tuple of `Tensor` of shape (3, n_points), and (n_points,).
176
+ """
177
+ # Use the second dimension as the length of the translation tensor, due to input shape (3, 73..).
178
+ length = translations.shape[1]
179
+ # Only create filler data if the length is shorter than the amount of points.
180
+ if length < n_points:
181
+ # Calculate the shape of the extra tensor, and pass it to the given augment function.
182
+ dif = (translations.shape[0], n_points - length)
183
+ extra = augment(dif, device=device)
184
+
185
+ # Concatenate all values together to get shape (3, pc_size).
186
+ translations = torch.cat((translations, extra), dim=1)
187
+ else:
188
+ translations = translations[:, :n_points]
189
+
190
+ # Shuffle if needed.
191
+ if apply_shuffle:
192
+ if shuffle is None:
193
+ shuffle = torch.randperm(n_points, device=device)
194
+
195
+ translations = torch.index_select(translations, 1, shuffle)
196
+
197
+ return translations, shuffle
198
+
199
+
200
+ def fill_point_clouds(actor_classes: Tensor, marker_classes: Tensor, translations: Tensor, frames: Tensor,
201
+ n_points: int = 1024, augment=torch.rand, apply_shuffle: bool = True, shuffle: Tensor = None,
202
+ device: str = 'cuda') \
203
+ -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor]:
204
+ """
205
+ Fill a point cloud with filler data, so it is as long as pc_size.
206
+ :param actor_classes: `Tensor` of shape (n_points,) that contains the actor classes.
207
+ :param marker_classes: `Tensor` of shape (n_points,) that contains the marker classes.
208
+ :param translations: `Tensor` of shape (3, n_points) that contains the marker translations.
209
+ :param frames: `Tensor` of shape (n_points,) that contains the animated frames.
210
+ :param n_points: `int` amount of total points that need to be in the output.
211
+ :param augment: Torch filler function to use for generating filler points, default `torch.rand`.
212
+ :param apply_shuffle: `bool` whether to shuffle the output.
213
+ :param shuffle: `Tensor` that contains a shuffled index order that needs to be used for shuffling. This does nothing if apply_shuffle is False.
214
+ :param device: `str` device on which to create the extra tensors.
215
+ :return: Tuple of `Tensor` of shape (n_points,), (n_points,), (3,n_points,), (n_points,), (n_points,)
216
+ that represent the actor classes, marker classes, translations, animated frames and the shuffled indices used.
217
+ """
218
+ # Use simple functions to create full tensors for the actors/markers/frames.
219
+ actor_classes = fill_1d_tensor_with_zeros(actor_classes, n_points, device=device)
220
+ marker_classes = fill_1d_tensor_with_zeros(marker_classes, n_points, device=device)
221
+ frames = fill_frames_tensor(frames, n_points, device=device)
222
+
223
+ # Extend the translation tensor.
224
+ length = translations.shape[1]
225
+ if length < n_points:
226
+ dif = (3, n_points - length)
227
+ extra = augment(dif, device=device)
228
+
229
+ # Concatenate all values together to get shape (pc_size,).
230
+ translations = torch.cat((translations, extra), dim=1)
231
+ else:
232
+ translations = translations[:, :n_points]
233
+
234
+ # Shuffle if needed.
235
+ if apply_shuffle:
236
+
237
+ if shuffle is None:
238
+ shuffle = torch.randperm(n_points, device=device)
239
+
240
+ actor_classes = torch.index_select(actor_classes, 0, shuffle)
241
+ marker_classes = torch.index_select(marker_classes, 0, shuffle)
242
+ translations = torch.index_select(translations, 1, shuffle)
243
+ frames = torch.index_select(frames, 0, shuffle)
244
+
245
+ # Returns a list of tensors of shape (n_points,), (n_points,), (3, n_points), (n_points,).
246
+ return actor_classes, marker_classes, translations, frames, shuffle
247
+
248
+
249
+ def remove_inf_markers(labeled: np.ndarray, device: str = 'cuda'):
250
+ """
251
+ Goes through the labeled data and removes all markers that have inf features. This will also scale the translations.
252
+ :param labeled: `np.ndarray` of shape (15, n_points) that contains the labeled data.
253
+ :param device: `str` device on which to create the extra tensors.
254
+ :return: Tuple of `tensor` that represent actors/markers/scaled translations/unscaled translations/frames.
255
+ """
256
+ # Check if the second feature (tx) is inf. This means it had no keyframe,
257
+ # and the NN should not classify this to avoid the network learning interpolated markers.
258
+ # Mask is True if it had a keyframe.
259
+ mask = ~np.isinf(labeled[2])
260
+
261
+ # Make tensors from the np arrays.
262
+ actor_cloud = torch.tensor(labeled[0][mask], dtype=torch.int, device=device)
263
+ marker_cloud = torch.tensor(labeled[1][mask], dtype=torch.int, device=device)
264
+ unscaled_t_cloud = labeled[2:5][:, mask]
265
+ frames = torch.tensor(labeled[-1][mask], dtype=torch.int, device=device)
266
+
267
+ # Scale the translations into a separate tensor.
268
+ scaled_t_cloud = fbx_handler.scale_translations(unscaled_t_cloud)
269
+ scaled_t_cloud = torch.tensor(scaled_t_cloud, dtype=torch.float32, device=device)
270
+
271
+ # After the scaled_t_cloud is made, we can convert the unscaled_t_cloud to a tensor too.
272
+ unscaled_t_cloud = torch.tensor(unscaled_t_cloud, dtype=torch.float32, device=device)
273
+ return actor_cloud, marker_cloud, scaled_t_cloud, unscaled_t_cloud, frames
274
+
275
+
276
+ def apply_translation(point_cloud: Tensor, t: float = 1.0, device: str = 'cuda') -> Tensor:
277
+ """
278
+ Apply a translation to all axes of a point cloud.
279
+ :param point_cloud: `Tensor` of shape (3, n_points) that contains the point cloud.
280
+ :param t: `float` that represents the translation.
281
+ :param device: `str` device on which to create the extra tensors.
282
+ :return: `Tensor` of shape (3, n_points) that contains the point cloud with the translation applied.
283
+ """
284
+ point_cloud[0] += torch.tensor(t, device=device)
285
+ point_cloud[1] += torch.tensor(t, device=device)
286
+ point_cloud[2] += torch.tensor(t, device=device)
287
+ return point_cloud
288
+
289
+
290
+ class TrainDataset(Dataset):
291
+ def __init__(self, file: Union[Path, np.array],
292
+ n_samples: Union[int, float] = 1.0,
293
+ n_attempts: int = 10,
294
+ pc_size: int = 1024,
295
  max_actors: int = 8,
296
+ use_random_max_actors: bool = True,
297
+ use_random_translation: bool = True,
298
+ use_random_rotation: bool = True,
299
+ shuffle_markers: bool = True,
300
+ translation_factor: float = 0.9,
301
+ max_overlap: Union[Tuple[float, float, float], float] = (0.2, 0.2, 0.2),
302
+ augment=torch.rand,
303
+ debug: int = -1,
304
+ device: str = 'cuda'):
305
+ self.debug = debug
306
+ self.device = device
307
+
308
+ # If the pc_size is a number under 73, we intend to use it as a multiplication.
309
+ if pc_size < 73:
310
+ pc_size *= 73
311
+ elif pc_size < max_actors * 73:
312
+ raise ValueError(f'pc_size must be large enough to contain 73 markers for {max_actors} actors '
313
+ f'({pc_size}/{max_actors * 73}).')
314
+
315
+ # Store most arguments as class properties, so they don't have to be passed to each function.
316
+ # These will all be deleted after the dataset is created.
317
+ self.n_attempts = n_attempts
318
+ self.pc_size = pc_size
319
  self.max_actors = max_actors
320
+ self.shuffle_markers = shuffle_markers
321
  self.translation_factor = translation_factor
322
+ self.max_overlap = convert_max_overlap(max_overlap)
 
 
 
323
 
324
+ # Isolate the dependent and independent variables.
325
+ if isinstance(file, np.ndarray):
326
+ self.all_data = file
327
+ else:
328
+ self.all_data = utils.h5_to_array4d(file)
329
+ # Shape (n_frames, 15, 73).
330
+ self.all_data = torch.tensor(self.all_data, dtype=torch.float32, device=device)
331
+ self.n_samples = convert_n_samples(n_samples, self.all_data.shape[0])
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
332
 
333
+ self._print(f'Loaded in {len(self.all_data)} poses, with n_samples = {n_samples}.', 0)
 
 
 
334
 
335
+ # Generate a random permutation of indices.
336
+ self.random_indices = torch.randperm(len(self.all_data))
337
+ self.random_idx = 0
338
+
339
+ # Initiate empty lists for all the different types of data.
340
+ actor_classes, marker_classes, translations, frames = [], [], [], []
341
+
342
+ # For each sample, create a random point cloud.
343
+ for _ in range(self.n_samples):
344
+ cur_max_actors = randint(1, max_actors) if use_random_max_actors else max_actors
345
+ actor_cloud, marker_cloud, translation_cloud, fs = self.create_sample(cur_max_actors,
346
+ use_random_rotation,
347
+ use_random_translation, augment)
348
+
349
+ actor_classes.append(actor_cloud)
350
+ marker_classes.append(marker_cloud)
351
+ translations.append(translation_cloud)
352
+ frames.append(fs)
353
+
354
+ # (n_samples, pc_size), (n_samples, pc_size), (n_samples, 3, pc_size), (n_samples,pc_size).
355
+ self.actor_classes = torch.stack(actor_classes)
356
+ self.marker_classes = torch.stack(marker_classes)
357
+ self.translations = torch.stack(translations)
358
+ self.frames = torch.stack(frames)
359
+
360
+ # Delete class properties that were only needed to create the dataset.
361
+ del self.pc_size, self.max_actors, self.shuffle_markers, self.translation_factor, self.n_samples, \
362
+ self.max_overlap, self.all_data, self.random_indices, self.random_idx, self.n_attempts
363
+
364
+ def _print(self, txt: str, lvl: int = 0) -> None:
365
+ if lvl <= self.debug:
366
+ print(txt)
367
+
368
+ def create_sample(self, max_actors: int, use_random_rotation: bool = True,
369
+ use_random_translation: bool = True, augment=torch.rand) -> Tuple[Tensor, Tensor, Tensor, Tensor]:
370
+ """
371
+ Create a random point cloud from the dataset.
372
+ :param max_actors: `int` amount of actors to aim for in this point cloud. Any missing markers will be filled.
373
+ :param use_random_rotation: `bool` whether to apply a random rotation to each actor's point cloud.
374
+ :param use_random_translation: `bool` whether to apply a random translation to each actor's point cloud.
375
+ :param augment: Torch function to use for the filler markers. Examples are `torch.rand`, `torch.ones`, etc.
376
+ :return: A tuple of tensors containing the actor point cloud, marker point cloud, and translation point cloud.
377
+ """
378
+ # Loop through all cur_max_actors, select a row from all_data, and concatenate it to the t_cloud.
379
+ actor_cloud, marker_cloud, t_cloud, frames = [], [], [], []
380
+ # For each actor, try 10 times to find a point cloud that does not overlap the accumulated cloud.
381
+ # If it fails all times, we will just have fewer actors in the point cloud.
382
+ for actor_idx in range(max_actors):
383
+ for attempt in range(self.n_attempts):
384
+ # In case we ever have lots of attempts, reset the random index if we have reached the end of the data.
385
+ if self.random_idx == len(self.all_data):
386
+ self.random_idx = 0
387
+
388
+ # Get a pose from the tensor using the shuffled index; shape (1, 14, 73).
389
+ row = self.all_data[self.random_indices[self.random_idx]]
390
+ self.random_idx += 1
391
+
392
+ # Collect relevant data from the row.
393
+ # Shapes: (73,).
394
+ a = row[0].to(torch.int)
395
+ m = row[1].to(torch.int)
396
+ f = row[-1].to(torch.int)
397
+
398
+ # Shape (3, 73).
399
+ t = row[2:5]
400
+ # Apply random rotation and translations if needed.
401
+ if use_random_rotation:
402
+ t = apply_y_rotation(t, device=self.device)
403
+ if use_random_translation:
404
+ t = self.apply_random_translation(t)
405
+
406
+ self._print(f'Checking overlap for {actor_idx} - {attempt}', 1)
407
+ if does_overlap(t_cloud, t, max_overlap=self.max_overlap):
408
+ # If the clouds overlap too much, we continue to the next attempt without adding this one.
409
+ print(f'Actor {actor_idx + 1} attempt {attempt + 1} failed.')
410
+ continue
411
+
412
+ # Add data to their respective lists if the clouds don't overlap.
413
+ actor_cloud.append(a)
414
+ marker_cloud.append(m)
415
+ t_cloud.append(t)
416
+ frames.append(f)
417
+
418
+ self._print(f'Actor {actor_idx + 1} attempt {attempt + 1} succeeded.', 1)
419
+ # If the clouds don't overlap too much,
420
+ # we break the loop because this attempt worked, and we don't need another one.
421
+ break
422
+
423
+ self._print(f'Total length: {len(t_cloud)}/{max_actors}', 0)
424
+ # Add all lists together to create long tensors.
425
+ # Shape (n_actors * 73,).
426
+ actor_cloud = torch.cat(actor_cloud, dim=0)
427
+ marker_cloud = torch.cat(marker_cloud, dim=0)
428
+ frames = torch.cat(frames, dim=0)
429
+ # Shape (3, n_actors * 73).
430
+ t_cloud = torch.cat(t_cloud, dim=1)
431
+
432
+ # Fill the clouds with more markers to get to pc_size.
433
+ # (1024,), (1024,), (1024, 3), (1024,).
434
+ actor_cloud, marker_cloud, t_cloud, frames, _ = fill_point_clouds(
435
+ actor_cloud, marker_cloud, t_cloud, frames, n_points=self.pc_size,
436
+ augment=augment, apply_shuffle=self.shuffle_markers, device=self.device)
437
+
438
+ return actor_cloud, marker_cloud, t_cloud, frames
439
+
440
+ def apply_random_translation(self, point_cloud: Tensor) -> Tensor:
441
+ """
442
+ Apply random translation to the point cloud.
443
+ :param point_cloud: `Tensor` of shape (3, n_points).
444
+ :return: Translated `Tensor` of shape (3, n_points).
445
+ """
446
+ x_translation = (torch.rand(1).item() - 0.5) * self.translation_factor
447
+ z_translation = (torch.rand(1).item() - 0.5) * self.translation_factor
448
+ point_cloud[0] += torch.tensor(x_translation, device=self.device)
449
+ point_cloud[2] += torch.tensor(z_translation, device=self.device)
450
+ return point_cloud
451
 
452
+ def __getitem__(self, index):
453
+ return self.actor_classes[index], self.marker_classes[index], self.translations[index], self.frames[index]
 
 
454
 
455
+ def __len__(self):
456
+ return len(self.actor_classes)
457
+
458
+
459
+ class InfDataset(Dataset):
460
+ def __init__(self, source: Union[Path, Tuple[np.ndarray, np.ndarray]],
461
+ pc_size: int = 1024,
462
+ n_samples: Union[int, float] = 1.0,
463
+ augment=torch.rand,
464
+ shuffle_markers: bool = False,
465
+ debug: int = -1,
466
+ device: str = 'cuda') -> None:
467
+ self.device = device
468
+ self.debug = debug
469
+
470
+ if isinstance(source, np.ndarray):
471
+ labeled_data, unlabeled_data = source
472
  else:
 
473
 
474
+ # if isinstance(source, Path):
475
+ # # if source.stem == 'ALL':
476
+ # # self.data = utils.combined_test_h5_to_array4d(source, pc_size)
477
+ # # else:
478
+ with h5py.File(source, 'r') as h5f:
479
+ labeled_data = np.array(h5f['labeled'])[:5]
480
+ unlabeled_data = np.array(h5f['unlabeled'])[:5]
481
+ # self.data = utils.merge_labeled_and_unlabeled_data(labeled_data, unlabeled_data, pc_size, augment)
482
+ # else:
483
+ # labeled_data, unlabeled_data = source
484
+ self.assemble_data(augment, labeled_data, unlabeled_data, pc_size, n_samples, shuffle_markers)
485
+
486
+ self._print(f'Actors: {self.actor_classes.shape}, markers: {self.marker_classes.shape}, '
487
+ f'translations: {self.translations.shape}', 0)
488
+ self._print(self.actor_classes[:, :10], 0)
489
+ self._print(self.marker_classes[:, :10], 0)
490
+ self._print(self.translations[:, :, :10], 0)
491
+ self._print(self.unscaled_translations[:, :, :10], 0)
492
+ self._print(self.frames[:, :10], 0)
493
+
494
+ def _print(self, txt: str, lvl: int = 0) -> None:
495
+ if lvl <= self.debug:
496
+ print(txt)
497
+
498
+ def assemble_data(self, augment, labeled_data: np.ndarray, unlabeled_data: np.ndarray, pc_size: int = 1024,
499
+ n_samples: int = 5, shuffle_markers: bool = False):
500
+ """
501
+ Assemble the various tensors.
502
+ :param augment: Torch function to use for the filler markers. Examples are `torch.rand`, `torch.ones`, etc.
503
+ :param labeled_data: `np.ndarray` that contains the data of the labeled markers.
504
+ :param unlabeled_data: `np.ndarray` that contains the data of the unlabeled markers.
505
+ :param pc_size: `int` amount of points to put in the point cloud.
506
+ :param n_samples: Total amount of samples to generate.
507
+ :param shuffle_markers: `bool` whether to shuffle the markers in the point cloud.
508
+ """
509
+ n_samples = convert_n_samples(n_samples, len(labeled_data))
510
+ # Initialize empty lists to store the data in.
511
+ actor_classes, marker_classes, translations, unscaled_translations, frames = [], [], [], [], []
512
+ for frame in range(n_samples):
513
+ labeled = labeled_data[frame]
514
+ unlabeled = unlabeled_data[frame]
515
+
516
+ actor_cloud, marker_cloud, scaled_t_cloud, unscaled_t_cloud, l_frames = remove_inf_markers(
517
+ labeled, device=self.device)
518
+
519
+ ul_actor_cloud, ul_marker_cloud, ul_scaled_t_cloud, ul_unscaled_t_cloud, ul_frames = \
520
+ remove_inf_markers(unlabeled, device=self.device)
521
+
522
+ merged_actors = torch.cat([actor_cloud, ul_actor_cloud], dim=0)
523
+ merged_markers = torch.cat([marker_cloud, ul_marker_cloud], dim=0)
524
+ merged_translations = torch.cat([scaled_t_cloud, ul_scaled_t_cloud], dim=1)
525
+ merged_unscaled_translations = torch.cat([unscaled_t_cloud, ul_unscaled_t_cloud], dim=1)
526
+ merged_frames = torch.cat([l_frames, ul_frames], dim=0)
527
+
528
+ # fill_point_clouds() uses the augment function to fill the point clouds, so we can't use it to
529
+ # fill the unscaled translations.
530
+ actor_cloud, marker_cloud, scaled_t_cloud, merged_frames, shuffled_idx = \
531
+ fill_point_clouds(merged_actors, merged_markers, merged_translations, merged_frames,
532
+ n_points=pc_size, augment=augment, apply_shuffle=shuffle_markers, device=self.device)
533
+
534
+ # use fill_translation_cloud to fill the unscaled translations.
535
+ # This is a separate function because fill_point_clouds() is also used in the TrainDataset class.
536
+ merged_unscaled_translations, _ = fill_translation_cloud(merged_unscaled_translations, n_points=pc_size,
537
+ augment=augment, apply_shuffle=shuffle_markers,
538
+ shuffle=shuffled_idx, device=self.device)
539
+
540
+ actor_classes.append(actor_cloud)
541
+ marker_classes.append(marker_cloud)
542
+ translations.append(scaled_t_cloud)
543
+ unscaled_translations.append(merged_unscaled_translations)
544
+ frames.append(merged_frames)
545
+
546
+ # (n_samples, pc_size), (n_samples, pc_size), (n_samples, 3, pc_size).
547
+ self.actor_classes = torch.stack(actor_classes)
548
+ self.marker_classes = torch.stack(marker_classes)
549
+ self.translations = torch.stack(translations)
550
+ self.unscaled_translations = torch.stack(unscaled_translations)
551
+ self.frames = torch.stack(frames)
552
 
553
  def __getitem__(self, index):
554
+ return self.actor_classes[index], self.marker_classes[index], \
555
+ self.translations[index], self.unscaled_translations[index], self.frames[index]
 
 
 
 
 
556
 
557
  def __len__(self):
558
+ return len(self.actor_classes)
559
+
560
+
561
+ def does_overlap(accumulated_point_cloud: List[Tensor], new_point_cloud: Tensor,
562
+ max_overlap: Tuple[float, float, float] = (0.2, 0.2, 0.2)) -> bool:
563
+ """
564
+ Checks if a new point cloud overlaps with any of the existing point clouds.
565
+ :param accumulated_point_cloud: List of `Tensor` of the accumulated point clouds.
566
+ :param new_point_cloud: `Tensor` point cloud to check overlap for.
567
+ :param max_overlap: Tuple of 3 floats to indicate allowed overlapping thresholds for each axis.
568
+ :return: `bool` whether the new point cloud overlaps with any of the existing point clouds.
569
+ """
570
+ def get_bounding_box(points: Tensor) -> Tuple[Tensor, Tensor]:
571
+ """
572
+ Gets the bounding box values (min, max) for each axis.
573
+ :param points: `Tensor` point cloud to analyze.
574
+ :return: Tuple of `Tensor` of minimum and maximum values.
575
+ """
576
+ min_values, _ = torch.min(points, dim=1)
577
+ max_values, _ = torch.max(points, dim=1)
578
  return min_values, max_values
579
 
580
+ def check_dimensional_overlap(bb1_min: Tensor, bb1_max: Tensor, bb2_min: Tensor, bb2_max: Tensor,
581
+ overlap_threshold: float = 0.2) -> bool:
582
+ """
583
+ Checks if two bounding boxes overlap in one axis.
584
+ :param bb1_min: `Tensor` of minimum value for the first bounding box.
585
+ :param bb1_max: `Tensor` of maximum value for the first bounding box.
586
+ :param bb2_min: `Tensor` of minimum value for the second bounding box.
587
+ :param bb2_max: `Tensor` of maximum value for the second bounding box.
588
+ :param overlap_threshold: `float` that indicates the maximum % of overlap allowed for this axis.
589
+ :return: `bool` whether the two bounding boxes overlap.
590
+ """
591
+ # Find the highest bbox minimum and the lowest bbox maximum.
592
+ overlap_min = torch.maximum(bb1_min, bb2_min)
593
+ overlap_max = torch.minimum(bb1_max, bb2_max)
594
+ # Calculate the overlap length. If the bounding boxes don't overlap, this length will be negative.
595
+ # Then we can return False right away.
596
+ overlap_length = overlap_max - overlap_min
597
+ if overlap_length <= 0:
598
+ return False
599
+
600
+ # Given that the overlap length is a positive number, we need to calculate how much overlap is happening.
601
+ # First find the outer bounds of the both bounding boxes (lowest minimum and highest maximum).
602
+ non_overlap_min = torch.minimum(bb1_min, bb2_min)
603
+ non_overlap_max = torch.maximum(bb1_max, bb2_max)
604
+ # Then calculate what fraction of the total length is the overlapping length.
605
+ total_length = non_overlap_max - non_overlap_min
606
+ overlap_ratio = overlap_length / total_length
607
+ # Return whether this ratio is higher than the allowed threshold.
608
+ return overlap_ratio > overlap_threshold
609
+
610
+ def check_3dimensional_overlap(bb1_min: Tensor, bb1_max: Tensor, bb2_min: Tensor, bb2_max: Tensor,
611
+ overlap_thresholds: Tuple[float, float, float]) -> bool:
612
+ """
613
+ Checks if two 3-dimensional bounding boxes overlap in the x and z axis.
614
+ :param bb1_min: `Tensor` of minimum values for the first bounding box.
615
+ :param bb1_max: `Tensor` of maximum values for the first bounding box.
616
+ :param bb2_min: `Tensor` of minimum values for the second bounding box.
617
+ :param bb2_max: `Tensor` of maximum values for the second bounding box.
618
+ :param overlap_thresholds: Tuple of 3 `float` that indicates the maximum % of overlap allowed for all axes.
619
+ :return: `bool` whether the two bounding boxes overlap.
620
+ """
621
+ x_overlap = check_dimensional_overlap(bb1_min[0], bb1_max[0], bb2_min[0], bb2_max[0], overlap_thresholds[0])
622
+ z_overlap = check_dimensional_overlap(bb1_min[2], bb1_max[2], bb2_min[2], bb2_max[2], overlap_thresholds[2])
623
+ # EXTRA: Check if the y axes are overlapping.
624
+ return x_overlap and z_overlap
625
+
626
+ # If this is the first attempt of checking an overlap, the accumulated point cloud is empty,
627
+ # so we don't need to check any overlap.
628
+ if not accumulated_point_cloud:
629
+ return False
630
+
631
+ # Find the bounding box values of the new point cloud.
632
+ new_min, new_max = get_bounding_box(new_point_cloud)
633
 
634
  overlaps = []
635
 
636
+ # Iterate through each point cloud in the accumulated list.
637
+ for idx, pc in enumerate(accumulated_point_cloud):
638
+ # Get the bounding box for the current cloud.
639
+ current_min, current_max = get_bounding_box(pc)
640
+ # Check if the new point cloud overlaps with the current cloud.
641
+ overlaps.append(check_3dimensional_overlap(current_min, current_max, new_min, new_max, max_overlap))
642
+
643
+ # If any axis of any point cloud overlapped, we don't want to add the point cloud.
644
+ return any(overlaps)
645
+
646
+
647
+ if __name__ == '__main__':
648
+ # train_dataset = TrainDataset(Path(r'G:\Firestorm\mocap-ai\data\h5\mes-1\train\IntroVideo_04_006.h5'),
649
+ # n_samples=1,
650
+ # max_actors=2,
651
+ # pc_size=2,
652
+ # use_random_max_actors=False,
653
+ # use_random_translation=True,
654
+ # use_random_rotation=False,
655
+ # shuffle_markers=False,
656
+ # max_overlap=.9)
657
+ # print(dir(train_dataset))
658
+ test_dataset = InfDataset(Path(r'G:\Firestorm\mocap-ai\data\h5\mes-1\test\HangoutSpot_1_001.h5'),
659
+ pc_size=150,
660
+ shuffle_markers=False,
661
+ debug=0)
 
 
 
 
 
 
preprocess_files.py CHANGED
@@ -1,15 +1,14 @@
1
  from pathlib import Path
2
- import shutil
3
  import multiprocessing
4
 
5
  # Import custom libs.
6
  import fbx_handler
7
  import utils
8
 
9
-
10
- source = Path('G:/Firestorm/mocap-ai/data/fbx/mes-2/')
11
- train_folder = Path('G:/Firestorm/mocap-ai/data/h5/mes-2/train')
12
- test_folder = Path('G:/Firestorm/mocap-ai/data/h5/mes-2/test')
13
 
14
 
15
  def process_fbx_file(fbx_file: Path):
@@ -25,7 +24,7 @@ def process_fbx_file(fbx_file: Path):
25
  print(fbx_file)
26
 
27
  # Create a new class object with the file path.
28
- my_obj = fbx_handler.FBXContainer(fbx_file, max_actors=4, pc_size=296, debug=0, save_init=True)
29
  # Init world transforms for labeled and unlabeled data. This will store all relevant transform info.
30
  with utils.Timer('Getting world transforms took'):
31
  try:
@@ -45,20 +44,18 @@ def process_fbx_file(fbx_file: Path):
45
 
46
  try:
47
  # Do the same thing for the test data.
48
- test_data = my_obj.export_test_data(export_test_path, merged=False)
49
  print(f'Test labeled shape: {test_data[0].shape}')
50
  print(f'Test unlabeled shape: {test_data[1].shape}')
51
- print(f'Minimum cloud size: {test_data[0].shape[1] + test_data[1].shape[1]}')
52
  except BaseException as e:
53
  print(e)
54
  return
55
 
56
 
57
- def process_fbx_files(source_folder: Path, v: int = 1):
58
  # Delete the existing folders and make them again, because the array4d_to_h5 function will append
59
  # # new data to any existing files.
60
- shutil.rmtree(train_folder)
61
- shutil.rmtree(test_folder)
62
  train_folder.mkdir(parents=True, exist_ok=True)
63
  test_folder.mkdir(parents=True, exist_ok=True)
64
 
@@ -69,7 +66,7 @@ def process_fbx_files(source_folder: Path, v: int = 1):
69
  # train_all = train_folder / 'ALL.h5'
70
  # test_all = test_folder / 'ALL.h5'
71
 
72
- with multiprocessing.Pool(4) as pool:
73
  pool.map(process_fbx_file, files)
74
 
75
  # print('--- FINAL ---')
 
1
  from pathlib import Path
 
2
  import multiprocessing
3
 
4
  # Import custom libs.
5
  import fbx_handler
6
  import utils
7
 
8
+ c = 'dowg'
9
+ source = Path(f'G:/Firestorm/mocap-ai/data/fbx/{c}/')
10
+ train_folder = Path(f'G:/Firestorm/mocap-ai/data/h5/{c}/train')
11
+ test_folder = Path(f'G:/Firestorm/mocap-ai/data/h5/{c}/test')
12
 
13
 
14
  def process_fbx_file(fbx_file: Path):
 
24
  print(fbx_file)
25
 
26
  # Create a new class object with the file path.
27
+ my_obj = fbx_handler.FBXContainer(fbx_file, debug=0)
28
  # Init world transforms for labeled and unlabeled data. This will store all relevant transform info.
29
  with utils.Timer('Getting world transforms took'):
30
  try:
 
44
 
45
  try:
46
  # Do the same thing for the test data.
47
+ test_data = my_obj.export_inf_data(export_test_path, merged=False)
48
  print(f'Test labeled shape: {test_data[0].shape}')
49
  print(f'Test unlabeled shape: {test_data[1].shape}')
50
+ print(f'Minimum cloud size: {test_data[0].shape[2] + test_data[1].shape[2]}')
51
  except BaseException as e:
52
  print(e)
53
  return
54
 
55
 
56
+ def process_fbx_files(source_folder: Path):
57
  # Delete the existing folders and make them again, because the array4d_to_h5 function will append
58
  # # new data to any existing files.
 
 
59
  train_folder.mkdir(parents=True, exist_ok=True)
60
  test_folder.mkdir(parents=True, exist_ok=True)
61
 
 
66
  # train_all = train_folder / 'ALL.h5'
67
  # test_all = test_folder / 'ALL.h5'
68
 
69
+ with multiprocessing.Pool(1) as pool:
70
  pool.map(process_fbx_file, files)
71
 
72
  # print('--- FINAL ---')
requirements.txt CHANGED
@@ -2,4 +2,7 @@ streamlit~=1.21.0
2
  pandas~=1.3.5
3
  numpy~=1.21.5
4
  torch~=1.13.1
5
- h5py
 
 
 
 
2
  pandas~=1.3.5
3
  numpy~=1.21.5
4
  torch~=1.13.1
5
+ h5py~=3.7.0
6
+ torchinfo~=1.7.2
7
+ seaborn~=0.12.2
8
+ matplotlib~=3.5.3
utils.py CHANGED
@@ -2,7 +2,7 @@ import cProfile
2
  import pstats
3
  import time
4
  from pathlib import Path
5
- from typing import List, Tuple
6
 
7
  import h5py
8
  import numpy as np
@@ -22,19 +22,6 @@ def append_suffix_to_file(file_path: Path, suffix: str = '_INF', ext: str = None
22
  return file_path.with_name(new_file_name)
23
 
24
 
25
- def is_int_in_list(n: int, l: List[int]) -> int:
26
- if l[0] > n:
27
- return 0
28
-
29
- for e in l:
30
- if e == n:
31
- return 1
32
- elif e > n:
33
- return 0
34
-
35
- return 0
36
-
37
-
38
  def array4d_to_h5(array_4ds: Tuple, output_file: Path, group: str = None, datasets: Tuple = 'array_data'):
39
  if len(array_4ds) != len(datasets):
40
  raise ValueError(f'Amount of arrays {len(array_4ds)} must match amount of dataset names {len(datasets)}.')
@@ -53,7 +40,7 @@ def h5_to_array4d(input_file: Path) -> np.array:
53
  return np.vstack([np.array(h5f[key]) for key in h5f.keys()])
54
 
55
 
56
- def combined_test_h5_to_array4d(input_file: Path, pc_size: int = 1024) -> np.array:
57
  with h5py.File(input_file, 'r') as h5f:
58
  data = []
59
  for grp_name in list(h5f.keys()):
@@ -69,10 +56,10 @@ def merge_labeled_and_unlabeled_data(labeled: np.array, unlabeled: np.array, pc_
69
  augment: str = None) -> np.array:
70
  missing = pc_size - (labeled.shape[2] + unlabeled.shape[2])
71
  if missing <= 0:
72
- # Returns shape (n_frames, self.pc_size, 14).
73
  return np.concatenate((unlabeled, labeled), axis=2)[:, :, -pc_size:]
74
 
75
- # This is similar to the way that fill_point_cloud() fills values.
76
  if augment is None:
77
  missing_markers = np.ones((labeled.shape[0], labeled.shape[1], missing))
78
  elif augment == 'normal':
@@ -83,7 +70,7 @@ def merge_labeled_and_unlabeled_data(labeled: np.array, unlabeled: np.array, pc_
83
  missing_markers[:, 0] = 0.
84
  missing_markers[:, 1] = 0.
85
 
86
- # Returns shape (n_frames, self.pc_size, 14).
87
  return np.concatenate((missing_markers,
88
  unlabeled,
89
  labeled), axis=2)
 
2
  import pstats
3
  import time
4
  from pathlib import Path
5
+ from typing import Tuple
6
 
7
  import h5py
8
  import numpy as np
 
22
  return file_path.with_name(new_file_name)
23
 
24
 
 
 
 
 
 
 
 
 
 
 
 
 
 
25
  def array4d_to_h5(array_4ds: Tuple, output_file: Path, group: str = None, datasets: Tuple = 'array_data'):
26
  if len(array_4ds) != len(datasets):
27
  raise ValueError(f'Amount of arrays {len(array_4ds)} must match amount of dataset names {len(datasets)}.')
 
40
  return np.vstack([np.array(h5f[key]) for key in h5f.keys()])
41
 
42
 
43
+ def combined_test_h5_to_array4d(input_file: Path, pc_size: int = 1024, merged: bool = True) -> np.array:
44
  with h5py.File(input_file, 'r') as h5f:
45
  data = []
46
  for grp_name in list(h5f.keys()):
 
56
  augment: str = None) -> np.array:
57
  missing = pc_size - (labeled.shape[2] + unlabeled.shape[2])
58
  if missing <= 0:
59
+ # Returns shape (n_frames, 15, self.pc_size).
60
  return np.concatenate((unlabeled, labeled), axis=2)[:, :, -pc_size:]
61
 
62
+ # This is similar to the way that TrainDataset.fill_point_cloud() fills values.
63
  if augment is None:
64
  missing_markers = np.ones((labeled.shape[0], labeled.shape[1], missing))
65
  elif augment == 'normal':
 
70
  missing_markers[:, 0] = 0.
71
  missing_markers[:, 1] = 0.
72
 
73
+ # Returns shape (n_frames, 15, self.pc_size).
74
  return np.concatenate((missing_markers,
75
  unlabeled,
76
  labeled), axis=2)