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

pytorch3d_stub/pytorch3d/transforms/__init__.py

@@ -1,173 +1,110 @@
-"""pytorch3d.transforms stub – pure PyTorch implementations."""
+"""pytorch3d.transforms stub with catch-all."""
 import torch
-import math
+import warnings
 
-def quaternion_multiply(q1: torch.Tensor, q2: torch.Tensor) -> torch.Tensor:
-    w1, x1, y1, z1 = q1.unbind(-1)
-    w2, x2, y2, z2 = q2.unbind(-1)
-    return torch.stack([
-        w1*w2 - x1*x2 - y1*y2 - z1*z2,
-        w1*x2 + x1*w2 + y1*z2 - z1*y2,
-        w1*y2 - x1*z2 + y1*w2 + z1*x2,
-        w1*z2 + x1*y2 - y1*x2 + z1*w2,
-    ], dim=-1)
 
-def quaternion_invert(q: torch.Tensor) -> torch.Tensor:
-    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 __getattr__(name):
+    if name.startswith("_"):
+        raise AttributeError(name)
+    warnings.warn(f"pytorch3d.transforms stub: {name} not implemented", stacklevel=2)
+    def _dummy(*a, **kw): return None
+    _dummy.__name__ = name
+    return _dummy
 
-def quaternion_to_matrix(quaternions: torch.Tensor) -> torch.Tensor:
-    r, i, j, k = quaternions.unbind(-1)
+
+def quaternion_to_matrix(quaternions):
+    r, i, j, k = quaternions[..., 0], quaternions[..., 1], quaternions[..., 2], quaternions[..., 3]
     two_s = 2.0 / (quaternions * quaternions).sum(-1)
-    return torch.stack([
+    o = 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))
+    ], -1)
+    return o.reshape(quaternions.shape[:-1] + (3, 3))
+
+
+def matrix_to_quaternion(matrix):
+    if matrix.shape[-2:] != (3, 3):
+        raise ValueError(f"Invalid rotation matrix shape {matrix.shape}")
+    batch_dim = matrix.shape[:-2]
+    m00, m01, m02 = matrix[..., 0, 0], matrix[..., 0, 1], matrix[..., 0, 2]
+    m10, m11, m12 = matrix[..., 1, 0], matrix[..., 1, 1], matrix[..., 1, 2]
+    m20, m21, m22 = matrix[..., 2, 0], matrix[..., 2, 1], matrix[..., 2, 2]
+    trace = m00 + m11 + m22
+    qw = torch.sqrt(torch.clamp(trace + 1, min=1e-8)) / 2
+    qx = (m21 - m12) / (4 * qw + 1e-8)
+    qy = (m02 - m20) / (4 * qw + 1e-8)
+    qz = (m10 - m01) / (4 * qw + 1e-8)
+    return torch.stack([qw, qx, qy, qz], dim=-1)
+
 
-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)
+def quaternion_multiply(a, b):
+    aw, ax, ay, az = a[..., 0], a[..., 1], a[..., 2], a[..., 3]
+    bw, bx, by, bz = b[..., 0], b[..., 1], b[..., 2], b[..., 3]
+    ow = aw*bw - ax*bx - ay*by - az*bz
+    ox = aw*bx + ax*bw + ay*bz - az*by
+    oy = aw*by - ax*bz + ay*bw + az*bx
+    oz = aw*bz + ax*by - ay*bx + az*bw
+    return torch.stack([ow, ox, oy, oz], dim=-1)
+
+
+def quaternion_invert(quaternion):
+    scaling = torch.tensor([1, -1, -1, -1], dtype=quaternion.dtype, device=quaternion.device)
+    return quaternion * scaling
 
 
 class Transform3d:
     def __init__(self, dtype=torch.float32, device="cpu", matrix=None):
+        self.device = device
+        self.dtype = dtype
         if matrix is not None:
-            self._matrix = matrix.to(dtype=dtype, device=device)
+            self._matrix = matrix.to(device=device, dtype=dtype)
         else:
             self._matrix = torch.eye(4, dtype=dtype, device=device).unsqueeze(0)
-
-    def get_matrix(self) -> torch.Tensor:
+    def get_matrix(self):
         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)
+    def compose(self, *others):
+        m = self._matrix
+        for o in others:
+            m = m @ o.get_matrix()
+        return Transform3d(matrix=m, device=self.device, dtype=self.dtype)
+    def transform_points(self, points):
+        if points.dim() == 2:
+            points = points.unsqueeze(0)
+        ones = torch.ones(*points.shape[:-1], 1, dtype=points.dtype, device=points.device)
+        pts4 = torch.cat([points, ones], dim=-1)
+        out = torch.bmm(pts4, self._matrix.expand(pts4.shape[0], -1, -1).transpose(-2, -1))
+        return out[..., :3]
+    def translate(self, x, y=None, z=None):
+        if isinstance(x, torch.Tensor) and x.dim() >= 1 and x.shape[-1] == 3:
+            t = x
         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)
+            t = torch.tensor([[x, y, z]], dtype=self.dtype, device=self.device)
+        if t.dim() == 1: t = t.unsqueeze(0)
+        T = torch.eye(4, dtype=self.dtype, device=self.device).unsqueeze(0).expand(t.shape[0], -1, -1).clone()
+        T[:, :3, 3] = t
+        new_m = self._matrix @ T
+        return Transform3d(matrix=new_m, device=self.device, dtype=self.dtype)
+    def scale(self, x, y=None, z=None):
+        if isinstance(x, torch.Tensor):
+            if x.dim() == 0:
+                s = x.expand(3)
+            elif x.shape[-1] == 3:
+                s = x.squeeze()
+            else:
+                s = x.expand(3)
         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)
-
-def matrix_to_quaternion(matrix: torch.Tensor) -> torch.Tensor:
-    """Convert rotation matrix to quaternion (w, x, y, z)."""
-    if matrix.dim() == 2:
-        matrix = matrix.unsqueeze(0)
-    batch_dim = matrix.shape[:-2]
-    m = matrix.reshape(-1, 3, 3)
-    
-    trace = m[:, 0, 0] + m[:, 1, 1] + m[:, 2, 2]
-    q = torch.zeros(m.shape[0], 4, dtype=m.dtype, device=m.device)
-    
-    s = torch.sqrt((trace + 1.0).clamp(min=1e-10)) * 2  # s = 4*w
-    q[:, 0] = 0.25 * s
-    q[:, 1] = (m[:, 2, 1] - m[:, 1, 2]) / s.clamp(min=1e-10)
-    q[:, 2] = (m[:, 0, 2] - m[:, 2, 0]) / s.clamp(min=1e-10)
-    q[:, 3] = (m[:, 1, 0] - m[:, 0, 1]) / s.clamp(min=1e-10)
-    
-    # Handle degenerate cases
-    mask1 = (m[:, 0, 0] > m[:, 1, 1]) & (m[:, 0, 0] > m[:, 2, 2]) & (trace <= 0)
-    if mask1.any():
-        s1 = torch.sqrt((1.0 + m[mask1, 0, 0] - m[mask1, 1, 1] - m[mask1, 2, 2]).clamp(min=1e-10)) * 2
-        q[mask1, 0] = (m[mask1, 2, 1] - m[mask1, 1, 2]) / s1.clamp(min=1e-10)
-        q[mask1, 1] = 0.25 * s1
-        q[mask1, 2] = (m[mask1, 0, 1] + m[mask1, 1, 0]) / s1.clamp(min=1e-10)
-        q[mask1, 3] = (m[mask1, 0, 2] + m[mask1, 2, 0]) / s1.clamp(min=1e-10)
-    
-    mask2 = (m[:, 1, 1] > m[:, 2, 2]) & ~mask1 & (trace <= 0)
-    if mask2.any():
-        s2 = torch.sqrt((1.0 + m[mask2, 1, 1] - m[mask2, 0, 0] - m[mask2, 2, 2]).clamp(min=1e-10)) * 2
-        q[mask2, 0] = (m[mask2, 0, 2] - m[mask2, 2, 0]) / s2.clamp(min=1e-10)
-        q[mask2, 1] = (m[mask2, 0, 1] + m[mask2, 1, 0]) / s2.clamp(min=1e-10)
-        q[mask2, 2] = 0.25 * s2
-        q[mask2, 3] = (m[mask2, 1, 2] + m[mask2, 2, 1]) / s2.clamp(min=1e-10)
-    
-    mask3 = ~mask1 & ~mask2 & (trace <= 0)
-    if mask3.any():
-        s3 = torch.sqrt((1.0 + m[mask3, 2, 2] - m[mask3, 0, 0] - m[mask3, 1, 1]).clamp(min=1e-10)) * 2
-        q[mask3, 0] = (m[mask3, 1, 0] - m[mask3, 0, 1]) / s3.clamp(min=1e-10)
-        q[mask3, 1] = (m[mask3, 0, 2] + m[mask3, 2, 0]) / s3.clamp(min=1e-10)
-        q[mask3, 2] = (m[mask3, 1, 2] + m[mask3, 2, 1]) / s3.clamp(min=1e-10)
-        q[mask3, 3] = 0.25 * s3
-    
-    return q.reshape(batch_dim + (4,))
+            if y is None: y = x
+            if z is None: z = x
+            s = torch.tensor([x, y, z], dtype=self.dtype, device=self.device)
+        if s.dim() == 1: s = s.unsqueeze(0)
+        S = torch.eye(4, dtype=self.dtype, device=self.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]
+        new_m = self._matrix @ S
+        return Transform3d(matrix=new_m, device=self.device, dtype=self.dtype)
+    def rotate(self, R):
+        if R.dim() == 2: R = R.unsqueeze(0)
+        T = torch.eye(4, dtype=self.dtype, device=self.device).unsqueeze(0).expand(R.shape[0], -1, -1).clone()
+        T[:, :3, :3] = R
+        new_m = self._matrix @ T
+        return Transform3d(matrix=new_m, device=self.device, dtype=self.dtype)