r"""Create a `ME.SparseTensor` coordinates from a sequence of coordinates Given a list of either numpy or pytorch tensor coordinates, return the batched coordinates suitable for `ME.SparseTensor`. Args: :attr:`coords` (a sequence of `torch.Tensor` or `numpy.ndarray`): a
(coords, dtype=torch.int32, device=None)
| 47 | |
| 48 | |
| 49 | def batched_coordinates(coords, dtype=torch.int32, device=None): |
| 50 | r"""Create a `ME.SparseTensor` coordinates from a sequence of coordinates |
| 51 | |
| 52 | Given a list of either numpy or pytorch tensor coordinates, return the |
| 53 | batched coordinates suitable for `ME.SparseTensor`. |
| 54 | |
| 55 | Args: |
| 56 | :attr:`coords` (a sequence of `torch.Tensor` or `numpy.ndarray`): a |
| 57 | list of coordinates. |
| 58 | |
| 59 | :attr:`dtype`: torch data type of the return tensor. torch.int32 by default. |
| 60 | |
| 61 | Returns: |
| 62 | :attr:`batched_coordindates` (`torch.Tensor`): a batched coordinates. |
| 63 | |
| 64 | .. warning:: |
| 65 | |
| 66 | From v0.4, the batch index will be prepended before all coordinates. |
| 67 | |
| 68 | """ |
| 69 | assert isinstance( |
| 70 | coords, collections.abc.Sequence |
| 71 | ), "The coordinates must be a sequence." |
| 72 | assert np.array( |
| 73 | [cs.ndim == 2 for cs in coords] |
| 74 | ).all(), "All coordinates must be in a 2D array." |
| 75 | D = np.unique(np.array([cs.shape[1] for cs in coords])) |
| 76 | assert len(D) == 1, f"Dimension of the array mismatch. All dimensions: {D}" |
| 77 | D = D[0] |
| 78 | if device is None: |
| 79 | if isinstance(coords, torch.Tensor): |
| 80 | device = coords[0].device |
| 81 | else: |
| 82 | device = "cpu" |
| 83 | assert dtype in [ |
| 84 | torch.int32, |
| 85 | torch.float32, |
| 86 | ], "Only torch.int32, torch.float32 supported for coordinates." |
| 87 | |
| 88 | # Create a batched coordinates |
| 89 | N = np.array([len(cs) for cs in coords]).sum() |
| 90 | bcoords = torch.zeros((N, D + 1), dtype=dtype, device=device) # uninitialized |
| 91 | |
| 92 | s = 0 |
| 93 | for b, cs in enumerate(coords): |
| 94 | if dtype == torch.int32: |
| 95 | if isinstance(cs, np.ndarray): |
| 96 | cs = torch.from_numpy(np.floor(cs)) |
| 97 | elif not ( |
| 98 | isinstance(cs, torch.IntTensor) or isinstance(cs, torch.LongTensor) |
| 99 | ): |
| 100 | cs = cs.floor() |
| 101 | |
| 102 | cs = cs.int() |
| 103 | else: |
| 104 | if isinstance(cs, np.ndarray): |
| 105 | cs = torch.from_numpy(cs) |
| 106 |
no outgoing calls