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

pytorch3d_stub/pytorch3d/transforms/__init__.py

@@ -128,3 +128,46 @@ class Scale(Transform3d):
         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,))