def normalize(cloud: Union[torch.Tensor, PointCloud], inplace: Optional[bool] = True): r"""Returns a normalized pointcloud with zero-mean and unit standard deviation. For batched clouds, each cloud is independently normalized. Args: cloud (torch.Tensor or PointCloud): Input pointcloud to be normalized (shape: :math:`B \times \cdots \times N \times D`, where :math:`B` is the batchsize (optional), :math:`N` is the number of points in the cloud, and :math:`D` is the dimensionality of the cloud. inplace (bool, optional): Bool to make the transform in-place. Returns: (torch.Tensor or PointCloud): The normalized pointcloud. """ if isinstance(cloud, np.ndarray): cloud = torch.from_numpy(cloud) helpers._assert_tensor(cloud) helpers._assert_dim_ge(cloud, 2) if not inplace: cloud = cloud.clone() cloud = (cloud - cloud.mean(-2).unsqueeze(-2))\ / (cloud.std(-2).unsqueeze(-2) + EPS) return cloud
def realign(src: Union[torch.Tensor, PointCloud], tgt: Union[torch.Tensor, PointCloud], inplace: Optional[bool] = True): r""" Aligns a pointcloud `src` to be in the same (axis-aligned) bounding box as that of pointcloud `tgt`. Args: src (torch.Tensor or PointCloud) : Source pointcloud to be transformed (shape: :math:`\cdots \times N \times D`, where :math:`N` is the number of points in the pointcloud, and :math:`D` is the dimensionality of each point in the cloud). tgt (torch.Tensor or PointCloud) : Target pointcloud to which `src`is to be transformed (The `src` cloud is transformed to the axis-aligned bounding box that the target cloud maps to). This cloud must have the same number of dimensions :math:`D` as in the source cloud. (shape: :math:`\cdots \times \cdots \times D`). inplace (bool, optional): Bool to make the transform in-place. Returns: (torch.Tensor): Pointcloud `src` realigned to fit in the (axis-aligned) bounding box of the `tgt` cloud. Example: >>> tgt = torch.rand(1000) >>> src = (tgt * 100) + 3 >>> src_realigned = realign(src, tgt) """ if isinstance(src, PointCloud): src = src.points if isinstance(tgt, PointCloud): tgt = tgt.points helpers._assert_tensor(src) helpers._assert_tensor(tgt) helpers._assert_dim_ge(src, 2) helpers._assert_dim_ge(tgt, 2) helpers._assert_shape_eq(src, tgt.shape, dim=-1) if not inplace: src = src.clone() # Compute the relative scaling factor and scale the src cloud. src_min, _ = src.min(-2) src_max, _ = src.max(-2) tgt_min, _ = tgt.min(-2) tgt_max, _ = tgt.max(-2) src_min = src_min.unsqueeze(-2) src_max = src_max.unsqueeze(-2) tgt_min = tgt_min.unsqueeze(-2) tgt_max = tgt_max.unsqueeze(-2) # Center the pointclouds. src = src - src.mean(-2).unsqueeze(-2) src = ((tgt_max - tgt_min) / (src_max - src_min + EPS)) * src # Undo the centering translation, and return the result. return src + tgt.mean(-2).unsqueeze(-2)
def rotate(cloud: Union[torch.Tensor, PointCloud], rotmat: torch.Tensor, inplace: Optional[bool] = True): """Rotates the the input pointcloud by a rotation matrix. Args: cloud (Tensor or np.array): pointcloud (ndims = 2 or 3) rotmat (Tensor or np.array): rotation matrix (3 x 3, 1 per cloud). inplace (bool, optional): Bool to make the transform in-place. Returns: cloud_rot (Tensor): rotated pointcloud of the same shape as input Shape: - cloud: :math:`(B x N x 3)` (or) :math:`(N x 3)`, where :math:`(B)` is the batchsize, :math:`(N)` is the number of points per cloud, and :math:`(3)` is the dimensionality of each cloud. - rotmat: :math:`(3, 3)` or :math:`(B, 3, 3)`. Example: >>> points = torch.rand(1000,3) >>> r_mat = torch.rand(3,3) >>> points2 = rotate(points, r_mat) """ if isinstance(cloud, np.ndarray): cloud = torch.from_numpy(cloud) if isinstance(cloud, PointCloud): cloud = cloud.points if isinstance(rotmat, np.ndarray): rotmat = torch.from_numpy(rotmat) helpers._assert_tensor(cloud) helpers._assert_tensor(rotmat) helpers._assert_dim_ge(cloud, 2) helpers._assert_dim_ge(rotmat, 2) # Rotation matrix must have last two dimensions of shape 3. helpers._assert_shape_eq(rotmat, (3, 3), dim=-1) helpers._assert_shape_eq(rotmat, (3, 3), dim=-2) if not inplace: cloud = cloud.clone() if rotmat.dim() == 2 and cloud.dim() == 2: cloud = torch.mm(rotmat, cloud.transpose(0, 1)).transpose(0, 1) else: if rotmat.dim() == 2: rotmat = rotmat.expand(cloud.shape[0], 3, 3) cloud = torch.bmm(rotmat, cloud.transpose(1, 2)).transpose(1, 2) return cloud
def scale(cloud: Union[torch.Tensor, PointCloud], scf: Union[float, int, torch.Tensor], inplace: Optional[bool] = True): """Scales the input pointcloud by a scaling factor. Args: cloud (torch.Tensor or kaolin.rep.PointCloud): pointcloud (ndims >= 2). scf (float, int, or torch.Tensor): scaling factor (scalar, or tensor). All elements of scf must be positive. inplace (bool, optional): Bool to make the transform in-place. Returns: (torch.Tensor): scaled pointcloud of the same shape as input. Shape: - cloud: :math:`(B x N x D)` (or) :math:`(N x D)`, where :math:`(B)` is the batchsize, :math:`(N)` is the number of points per cloud, and :math:`(D)` is the dimensionality of each cloud. - scf: :math:`(1)` or :math:`(B)`. Example: >>> points = torch.rand(1000,3) >>> points2 = scale(points, torch.FloatTensor([3])) """ if isinstance(cloud, np.ndarray): cloud = torch.from_numpy(cloud) if isinstance(scf, np.ndarray): scf = torch.from_numpy(scf) if isinstance(cloud, PointCloud): cloud = cloud.points if isinstance(scf, int) or isinstance(scf, float): scf = torch.Tensor([scf]).to(cloud.device) helpers._assert_tensor(cloud) helpers._assert_tensor(scf) helpers._assert_dim_ge(cloud, 2) helpers._assert_gt(scf, 0.) if not inplace: cloud = cloud.clone() return scf * cloud
def __init__(self, points: Optional[torch.Tensor] = None, normals: torch.Tensor = None, device: Optional[str] = 'cpu', copy: Optional[bool] = False): r"""Initialize a PointCloud object, given a tensor of points, and optionally, a tensor representing poincloud normals. Args: pts (torch.Tensor): Points that make up the pointcloud (shape: :math:`... \times N \times D`), where :math:`N` denotes the number of points in the cloud, and :math:`D` denotes the dimensionality of each point. normals (torch.Tensor): Normals for each point in the cloud (shape: :math:`N \times D`, where `D` = 2 or `D` = 3). That is, normals can only be provided for 2D or 3D pointclouds. device (str, Optional): Device to store the pointcloud object on (default: 'cpu'). Must be a valid `torch.device` type. copy (bool, Optional): Whether or not to create a deep copy of the Tensor(s) used to initialze class members. """ if points is None: self.points = None else: helpers._assert_tensor(points) helpers._assert_dim_ge(points, 2) self.points = points.clone() if copy else points self.points = self.points.to(device) if normals is None: self.normals = None else: helpers._assert_tensor(normals) if points.dim() == 2: helpers._assert_shape_eq(normals, (points.shape[-2], 3)) self.normals = normals.clone() if copy else normals self.normals = self.normals.to(device)
def __init__(self, voxels: Optional[torch.Tensor] = None, copy: Optional[bool] = False): r"""Initialize a voxel grid, given a tensor of voxel `features`. Args: voxels (torch.Tensor, optional): Tensor containing voxel features (shape: Any shape that has >= 3 dims). copy (bool, optional): Whether or not to create a deep copy of the Tensor(s) used to initialize class member(s). Note: By default, the created VoxelGrid object stores a reference to the input `voxels` tensor. To create a deep copy of the voxels, set the `copy` argument to `True`. """ super(VoxelGrid, self).__init__() if voxels is None: self.voxels = None else: helpers._assert_tensor(voxels) helpers._assert_dim_ge(voxels, 3) self.voxels = voxels.clone() if copy else voxels
def sample_triangle_mesh(vertices: torch.Tensor, faces: torch.Tensor, num_samples: int, eps: float = 1e-10): r""" Uniformly samples the surface of a mesh. Args: vertices (torch.Tensor): Vertices of the mesh (shape: :math:`N \times 3`, where :math:`N` is the number of vertices) faces (torch.LongTensor): Faces of the mesh (shape: :math:`F \times 3`, where :math:`F` is the number of faces). num_samples (int): Number of points to sample eps (float): A small number to prevent division by zero for small surface areas. Returns: (torch.Tensor): Uniformly sampled points from the triangle mesh. Example: >>> points = sample_triangle_mesh(vertices, faces, 10) >>> points tensor([[ 0.0293, 0.2179, 0.2168], [ 0.2003, -0.3367, 0.2187], [ 0.2152, -0.0943, 0.1907], [-0.1852, 0.1686, -0.0522], [-0.2167, 0.3171, 0.0737], [ 0.2219, -0.0289, 0.1531], [ 0.2217, -0.0115, 0.1247], [-0.1400, 0.0364, -0.1618], [ 0.0658, -0.0310, -0.2198], [ 0.1926, -0.1867, -0.2153]]) """ helpers._assert_tensor(vertices) helpers._assert_tensor(faces) helpers._assert_dim_ge(vertices, 2) helpers._assert_dim_ge(faces, 2) # We want the last dimension of vertices to be of shape 3. helpers._assert_shape_eq(vertices, (-1, 3), dim=-1) dist_uni = torch.distributions.Uniform(torch.zeros((1,), device=vertices.device), 1.) # calculate area of each face x1, x2, x3 = torch.split(torch.index_select( vertices, 0, faces[:, 0]) - torch.index_select( vertices, 0, faces[:, 1]), 1, dim=1) y1, y2, y3 = torch.split(torch.index_select( vertices, 0, faces[:, 1]) - torch.index_select( vertices, 0, faces[:, 2]), 1, dim=1) a = (x2 * y3 - x3 * y2) ** 2 b = (x3 * y1 - x1 * y3) ** 2 c = (x1 * y2 - x2 * y1) ** 2 Areas = torch.sqrt(a + b + c) / 2 # percentage of each face w.r.t. full surface area Areas = Areas / (torch.sum(Areas) + eps) # define descrete distribution w.r.t. face area ratios caluclated cat_dist = torch.distributions.Categorical(Areas.view(-1)) face_choices = cat_dist.sample([num_samples]) # from each face sample a point select_faces = faces[face_choices] xs = torch.index_select(vertices, 0, select_faces[:, 0]) ys = torch.index_select(vertices, 0, select_faces[:, 1]) zs = torch.index_select(vertices, 0, select_faces[:, 2]) u = torch.sqrt(dist_uni.sample([num_samples])) v = dist_uni.sample([num_samples]) points = (1 - u) * xs + (u * (1 - v)) * ys + u * v * zs return points