def test_smoke(device='cpu'): g = kal.models.Voxel3DIWGAN.Voxel3DIWGenerator() d = kal.models.Voxel3DIWGAN.Voxel3DIWDiscriminator() g = g.to(device) d = d.to(device) x = torch.randn(4, 200) x_ = g(x) x_ = d(x_) helpers._assert_shape_eq(x, x_.shape, dim=0)
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 test_smoke(): model = kal.models.PointNet2.PointNet2Classifier( in_features=4, num_classes=13, batchnorm=True, ).cuda() x = torch.randn(2, 128, 7).cuda() out = model(x) helpers._assert_shape_eq(out, (2, 13)) model = kal.models.PointNet2.PointNet2Segmenter( in_features=4, num_classes=13, batchnorm=True, ).cuda() x = torch.randn(2, 128, 7).cuda() out = model(x) helpers._assert_shape_eq(out, (2, 128, 13))
def test_downsample_voxelgrid(device='cpu'): voxel = torch.ones([32, 32, 32]).to(device) down = kal.transforms.DownsampleVoxelGrid([2,2,2], inplace=False) helpers._assert_shape_eq(down(voxel), (16,16,16)) down = kal.transforms.DownsampleVoxelGrid([3,3,3], inplace=False) helpers._assert_shape_eq(down(voxel), (10,10,10)) down = kal.transforms.DownsampleVoxelGrid([3,2,1], inplace=False) helpers._assert_shape_eq(down(voxel), (10,16,32))
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 rotate(mesh: Type[Mesh], rotmat: torch.Tensor, inplace: Optional[bool] = True): r"""Rotate a mesh given a 3 x 3 rotation matrix. Args: mesh (Mesh): Mesh to be rotated. rotmat (torch.Tensor): Rotation matrix (shape: :math:`3 \times 3`). inplace (bool, optional): Bool to make this operation in-place. Returns: (Mesh): Rotatted mesh. """ if not isinstance(mesh, Mesh): raise TypeError('Input mesh must be of type Mesh. ' 'Got {0} instead.'.format(type(mesh))) if not inplace: mesh = mesh.clone() helpers._assert_tensor(rotmat) helpers._assert_shape_eq(rotmat, (3, 3)) mesh.vertices = torch.matmul(rotmat, mesh.vertices.t()).t() return mesh
def test_smoke(device='cpu'): model = kal.models.PointNet.PointNetClassifier( in_channels=3, feat_size=1024, num_classes=13, batchnorm=True, ) x = torch.randn(2, 128, 3) out = model(x) helpers._assert_shape_eq(out, (2, 13)) model = kal.models.PointNet.PointNetClassifier( in_channels=3, feat_size=1024, num_classes=13, batchnorm=True, transposed_input=True, ) x = torch.randn(2, 3, 128) out = model(x) helpers._assert_shape_eq(out, (2, 13)) model = kal.models.PointNet.PointNetSegmenter( in_channels=3, feat_size=1024, num_classes=13, batchnorm=True, ) x = torch.randn(2, 128, 3) out = model(x) helpers._assert_shape_eq(out, (2, 128, 13)) model = kal.models.PointNet.PointNetSegmenter( in_channels=3, feat_size=1024, num_classes=13, batchnorm=True, transposed_input=True, ) x = torch.randn(2, 3, 128) out = model(x) helpers._assert_shape_eq(out, (2, 128, 13))
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
def test_triangle_mesh_to_sdf(device='cpu'): mesh = TriangleMesh.from_obj('tests/model.obj') mesh.to(device) mesh2sdf = kal.transforms.TriangleMeshToSDF(100) helpers._assert_shape_eq(mesh2sdf(mesh), (100,), dim=-1)
def test_triangle_mesh_to_voxelgrid(device='cpu'): mesh = TriangleMesh.from_obj('tests/model.obj') mesh.to(device) mesh2voxel = kal.transforms.TriangleMeshToVoxelGrid(32) helpers._assert_shape_eq(mesh2voxel(mesh), (32, 32, 32))
def test_triangle_mesh_to_pointcloud(device='cpu'): mesh = TriangleMesh.from_obj('tests/model.obj') mesh.to(device) mesh2cloud = kal.transforms.TriangleMeshToPointCloud(10000) pts = mesh2cloud(mesh) helpers._assert_shape_eq(pts, (10000, 3))
def test_upsample_voxelgrid(device='cpu'): voxel = torch.ones([32, 32, 32]).to(device) up = kal.transforms.UpsampleVoxelGrid(64) helpers._assert_shape_eq(up(voxel), (64, 64, 64)) up = kal.transforms.UpsampleVoxelGrid(33) helpers._assert_shape_eq(up(voxel), (33, 33, 33))
def test_smoke(device='cpu'): net = kal.models.VoxelSuperresSimple.EncoderDecoder() x = torch.randn(1, 1, 32, 32, 32) x_ = net(x) helpers._assert_shape_eq(x, x_.shape, dim=0)
def test_smoke(device='cpu'): net = kal.models.VoxelSuperresODM.SuperresNetwork(30, 15) x = torch.randn(1, 6, 128, 128) x_ = net(x) helpers._assert_shape_eq(x, x_.shape, dim=0)