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

pytorch3d_stub/pytorch3d/transforms/__init__.py

@@ -55,38 +55,60 @@ def quaternion_invert(quaternion):
 
 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(device=device, dtype=dtype)
+            self.device = self._matrix.device
+            self.dtype = self._matrix.dtype
         else:
             self._matrix = torch.eye(4, dtype=dtype, device=device).unsqueeze(0)
+            self.device = self._matrix.device
+            self.dtype = dtype
+
     def get_matrix(self):
         return self._matrix
+
     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)
+            om = o.get_matrix().to(device=m.device, dtype=m.dtype)
+            if om.shape[0] != m.shape[0]:
+                if m.shape[0] == 1:
+                    m = m.expand(om.shape[0], -1, -1)
+                elif om.shape[0] == 1:
+                    om = om.expand(m.shape[0], -1, -1)
+            m = m @ om
+        return Transform3d(matrix=m, device=str(m.device), dtype=m.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)
+        mat = self._matrix
+        points = points.to(device=mat.device, dtype=mat.dtype)
+        ones = torch.ones(*points.shape[:-1], 1, dtype=mat.dtype, device=mat.device)
         pts4 = torch.cat([points, ones], dim=-1)
-        out = torch.bmm(pts4, self._matrix.expand(pts4.shape[0], -1, -1).transpose(-2, -1))
+        m = mat.expand(pts4.shape[0], -1, -1)
+        out = torch.bmm(pts4, m.transpose(-2, -1))
         return out[..., :3]
+
     def translate(self, x, y=None, z=None):
+        dev = self._matrix.device
+        dt = self._matrix.dtype
         if isinstance(x, torch.Tensor) and x.dim() >= 1 and x.shape[-1] == 3:
-            t = x
+            t = x.to(device=dev, dtype=dt)
         else:
-            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 = torch.tensor([[x, y, z]], dtype=dt, device=dev)
+        if t.dim() == 1:
+            t = t.unsqueeze(0)
+        T = torch.eye(4, dtype=dt, device=dev).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)
+        return Transform3d(matrix=new_m, device=str(dev), dtype=dt)
+
     def scale(self, x, y=None, z=None):
+        dev = self._matrix.device
+        dt = self._matrix.dtype
         if isinstance(x, torch.Tensor):
+            x = x.to(device=dev, dtype=dt)
             if x.dim() == 0:
                 s = x.expand(3)
             elif x.shape[-1] == 3:
@@ -96,23 +118,49 @@ class Transform3d:
         else:
             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 = torch.tensor([x, y, z], dtype=dt, device=dev)
+        if s.dim() == 1:
+            s = s.unsqueeze(0)
+        S = torch.eye(4, dtype=dt, device=dev).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)
+        return Transform3d(matrix=new_m, device=str(dev), dtype=dt)
+
     def to(self, device=None, dtype=None):
-        if device is not None: self.device = device
-        if dtype is not None: self.dtype = dtype
         self._matrix = self._matrix.to(device=device, dtype=dtype)
+        self.device = self._matrix.device
+        if dtype is not None:
+            self.dtype = dtype
         return self
+
     def inverse(self):
         inv_m = torch.inverse(self._matrix)
-        return Transform3d(matrix=inv_m, device=self.device, dtype=self.dtype)
+        return Transform3d(matrix=inv_m, device=str(self._matrix.device), dtype=self._matrix.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()
+        dev = self._matrix.device
+        dt = self._matrix.dtype
+        R = R.to(device=dev, dtype=dt)
+        if R.dim() == 2:
+            R = R.unsqueeze(0)
+        T = torch.eye(4, dtype=dt, device=dev).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)
+        return Transform3d(matrix=new_m, device=str(dev), dtype=dt)
+
+    def stack(self, *others):
+        matrices = [self._matrix] + [o.get_matrix().to(device=self._matrix.device, dtype=self._matrix.dtype) for o in others]
+        stacked = torch.cat(matrices, dim=0)
+        return Transform3d(matrix=stacked, device=str(self._matrix.device), dtype=self._matrix.dtype)
+
+    def clone(self):
+        return Transform3d(matrix=self._matrix.clone(), device=str(self._matrix.device), dtype=self._matrix.dtype)
+
+    def __len__(self):
+        return self._matrix.shape[0]
+
+    def __getitem__(self, index):
+        m = self._matrix[index]
+        if m.dim() == 2:
+            m = m.unsqueeze(0)
+        return Transform3d(matrix=m, device=str(self._matrix.device), dtype=self._matrix.dtype)