Upload pytorch3d_stub/pytorch3d/transforms/__init__.py with huggingface_hub

pytorch3d_stub/pytorch3d/transforms/__init__.py

@@ -53,6 +53,31 @@ def quaternion_invert(quaternion):
     return quaternion * scaling
 
 
+def axis_angle_to_quaternion(axis_angle):
+    """Convert rotations given as axis/angle to quaternions.
+    Args:
+        axis_angle: Rotations given as a vector in axis angle form,
+            as a tensor of shape (..., 3), where the magnitude is
+            the angle turned anticlockwise in radians around the
+            vector's direction.
+    Returns:
+        quaternions with real parts first, as tensor of shape (..., 4).
+    """
+    angles = torch.norm(axis_angle, p=2, dim=-1, keepdim=True)
+    half_angles = angles * 0.5
+    eps = 1e-6
+    small_angles = angles.abs() < eps
+    sin_half = torch.where(small_angles, 0.5 * torch.ones_like(angles), torch.sin(half_angles) / angles)
+    cos_half = torch.where(small_angles, torch.ones_like(angles), torch.cos(half_angles))
+    quaternions = torch.cat([cos_half, axis_angle * sin_half], dim=-1)
+    return quaternions
+
+
+def axis_angle_to_matrix(axis_angle):
+    """Convert rotations given as axis/angle to rotation matrices."""
+    return quaternion_to_matrix(axis_angle_to_quaternion(axis_angle))
+
+
 class Transform3d:
     def __init__(self, dtype=torch.float32, device="cpu", matrix=None):
         if matrix is not None:
@@ -164,3 +189,56 @@ class Transform3d:
         if m.dim() == 2:
             m = m.unsqueeze(0)
         return Transform3d(matrix=m, device=str(self._matrix.device), dtype=self._matrix.dtype)
+
+
+class Rotate(Transform3d):
+    """Transform3d initialized with a rotation matrix."""
+    def __init__(self, R=None, dtype=torch.float32, device="cpu"):
+        super().__init__(dtype=dtype, device=device)
+        if R is not None:
+            if isinstance(R, torch.Tensor):
+                R = R.to(device=device, dtype=dtype)
+                if R.dim() == 2:
+                    R = R.unsqueeze(0)
+                T = torch.eye(4, dtype=dtype, device=device).unsqueeze(0).expand(R.shape[0], -1, -1).clone()
+                T[:, :3, :3] = R
+                self._matrix = T
+                self.device = self._matrix.device
+
+
+class Translate(Transform3d):
+    """Transform3d initialized with a translation."""
+    def __init__(self, x=0.0, y=0.0, z=0.0, dtype=torch.float32, device="cpu"):
+        super().__init__(dtype=dtype, device=device)
+        if isinstance(x, torch.Tensor):
+            t = x.to(device=device, dtype=dtype)
+            if t.dim() == 1:
+                t = t.unsqueeze(0)
+        else:
+            t = torch.tensor([[x, y, z]], dtype=dtype, device=device)
+        T = torch.eye(4, dtype=dtype, device=device).unsqueeze(0).expand(t.shape[0], -1, -1).clone()
+        T[:, :3, 3] = t
+        self._matrix = T
+        self.device = self._matrix.device
+
+
+class Scale(Transform3d):
+    """Transform3d initialized with a scale."""
+    def __init__(self, x=1.0, y=None, z=None, dtype=torch.float32, device="cpu"):
+        super().__init__(dtype=dtype, device=device)
+        if isinstance(x, torch.Tensor):
+            s = x.to(device=device, dtype=dtype)
+            if s.dim() == 0:
+                s = s.expand(3).unsqueeze(0)
+            elif s.dim() == 1 and s.shape[0] == 3:
+                s = s.unsqueeze(0)
+            elif s.dim() == 1 and s.shape[0] == 1:
+                s = s.expand(3).unsqueeze(0)
+        else:
+            if y is None: y = x
+            if z is None: z = x
+            s = torch.tensor([[x, y, z]], dtype=dtype, device=device)
+        S = torch.eye(4, dtype=dtype, device=device).unsqueeze(0).expand(s.shape[0], -1, -1).clone()
+        S[:, 0, 0] = s[:, 0]; S[:, 1, 1] = s[:, 1]; S[:, 2, 2] = s[:, 2]
+        self._matrix = S
+        self.device = self._matrix.device