Added docs and fixed the math for the data augmentation.
Browse files- fbx_handler.md +2 -1
- fbx_handler.py +72 -74
- globals.py +17 -0
- labeler/data_setup.py +636 -136
- preprocess_files.py +9 -12
- requirements.txt +4 -1
- 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(
|
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
|
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 |
-
#
|
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 (
|
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 |
-
#
|
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
|
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
|
724 |
-
self.init_unlabeled_world_transforms(r=r
|
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
|
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,
|
757 |
wide_layout = np.array(labeled_data)
|
758 |
-
# Transpose the array so that the order becomes (n_frames, n_actors,
|
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
|
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,
|
787 |
wide_layout = np.array(unlabeled_data)
|
788 |
-
# Transpose the array so that the order becomes (n_frames,
|
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
|
|
|
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
|
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,
|
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 |
-
|
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
|
876 |
if self.unlabeled_world_transforms is None:
|
877 |
# Note: Unlabeled data is already flattened.
|
878 |
-
self.init_unlabeled_world_transforms(r=r
|
879 |
|
880 |
-
# Starting with (n_frames, n_actors,
|
881 |
# Flatten the array, so we get a list of frames.
|
882 |
-
# Returns shape (n_frames,
|
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,
|
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,
|
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
|
7 |
-
import
|
8 |
-
|
9 |
-
|
10 |
-
|
11 |
-
|
12 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
13 |
|
14 |
# Create the rotation matrix for the y-axis
|
15 |
-
rotation_matrix = torch.tensor([
|
16 |
-
|
17 |
-
|
18 |
-
|
19 |
-
|
20 |
-
|
21 |
-
|
22 |
-
|
23 |
-
|
24 |
-
|
25 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
26 |
max_actors: int = 8,
|
27 |
-
|
28 |
-
|
29 |
-
|
30 |
-
|
31 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
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 |
-
|
40 |
-
|
41 |
-
|
42 |
-
|
43 |
-
|
44 |
-
|
45 |
-
|
46 |
-
|
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 |
-
|
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 |
-
|
73 |
-
|
74 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
75 |
|
76 |
-
|
77 |
-
|
78 |
-
random_indices = torch.randint(0, current_num_points, (num_points_to_add,))
|
79 |
-
additional_points = point_cloud[:, random_indices, :]
|
80 |
|
81 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
82 |
else:
|
83 |
-
filled_point_cloud = point_cloud
|
84 |
|
85 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
86 |
|
87 |
def __getitem__(self, index):
|
88 |
-
|
89 |
-
|
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.
|
98 |
-
|
99 |
-
|
100 |
-
def does_overlap(accumulated_point_cloud, new_point_cloud
|
101 |
-
|
102 |
-
|
103 |
-
|
104 |
-
|
105 |
-
|
106 |
-
|
107 |
-
|
108 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
109 |
return min_values, max_values
|
110 |
|
111 |
-
def
|
112 |
-
|
113 |
-
|
114 |
-
|
115 |
-
|
116 |
-
|
117 |
-
|
118 |
-
|
119 |
-
|
120 |
-
|
121 |
-
|
122 |
-
|
123 |
-
|
124 |
-
|
125 |
-
|
126 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
127 |
|
128 |
overlaps = []
|
129 |
|
130 |
-
|
131 |
-
|
132 |
-
|
133 |
-
|
134 |
-
|
135 |
-
|
136 |
-
|
137 |
-
|
138 |
-
|
139 |
-
|
140 |
-
|
141 |
-
|
142 |
-
|
143 |
-
|
144 |
-
|
145 |
-
|
146 |
-
|
147 |
-
|
148 |
-
|
149 |
-
|
150 |
-
|
151 |
-
|
152 |
-
|
153 |
-
|
154 |
-
|
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/
|
11 |
-
train_folder = Path('G:/Firestorm/mocap-ai/data/h5/
|
12 |
-
test_folder = Path('G:/Firestorm/mocap-ai/data/h5/
|
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,
|
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.
|
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[
|
52 |
except BaseException as e:
|
53 |
print(e)
|
54 |
return
|
55 |
|
56 |
|
57 |
-
def process_fbx_files(source_folder: Path
|
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(
|
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
|
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
|
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
|
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)
|