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

pytorch3d_stub/pytorch3d/transforms/__init__.py

@@ -1,8 +1,8 @@
-"""Minimal pytorch3d.transforms stub – quaternion operations only."""
+"""pytorch3d.transforms stub – pure PyTorch implementations."""
 import torch
+import math
 
 def quaternion_multiply(q1: torch.Tensor, q2: torch.Tensor) -> torch.Tensor:
-    """Hamilton product of two quaternions (w, x, y, z)."""
     w1, x1, y1, z1 = q1.unbind(-1)
     w2, x2, y2, z2 = q2.unbind(-1)
     return torch.stack([
@@ -13,6 +13,118 @@ def quaternion_multiply(q1: torch.Tensor, q2: torch.Tensor) -> torch.Tensor:
     ], dim=-1)
 
 def quaternion_invert(q: torch.Tensor) -> torch.Tensor:
-    """Invert a quaternion (w, x, y, z) -> (w, -x, -y, -z) / ||q||^2."""
     scaling = torch.tensor([1, -1, -1, -1], device=q.device, dtype=q.dtype)
     return q * scaling / (q * q).sum(dim=-1, keepdim=True).clamp(min=1e-10)
+
+def quaternion_to_matrix(quaternions: torch.Tensor) -> torch.Tensor:
+    r, i, j, k = quaternions.unbind(-1)
+    two_s = 2.0 / (quaternions * quaternions).sum(-1)
+    return torch.stack([
+        1 - two_s * (j*j + k*k), two_s * (i*j - k*r), two_s * (i*k + j*r),
+        two_s * (i*j + k*r), 1 - two_s * (i*i + k*k), two_s * (j*k - i*r),
+        two_s * (i*k - j*r), two_s * (j*k + i*r), 1 - two_s * (i*i + j*j),
+    ], dim=-1).reshape(quaternions.shape[:-1] + (3, 3))
+
+def axis_angle_to_quaternion(axis_angle: torch.Tensor) -> torch.Tensor:
+    angles = torch.norm(axis_angle, p=2, dim=-1, keepdim=True)
+    half_angles = 0.5 * angles
+    eps = 1e-6
+    small = (angles.abs() < eps).squeeze(-1)
+    sin_half = torch.sin(half_angles)
+    cos_half = torch.cos(half_angles)
+    k = torch.where(angles > eps, sin_half / angles, 0.5 * torch.ones_like(angles))
+    w = cos_half.squeeze(-1)
+    xyz = axis_angle * k
+    return torch.cat([w.unsqueeze(-1), xyz], dim=-1)
+
+
+class Transform3d:
+    def __init__(self, dtype=torch.float32, device="cpu", matrix=None):
+        if matrix is not None:
+            self._matrix = matrix.to(dtype=dtype, device=device)
+        else:
+            self._matrix = torch.eye(4, dtype=dtype, device=device).unsqueeze(0)
+
+    def get_matrix(self) -> torch.Tensor:
+        return self._matrix
+
+    def _compose(self, other_matrix: torch.Tensor) -> "Transform3d":
+        new_matrix = torch.bmm(self._matrix, other_matrix)
+        return Transform3d(matrix=new_matrix)
+
+    def scale(self, s) -> "Transform3d":
+        if isinstance(s, (int, float)):
+            s = torch.tensor([[s, s, s]], dtype=self._matrix.dtype, device=self._matrix.device)
+        elif isinstance(s, torch.Tensor):
+            if s.dim() == 1:
+                s = s.unsqueeze(0)
+        S = torch.eye(4, dtype=self._matrix.dtype, device=self._matrix.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]
+        return self._compose(S)
+
+    def rotate(self, R) -> "Transform3d":
+        if isinstance(R, torch.Tensor):
+            if R.dim() == 2:
+                R = R.unsqueeze(0)
+            M = torch.eye(4, dtype=self._matrix.dtype, device=self._matrix.device).unsqueeze(0).expand(R.shape[0], -1, -1).clone()
+            M[:, :3, :3] = R
+            return self._compose(M)
+        return self
+
+    def translate(self, t) -> "Transform3d":
+        if isinstance(t, torch.Tensor):
+            if t.dim() == 1:
+                t = t.unsqueeze(0)
+            M = torch.eye(4, dtype=self._matrix.dtype, device=self._matrix.device).unsqueeze(0).expand(t.shape[0], -1, -1).clone()
+            M[:, 3, :3] = t
+            return self._compose(M)
+        return self
+
+    def compose(self, other: "Transform3d") -> "Transform3d":
+        return self._compose(other._matrix)
+
+
+class Rotate(Transform3d):
+    def __init__(self, R, dtype=torch.float32, device="cpu"):
+        if isinstance(R, torch.Tensor):
+            if R.dim() == 2:
+                R = R.unsqueeze(0)
+            M = torch.eye(4, dtype=dtype, device=device).unsqueeze(0).expand(R.shape[0], -1, -1).clone()
+            M[:, :3, :3] = R.to(dtype=dtype, device=device)
+            super().__init__(dtype=dtype, device=device, matrix=M)
+        else:
+            super().__init__(dtype=dtype, device=device)
+
+class Translate(Transform3d):
+    def __init__(self, *args, dtype=torch.float32, device="cpu"):
+        if len(args) == 1 and isinstance(args[0], torch.Tensor):
+            t = args[0]
+            if t.dim() == 1:
+                t = t.unsqueeze(0)
+        elif len(args) == 3:
+            t = torch.tensor([[args[0], args[1], args[2]]], dtype=dtype, device=device)
+        else:
+            super().__init__(dtype=dtype, device=device)
+            return
+        M = torch.eye(4, dtype=dtype, device=device).unsqueeze(0).expand(t.shape[0], -1, -1).clone()
+        M[:, 3, :3] = t.to(dtype=dtype, device=device)
+        super().__init__(dtype=dtype, device=device, matrix=M)
+
+class Scale(Transform3d):
+    def __init__(self, s, dtype=torch.float32, device="cpu"):
+        if isinstance(s, (int, float)):
+            s = torch.tensor([[s, s, s]], dtype=dtype, device=device)
+        elif isinstance(s, torch.Tensor):
+            if s.dim() == 0:
+                s = s.reshape(1, 1).expand(1, 3)
+            elif s.dim() == 1 and s.shape[0] == 3:
+                s = s.unsqueeze(0)
+            elif s.dim() == 1 and s.shape[0] != 3:
+                s = s.unsqueeze(-1).expand(-1, 3)
+        M = torch.eye(4, dtype=dtype, device=device).unsqueeze(0).expand(s.shape[0], -1, -1).clone()
+        M[:, 0, 0] = s[:, 0]
+        M[:, 1, 1] = s[:, 1]
+        M[:, 2, 2] = s[:, 2]
+        super().__init__(dtype=dtype, device=device, matrix=M)