Merge branch 'main' of https://github.com/moverseai/moai

moai/monads/geometry/project.py

@@ -1,6 +1,27 @@
 import torch
 
 
+## Project2D
+class Project2D(torch.nn.Module):
+    """
+    returns tensor of size [1,V,K,2]
+    """
+
+    def __init__(self, epsilon: float = 1e-7) -> None:
+        # ADD EPSILON
+        super(Project2D, self).__init__()
+        self.epsilon = epsilon
+
+    def forward(
+        self, intrisics: torch.Tensor, points: torch.Tensor  # [B,V,3,3]
+    ):  # [B,V,J,3]
+
+        kp_uv = torch.einsum("bvmn,bvjn->bvjm", intrisics, points)
+        kp_uv = kp_uv[..., :2] / (kp_uv[..., -1:] + self.epsilon)
+
+        return kp_uv
+
+
 class Projection(torch.nn.Module):
 
     __XYZ_AT__ = ["channel", "row"]

moai/monads/geometry/transform_points.py

@@ -3,6 +3,24 @@
 import torch
 
 
+class Transform3D(torch.nn.Module):
+    """
+    returns  tensor of size  [1,V,K,3]
+    """
+
+    def __init__(self) -> None:
+        super(Transform3D, self).__init__()
+
+    def forward(
+        self,
+        transforms: torch.Tensor,  # [B,V,4,4]
+        points: torch.Tensor,  # [B,J,3]
+    ):
+        R = transforms[..., :3, :3]
+        T = torch.unsqueeze(transforms[..., :3, 3], dim=-2)
+        return torch.einsum("bvmn,bjn->bvjm", R, points) + T
+
+
 class Transformation(torch.nn.Module):
     __XYZ_AT__ = ["channel", "row"]