Handle the case of tensor inputs and squeeze the singleton dimensions

src/skelcast/data/transforms.py

@@ -1,3 +1,5 @@
+import logging
+
 import torch
 from typing import Any, Tuple
 
@@ -50,8 +52,11 @@ def __init__(self, parents: list = None) -> None:
         else:
             self.parents = parents
 
-    def __call__(self, x) -> torch.Tensor:
-        return xyz_to_expmap(x, self.parents)
+    def __call__(self, x) -> torch.Tensor: # Really? Torch tensor? Check it again
+        logging.info(f'The type of x is: {type(x)}')
+        if isinstance(x, torch.Tensor):
+            return xyz_to_expmap(x.squeeze(1).numpy(), self.parents)
+        return xyz_to_expmap(x.squeeze(1), self.parents)
 
 
 @TRANSFORMS.register_module()
@@ -74,8 +79,11 @@ def __init__(self, parents: list = None) -> None:
             self.parents = parents
 
     def __call__(self, x) -> Any:
+        if isinstance(x, torch.Tensor):
+            x = x.squeeze(1).numpy()
         _exps = xyz_to_expmap(x, self.parents)
-        return exps_to_quats(_exps)
+        # Because it returns a tensor we need to convert it to numpy
+        return exps_to_quats(_exps.squeeze(1).numpy())
 
 class Compose:
     def __init__(self, transforms: list) -> None:

src/skelcast/data/utils.py

@@ -1,4 +1,5 @@
 """Copied from https://github.com/qxcv/pose-prediction/blob/master/expmap.py"""
+import torch
 import numpy as np
 
 
@@ -85,7 +86,7 @@ def xyz_to_expmap(xyz_seq, parents):
     root = toposorted[0]
     exp_seq[1:, root] = xyz_seq[1:, root] - xyz_seq[:-1, root]
 
-    return exp_seq
+    return torch.from_numpy(exp_seq).unsqueeze(1)
 
 def exp_to_rotmat(exp):
     """Convert rotation paramterised as exponential map into ordinary 3x3
@@ -133,4 +134,4 @@ def exps_to_quats(exps):
     rv_flat[nonzero_mask, 1:] = nonzero_normed * sines[..., None]
 
     rv_shape = exps.shape[:-1] + (4, )
-    return rv_flat.reshape(rv_shape)
+    return torch.from_numpy(rv_flat.reshape(rv_shape)).unsqueeze(1)