[Fix] Fix wrong displacement shape mentioned by issue (#112)

* Fix wrong displacement shape mentioned by issue

Author: LazyBusyYang

tests/data_structure/body_model/test_smplxd_data.py

@@ -30,7 +30,7 @@ def test_new():
         transl=np.zeros(shape=(2, 3)),
         betas=np.zeros(shape=(2, 10)),
         mask=np.ones(shape=(2, )),
-        displacement=np.zeros(shape=(2, 10475)),
+        displacement=np.zeros(shape=(2, 10475, 3)),
         logger='root')
     assert smplxd_data['betas'][0, 0] == 0
     assert smplxd_data.get_expression().shape == (2, 10)
@@ -100,7 +100,7 @@ def test_setitem():
     smplxd_data['transl'] = np.zeros(shape=[2, 3])
     smplxd_data['fullpose'] = np.zeros(shape=[2, 55, 3])
     smplxd_data['gender'] = 'neutral'
-    smplxd_data['displacement'] = np.zeros(shape=(2, 10475))
+    smplxd_data['displacement'] = np.zeros(shape=(2, 10475, 3))
     # set arbitrary key
     smplxd_data['frame_number'] = 1000
 

tests/model/registrant/test_smplifyxd.py

@@ -98,7 +98,7 @@ def test_smplifyxd_keypoints3d():
             np_v = v.detach().cpu().numpy()
             assert not np.any(np.isnan(np_v)), f'{k} fails.'
     smplxd_data.from_param_dict(smplifyxd_output)
-    assert len(smplxd_data.get_displacement().shape) == 2
+    assert len(smplxd_data.get_displacement().shape) == 3
     result_path = os.path.join(output_dir, 'smplxd_result.npz')
     smplxd_data.dump(result_path)
     # test not use_one_betas_per_video and return values

xrmocap/data_structure/body_model/smplxd_data.py

@@ -54,7 +54,7 @@ def __init__(self,
                 zero-tensor in shape [frame_num, 10] will be created.
             displacement (Union[np.ndarray, torch.Tensor, None], optional):
                 A tensor or ndarray for displacement,
-                in shape [frame_num, NUM_VERTS].
+                in shape [frame_num, NUM_VERTS, 3].
                 Defaults to None,
                 zero-tensor in shape [frame_num, NUM_VERTS] will be created.
             mask (Union[np.ndarray, torch.Tensor, None], optional):
@@ -77,7 +77,7 @@ def __init__(self,
             logger=logger)
         if displacement is None and 'displacement' not in self:
             displacement = np.zeros(
-                shape=(self.get_batch_size(), self.__class__.NUM_VERTS))
+                shape=(self.get_batch_size(), self.__class__.NUM_VERTS, 3))
         if displacement is not None:
             self.set_displacement(displacement)
 
@@ -117,7 +117,7 @@ def set_displacement(
         Args:
             displacement (Union[np.ndarray, torch.Tensor]):
                 Displacement parameters in ndarray or tensor,
-                in shape [batch_size, NUM_VERTS].
+                in shape [batch_size, NUM_VERTS, 3].
 
         Raises:
             TypeError: Type of displacement is not correct.
@@ -128,18 +128,18 @@ def set_displacement(
             self.logger.error('Type of displacement is not correct.\n' +
                               f'Type: {type(displacement)}.')
             raise TypeError
-        if len(displacement.shape) == 1:
-            displacement = displacement[np.newaxis, ...]
-        displacement_dim = displacement.shape[-1]
-        displacement_np = displacement.reshape(-1, displacement_dim)
-        dict.__setitem__(self, 'displacement', displacement_np)
+        if len(displacement.shape) < 3:
+            self.logger.error('Shape of displacement is not correct.\n' +
+                              f'Shape: {type(displacement.shape)}.')
+            raise ValueError
+        dict.__setitem__(self, 'displacement', displacement)
 
     def get_displacement(self) -> np.ndarray:
         """Get displacement.
 
         Returns:
             ndarray:
-                Displacement in shape [batch_size, NUM_VERTS].
+                Displacement in shape [batch_size, NUM_VERTS, 3].
         """
         displacement = self.__getitem__('displacement')
         return displacement