learning3d 0.0.1__py3-none-any.whl

learning3d/utils/pointconv_util.py

@@ -0,0 +1,382 @@
+"""
+Utility function for PointConv
+Originally from : https://github.com/yanx27/Pointnet_Pointnet2_pytorch/blob/master/utils.py
+Modify by Wenxuan Wu
+Date: September 2019
+"""
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from time import time
+import numpy as np
+from sklearn.neighbors._kde import KernelDensity
+
+def timeit(tag, t):
+    print("{}: {}s".format(tag, time() - t))
+    return time()
+
+def square_distance(src, dst):
+    """
+    Calculate Euclid distance between each two points.
+
+    src^T * dst = xn * xm + yn * ym + zn * zm；
+    sum(src^2, dim=-1) = xn*xn + yn*yn + zn*zn;
+    sum(dst^2, dim=-1) = xm*xm + ym*ym + zm*zm;
+    dist = (xn - xm)^2 + (yn - ym)^2 + (zn - zm)^2
+         = sum(src**2,dim=-1)+sum(dst**2,dim=-1)-2*src^T*dst
+
+    Input:
+        src: source points, [B, N, C]
+        dst: target points, [B, M, C]
+    Output:
+        dist: per-point square distance, [B, N, M]
+    """
+    B, N, _ = src.shape
+    _, M, _ = dst.shape
+    dist = -2 * torch.matmul(src, dst.permute(0, 2, 1))
+    dist += torch.sum(src ** 2, -1).view(B, N, 1)
+    dist += torch.sum(dst ** 2, -1).view(B, 1, M)
+    return dist
+
+def index_points(points, idx):
+    """
+
+    Input:
+        points: input points data, [B, N, C]
+        idx: sample index data, [B, S]
+    Return:
+        new_points:, indexed points data, [B, S, C]
+    """
+    device = points.device
+    B = points.shape[0]
+    view_shape = list(idx.shape)
+    view_shape[1:] = [1] * (len(view_shape) - 1)
+    repeat_shape = list(idx.shape)
+    repeat_shape[0] = 1
+    batch_indices = torch.arange(B, dtype=torch.long).to(device).view(view_shape).repeat(repeat_shape)
+    new_points = points[batch_indices, idx, :]
+    return new_points
+
+def farthest_point_sample(xyz, npoint):
+    """
+    Input:
+        xyz: pointcloud data, [B, N, C]
+        npoint: number of samples
+    Return:
+        centroids: sampled pointcloud index, [B, npoint]
+    """
+    #import ipdb; ipdb.set_trace()
+    device = xyz.device
+    B, N, C = xyz.shape
+    centroids = torch.zeros(B, npoint, dtype=torch.long).to(device)
+    distance = torch.ones(B, N).to(device) * 1e10
+    #farthest = torch.randint(0, N, (B,), dtype=torch.long).to(device)
+    farthest = torch.zeros(B, dtype=torch.long).to(device)
+    batch_indices = torch.arange(B, dtype=torch.long).to(device)
+    for i in range(npoint):
+        centroids[:, i] = farthest
+        centroid = xyz[batch_indices, farthest, :].view(B, 1, 3)
+        dist = torch.sum((xyz - centroid) ** 2, -1)
+        mask = dist < distance
+        distance[mask] = dist[mask]
+        farthest = torch.max(distance, -1)[1]
+    return centroids
+
+def query_ball_point(radius, nsample, xyz, new_xyz):
+    """
+    Input:
+        radius: local region radius
+        nsample: max sample number in local region
+        xyz: all points, [B, N, C]
+        new_xyz: query points, [B, S, C]
+    Return:
+        group_idx: grouped points index, [B, S, nsample]
+    """
+    device = xyz.device
+    B, N, C = xyz.shape
+    _, S, _ = new_xyz.shape
+    group_idx = torch.arange(N, dtype=torch.long).to(device).view(1, 1, N).repeat([B, S, 1])
+    sqrdists = square_distance(new_xyz, xyz)
+    group_idx[sqrdists > radius ** 2] = N
+    group_idx = group_idx.sort(dim=-1)[0][:, :, :nsample]
+    group_first = group_idx[:, :, 0].view(B, S, 1).repeat([1, 1, nsample])
+    mask = group_idx == N
+    group_idx[mask] = group_first[mask]
+    return group_idx
+
+def knn_point(nsample, xyz, new_xyz):
+    """
+    Input:
+        nsample: max sample number in local region
+        xyz: all points, [B, N, C]
+        new_xyz: query points, [B, S, C]
+    Return:
+        group_idx: grouped points index, [B, S, nsample]
+    """
+    sqrdists = square_distance(new_xyz, xyz)
+    _, group_idx = torch.topk(sqrdists, nsample, dim = -1, largest=False, sorted=False)
+    return group_idx
+
+def sample_and_group(npoint, nsample, xyz, points, density_scale = None):
+    """
+    Input:
+        npoint:
+        nsample:
+        xyz: input points position data, [B, N, C]
+        points: input points data, [B, N, D]
+    Return:
+        new_xyz: sampled points position data, [B, 1, C]
+        new_points: sampled points data, [B, 1, N, C+D]
+    """
+    B, N, C = xyz.shape
+    S = npoint
+    fps_idx = farthest_point_sample(xyz, npoint) # [B, npoint, C]
+    new_xyz = index_points(xyz, fps_idx)
+    idx = knn_point(nsample, xyz, new_xyz)
+    grouped_xyz = index_points(xyz, idx) # [B, npoint, nsample, C]
+    grouped_xyz_norm = grouped_xyz - new_xyz.view(B, S, 1, C)
+    if points is not None:
+        grouped_points = index_points(points, idx)
+        new_points = torch.cat([grouped_xyz_norm, grouped_points], dim=-1) # [B, npoint, nsample, C+D]
+    else:
+        new_points = grouped_xyz_norm
+
+    if density_scale is None:
+        return new_xyz, new_points, grouped_xyz_norm, idx
+    else:
+        grouped_density = index_points(density_scale, idx)
+        return new_xyz, new_points, grouped_xyz_norm, idx, grouped_density
+
+def sample_and_group_all(xyz, points, density_scale = None):
+    """
+    Input:
+        xyz: input points position data, [B, N, C]
+        points: input points data, [B, N, D]
+    Return:
+        new_xyz: sampled points position data, [B, 1, C]
+        new_points: sampled points data, [B, 1, N, C+D]
+    """
+    device = xyz.device
+    B, N, C = xyz.shape
+    #new_xyz = torch.zeros(B, 1, C).to(device)
+    new_xyz = xyz.mean(dim = 1, keepdim = True)
+    grouped_xyz = xyz.view(B, 1, N, C) - new_xyz.view(B, 1, 1, C)
+    if points is not None:
+        new_points = torch.cat([grouped_xyz, points.view(B, 1, N, -1)], dim=-1)
+    else:
+        new_points = grouped_xyz
+    if density_scale is None:
+        return new_xyz, new_points, grouped_xyz
+    else:
+        grouped_density = density_scale.view(B, 1, N, 1)
+        return new_xyz, new_points, grouped_xyz, grouped_density
+
+def group(nsample, xyz, points):
+    """
+    Input:
+        npoint:
+        nsample:
+        xyz: input points position data, [B, N, C]
+        points: input points data, [B, N, D]
+    Return:
+        new_xyz: sampled points position data, [B, 1, C]
+        new_points: sampled points data, [B, 1, N, C+D]
+    """
+    B, N, C = xyz.shape
+    S = N
+    new_xyz = xyz
+    idx = knn_point(nsample, xyz, new_xyz)
+    grouped_xyz = index_points(xyz, idx) # [B, npoint, nsample, C]
+    grouped_xyz_norm = grouped_xyz - new_xyz.view(B, S, 1, C)
+    if points is not None:
+        grouped_points = index_points(points, idx)
+        new_points = torch.cat([grouped_xyz_norm, grouped_points], dim=-1) # [B, npoint, nsample, C+D]
+    else:
+        new_points = grouped_xyz_norm
+
+    return new_points, grouped_xyz_norm
+
+def compute_density(xyz, bandwidth):
+    '''
+    xyz: input points position data, [B, N, C]
+    '''
+    #import ipdb; ipdb.set_trace()
+    B, N, C = xyz.shape
+    sqrdists = square_distance(xyz, xyz)
+    gaussion_density = torch.exp(- sqrdists / (2.0 * bandwidth * bandwidth)) / (2.5 * bandwidth)
+    xyz_density = gaussion_density.mean(dim = -1)
+
+    return xyz_density
+
+class DensityNet(nn.Module):
+    def __init__(self, hidden_unit = [16, 8]):
+        super(DensityNet, self).__init__()
+        self.mlp_convs = nn.ModuleList()
+        self.mlp_bns = nn.ModuleList() 
+
+        self.mlp_convs.append(nn.Conv2d(1, hidden_unit[0], 1))
+        self.mlp_bns.append(nn.BatchNorm2d(hidden_unit[0]))
+        for i in range(1, len(hidden_unit)):
+            self.mlp_convs.append(nn.Conv2d(hidden_unit[i - 1], hidden_unit[i], 1))
+            self.mlp_bns.append(nn.BatchNorm2d(hidden_unit[i]))
+        self.mlp_convs.append(nn.Conv2d(hidden_unit[-1], 1, 1))
+        self.mlp_bns.append(nn.BatchNorm2d(1))
+
+    def forward(self, density_scale):
+        for i, conv in enumerate(self.mlp_convs):
+            bn = self.mlp_bns[i]
+            density_scale =  bn(conv(density_scale))
+            if i == len(self.mlp_convs):
+                density_scale = F.sigmoid(density_scale)
+            else:
+                density_scale = F.relu(density_scale)
+        
+        return density_scale
+
+class WeightNet(nn.Module):
+
+    def __init__(self, in_channel, out_channel, hidden_unit = [8, 8]):
+        super(WeightNet, self).__init__()
+
+        self.mlp_convs = nn.ModuleList()
+        self.mlp_bns = nn.ModuleList()
+        if hidden_unit is None or len(hidden_unit) == 0:
+            self.mlp_convs.append(nn.Conv2d(in_channel, out_channel, 1))
+            self.mlp_bns.append(nn.BatchNorm2d(out_channel))
+        else:
+            self.mlp_convs.append(nn.Conv2d(in_channel, hidden_unit[0], 1))
+            self.mlp_bns.append(nn.BatchNorm2d(hidden_unit[0]))
+            for i in range(1, len(hidden_unit)):
+                self.mlp_convs.append(nn.Conv2d(hidden_unit[i - 1], hidden_unit[i], 1))
+                self.mlp_bns.append(nn.BatchNorm2d(hidden_unit[i]))
+            self.mlp_convs.append(nn.Conv2d(hidden_unit[-1], out_channel, 1))
+            self.mlp_bns.append(nn.BatchNorm2d(out_channel))
+        
+    def forward(self, localized_xyz):
+        #xyz : BxCxKxN
+
+        weights = localized_xyz
+        for i, conv in enumerate(self.mlp_convs):
+            bn = self.mlp_bns[i]
+            weights =  F.relu(bn(conv(weights)))
+
+        return weights
+
+class PointConvSetAbstraction(nn.Module):
+    def __init__(self, npoint, nsample, in_channel, mlp, group_all):
+        super(PointConvSetAbstraction, self).__init__()
+        self.npoint = npoint
+        self.nsample = nsample
+        self.mlp_convs = nn.ModuleList()
+        self.mlp_bns = nn.ModuleList()
+        last_channel = in_channel
+        for out_channel in mlp:
+            self.mlp_convs.append(nn.Conv2d(last_channel, out_channel, 1))
+            self.mlp_bns.append(nn.BatchNorm2d(out_channel))
+            last_channel = out_channel
+
+        self.weightnet = WeightNet(3, 16)
+        self.linear = nn.Linear(16 * mlp[-1], mlp[-1])
+        self.bn_linear = nn.BatchNorm1d(mlp[-1])
+        self.group_all = group_all
+
+    def forward(self, xyz, points):
+        """
+        Input:
+            xyz: input points position data, [B, C, N]
+            points: input points data, [B, D, N]
+        Return:
+            new_xyz: sampled points position data, [B, C, S]
+            new_points_concat: sample points feature data, [B, D', S]
+        """
+        B = xyz.shape[0]
+        xyz = xyz.permute(0, 2, 1)
+        if points is not None:
+            points = points.permute(0, 2, 1)
+
+        if self.group_all:
+            new_xyz, new_points, grouped_xyz_norm = sample_and_group_all(xyz, points)
+        else:
+            new_xyz, new_points, grouped_xyz_norm, _ = sample_and_group(self.npoint, self.nsample, xyz, points)
+        # new_xyz: sampled points position data, [B, npoint, C]
+        # new_points: sampled points data, [B, npoint, nsample, C+D]
+        new_points = new_points.permute(0, 3, 2, 1) # [B, C+D, nsample,npoint]
+        for i, conv in enumerate(self.mlp_convs):
+            bn = self.mlp_bns[i]
+            new_points =  F.relu(bn(conv(new_points)))
+
+        grouped_xyz = grouped_xyz_norm.permute(0, 3, 2, 1)
+        weights = self.weightnet(grouped_xyz)
+        new_points = torch.matmul(input=new_points.permute(0, 3, 1, 2), other = weights.permute(0, 3, 2, 1)).view(B, self.npoint, -1)
+        new_points = self.linear(new_points)
+        new_points = self.bn_linear(new_points.permute(0, 2, 1))
+        new_points = F.relu(new_points)
+        new_xyz = new_xyz.permute(0, 2, 1)
+
+        return new_xyz, new_points
+
+class PointConvDensitySetAbstraction(nn.Module):
+    def __init__(self, npoint, nsample, in_channel, mlp, bandwidth, group_all):
+        super(PointConvDensitySetAbstraction, self).__init__()
+        self.npoint = npoint
+        self.nsample = nsample
+        self.mlp_convs = nn.ModuleList()
+        self.mlp_bns = nn.ModuleList()
+        last_channel = in_channel
+        for out_channel in mlp:
+            self.mlp_convs.append(nn.Conv2d(last_channel, out_channel, 1))
+            self.mlp_bns.append(nn.BatchNorm2d(out_channel))
+            last_channel = out_channel
+
+        self.weightnet = WeightNet(3, 16)
+        self.linear = nn.Linear(16 * mlp[-1], mlp[-1])
+        self.bn_linear = nn.BatchNorm1d(mlp[-1])
+        self.densitynet = DensityNet()
+        self.group_all = group_all
+        self.bandwidth = bandwidth
+
+    def forward(self, xyz, points):
+        """
+        Input:
+            xyz: input points position data, [B, C, N]
+            points: input points data, [B, D, N]
+        Return:
+            new_xyz: sampled points position data, [B, C, S]
+            new_points_concat: sample points feature data, [B, D', S]
+        """
+        B = xyz.shape[0]
+        N = xyz.shape[2]
+        xyz = xyz.permute(0, 2, 1)
+        if points is not None:
+            points = points.permute(0, 2, 1)
+
+        xyz_density = compute_density(xyz, self.bandwidth)
+        inverse_density = 1.0 / xyz_density 
+
+        if self.group_all:
+            new_xyz, new_points, grouped_xyz_norm, grouped_density = sample_and_group_all(xyz, points, inverse_density.view(B, N, 1))
+        else:
+            new_xyz, new_points, grouped_xyz_norm, _, grouped_density = sample_and_group(self.npoint, self.nsample, xyz, points, inverse_density.view(B, N, 1))
+        # new_xyz: sampled points position data, [B, npoint, C]
+        # new_points: sampled points data, [B, npoint, nsample, C+D]
+        new_points = new_points.permute(0, 3, 2, 1) # [B, C+D, nsample,npoint]
+        for i, conv in enumerate(self.mlp_convs):
+            bn = self.mlp_bns[i]
+            new_points =  F.relu(bn(conv(new_points)))
+
+        inverse_max_density = grouped_density.max(dim = 2, keepdim=True)[0]
+        density_scale = grouped_density / inverse_max_density
+        density_scale = self.densitynet(density_scale.permute(0, 3, 2, 1))
+        new_points = new_points * density_scale
+
+        grouped_xyz = grouped_xyz_norm.permute(0, 3, 2, 1)
+        weights = self.weightnet(grouped_xyz)     
+        new_points = torch.matmul(input=new_points.permute(0, 3, 1, 2), other = weights.permute(0, 3, 2, 1)).view(B, self.npoint, -1)
+        new_points = self.linear(new_points)
+        new_points = self.bn_linear(new_points.permute(0, 2, 1))
+        new_points = F.relu(new_points)
+        new_xyz = new_xyz.permute(0, 2, 1)
+
+        return new_xyz, new_points
+
+        

learning3d/utils/ppfnet_util.py

@@ -0,0 +1,244 @@
+"""Utilities for PointNet related functions
+
+Modified from:
+	Pytorch Implementation of PointNet and PointNet++
+	https://github.com/yanx27/Pointnet_Pointnet2_pytorch
+"""
+
+import torch
+
+
+def angle_difference(src, dst):
+	"""Calculate angle between each pair of vectors.
+	Assumes points are l2-normalized to unit length.
+
+	Input:
+		src: source points, [B, N, C]
+		dst: target points, [B, M, C]
+	Output:
+		dist: per-point square distance, [B, N, M]
+	"""
+	B, N, _ = src.shape
+	_, M, _ = dst.shape
+	dist = torch.matmul(src, dst.permute(0, 2, 1))
+	dist = torch.acos(dist)
+
+	return dist
+
+
+def square_distance(src, dst):
+	"""Calculate Euclid distance between each two points.
+		src^T * dst = xn * xm + yn * ym + zn * zm；
+		sum(src^2, dim=-1) = xn*xn + yn*yn + zn*zn;
+		sum(dst^2, dim=-1) = xm*xm + ym*ym + zm*zm;
+		dist = (xn - xm)^2 + (yn - ym)^2 + (zn - zm)^2
+			 = sum(src**2,dim=-1)+sum(dst**2,dim=-1)-2*src^T*dst
+
+	Args:
+		src: source points, [B, N, C]
+		dst: target points, [B, M, C]
+	Returns:
+		dist: per-point square distance, [B, N, M]
+	"""
+	B, N, _ = src.shape
+	_, M, _ = dst.shape
+	dist = -2 * torch.matmul(src, dst.permute(0, 2, 1))
+	dist += torch.sum(src ** 2, dim=-1)[:, :, None]
+	dist += torch.sum(dst ** 2, dim=-1)[:, None, :]
+	return dist
+
+
+def index_points(points, idx):
+	"""Array indexing, i.e. retrieves relevant points based on indices
+
+	Args:
+		points: input points data_loader, [B, N, C]
+		idx: sample index data_loader, [B, S]. S can be 2 dimensional
+	Returns:
+		new_points:, indexed points data_loader, [B, S, C]
+	"""
+	device = points.device
+	B = points.shape[0]
+	view_shape = list(idx.shape)
+	view_shape[1:] = [1] * (len(view_shape) - 1)
+	repeat_shape = list(idx.shape)
+	repeat_shape[0] = 1
+	batch_indices = torch.arange(B, dtype=torch.long).to(device).view(view_shape).repeat(repeat_shape)
+	new_points = points[batch_indices, idx, :]
+	return new_points
+
+
+def farthest_point_sample(xyz, npoint):
+	"""Iterative farthest point sampling
+
+	Args:
+		xyz: pointcloud data_loader, [B, N, C]
+		npoint: number of samples
+	Returns:
+		centroids: sampled pointcloud index, [B, npoint]
+	"""
+	device = xyz.device
+	B, N, C = xyz.shape
+	centroids = torch.zeros(B, npoint, dtype=torch.long).to(device)
+	distance = torch.ones(B, N).to(device) * 1e10
+	farthest = torch.randint(0, N, (B,), dtype=torch.long).to(device)
+	batch_indices = torch.arange(B, dtype=torch.long).to(device)
+	for i in range(npoint):
+		centroids[:, i] = farthest
+		centroid = xyz[batch_indices, farthest, :].view(B, 1, 3)
+		dist = torch.sum((xyz - centroid) ** 2, -1)
+		mask = dist < distance
+		distance[mask] = dist[mask]
+		farthest = torch.max(distance, -1)[1]
+	return centroids
+
+
+def query_ball_point(radius, nsample, xyz, new_xyz, itself_indices=None):
+	""" Grouping layer in PointNet++.
+
+	Inputs:
+		radius: local region radius
+		nsample: max sample number in local region
+		xyz: all points, (B, N, C)
+		new_xyz: query points, (B, S, C)
+		itself_indices (Optional): Indices of new_xyz into xyz (B, S).
+		  Used to try and prevent grouping the point itself into the neighborhood.
+		  If there is insufficient points in the neighborhood, or if left is none, the resulting cluster will
+		  still contain the center point.
+	Returns:
+		group_idx: grouped points index, [B, S, nsample]
+	"""
+	device = xyz.device
+	B, N, C = xyz.shape
+	_, S, _ = new_xyz.shape
+	group_idx = torch.arange(N, dtype=torch.long).to(device).view(1, 1, N).repeat([B, S, 1])  # (B, S, N)
+	sqrdists = square_distance(new_xyz, xyz)
+
+	if itself_indices is not None:
+		# Remove indices of the center points so that it will not be chosen
+		batch_indices = torch.arange(B, dtype=torch.long).to(device)[:, None].repeat(1, S)  # (B, S)
+		row_indices = torch.arange(S, dtype=torch.long).to(device)[None, :].repeat(B, 1)  # (B, S)
+		group_idx[batch_indices, row_indices, itself_indices] = N
+
+	group_idx[sqrdists > radius ** 2] = N
+	group_idx = group_idx.sort(dim=-1)[0][:, :, :nsample]
+	if itself_indices is not None:
+		group_first = itself_indices[:, :, None].repeat([1, 1, nsample])
+	else:
+		group_first = group_idx[:, :, 0].view(B, S, 1).repeat([1, 1, nsample])
+	mask = group_idx == N
+	group_idx[mask] = group_first[mask]
+	return group_idx
+
+
+def sample_and_group(npoint: int, radius: float, nsample: int, xyz: torch.Tensor, points: torch.Tensor,
+					 returnfps: bool=False):
+	"""
+	Args:
+		npoint (int): Set to negative to compute for all points
+		radius:
+		nsample:
+		xyz: input points position data_loader, [B, N, C]
+		points: input points data_loader, [B, N, D]
+		returnfps (bool) Whether to return furthest point indices
+	Returns:
+		new_xyz: sampled points position data_loader, [B, 1, C]
+		new_points: sampled points data_loader, [B, 1, N, C+D]
+	"""
+	B, N, C = xyz.shape
+
+	if npoint > 0:
+		S = npoint
+		fps_idx = farthest_point_sample(xyz, npoint)  # [B, npoint, C]
+		new_xyz = index_points(xyz, fps_idx)
+	else:
+		S = xyz.shape[1]
+		fps_idx = torch.arange(0, xyz.shape[1])[None, ...].repeat(xyz.shape[0], 1)
+		new_xyz = xyz
+
+	idx = query_ball_point(radius, nsample, xyz, new_xyz)  # (B, N, nsample)
+	grouped_xyz = index_points(xyz, idx)  # (B, npoint, nsample, C)
+	grouped_xyz_norm = grouped_xyz - new_xyz.view(B, S, 1, C)
+	if points is not None:
+		grouped_points = index_points(points, idx)
+		new_points = torch.cat([grouped_xyz_norm, grouped_points], dim=-1) # [B, npoint, nsample, C+D]
+	else:
+		new_points = grouped_xyz_norm
+	if returnfps:
+		return new_xyz, new_points, grouped_xyz, fps_idx
+	else:
+		return new_xyz, new_points
+
+
+def angle(v1: torch.Tensor, v2: torch.Tensor):
+	"""Compute angle between 2 vectors
+
+	For robustness, we use the same formulation as in PPFNet, i.e.
+		angle(v1, v2) = atan2(cross(v1, v2), dot(v1, v2)).
+	This handles the case where one of the vectors is 0.0, since torch.atan2(0.0, 0.0)=0.0
+
+	Args:
+		v1: (B, *, 3)
+		v2: (B, *, 3)
+
+	Returns:
+
+	"""
+
+	cross_prod = torch.stack([v1[..., 1] * v2[..., 2] - v1[..., 2] * v2[..., 1],
+							  v1[..., 2] * v2[..., 0] - v1[..., 0] * v2[..., 2],
+							  v1[..., 0] * v2[..., 1] - v1[..., 1] * v2[..., 0]], dim=-1)
+	cross_prod_norm = torch.norm(cross_prod, dim=-1)
+	dot_prod = torch.sum(v1 * v2, dim=-1)
+
+	return torch.atan2(cross_prod_norm, dot_prod)
+
+
+def sample_and_group_multi(npoint: int, radius: float, nsample: int, xyz: torch.Tensor, normals: torch.Tensor,
+						   returnfps: bool = False):
+	"""Sample and group for xyz, dxyz and ppf features
+
+	Args:
+		npoint(int): Number of clusters (equivalently, keypoints) to sample.
+					 Set to negative to compute for all points
+		radius(int): Radius of cluster for computing local features
+		nsample: Maximum number of points to consider per cluster
+		xyz: XYZ coordinates of the points
+		normals: Corresponding normals for the points (required for ppf computation)
+		returnfps: Whether to return indices of FPS points and their neighborhood
+
+	Returns:
+		Dictionary containing the following fields ['xyz', 'dxyz', 'ppf'].
+		If returnfps is True, also returns: grouped_xyz, fps_idx
+	"""
+
+	B, N, C = xyz.shape
+
+	if npoint > 0:
+		S = npoint
+		fps_idx = farthest_point_sample(xyz, npoint)  # [B, npoint, C]
+		new_xyz = index_points(xyz, fps_idx)
+		nr = index_points(normals, fps_idx)[:, :, None, :]
+	else:
+		S = xyz.shape[1]
+		fps_idx = torch.arange(0, xyz.shape[1])[None, ...].repeat(xyz.shape[0], 1).to(xyz.device)
+		new_xyz = xyz
+		nr = normals[:, :, None, :]
+
+	idx = query_ball_point(radius, nsample, xyz, new_xyz, fps_idx)  # (B, npoint, nsample)
+	grouped_xyz = index_points(xyz, idx)  # (B, npoint, nsample, C)
+	d = grouped_xyz - new_xyz.view(B, S, 1, C)  # d = p_r - p_i  (B, npoint, nsample, 3)
+	ni = index_points(normals, idx)
+
+	nr_d = angle(nr, d)
+	ni_d = angle(ni, d)
+	nr_ni = angle(nr, ni)
+	d_norm = torch.norm(d, dim=-1)
+
+	xyz_feat = d  # (B, npoint, n_sample, 3)
+	ppf_feat = torch.stack([nr_d, ni_d, nr_ni, d_norm], dim=-1)  # (B, npoint, n_sample, 4)
+
+	if returnfps:
+		return {'xyz': new_xyz, 'dxyz': xyz_feat, 'ppf': ppf_feat}, grouped_xyz, fps_idx
+	else:
+		return {'xyz': new_xyz, 'dxyz': xyz_feat, 'ppf': ppf_feat}

learning3d/utils/svd.py

@@ -0,0 +1,59 @@
+import torch
+import torch.nn as nn
+import math
+
+class SVDHead(nn.Module):
+    def __init__(self, emb_dims, input_shape="bnc"):
+        super(SVDHead, self).__init__()
+        self.emb_dims = emb_dims
+        self.reflect = nn.Parameter(torch.eye(3), requires_grad=False)
+        self.reflect[2, 2] = -1
+        self.input_shape = input_shape
+
+    def forward(self, *input):
+        src_embedding = input[0]
+        tgt_embedding = input[1]
+        src = input[2]
+        tgt = input[3]
+        batch_size = src.size(0)
+        if self.input_shape == "bnc":
+            src = src.permute(0, 2, 1)
+            tgt = tgt.permute(0, 2, 1)
+
+        d_k = src_embedding.size(1)
+        scores = torch.matmul(src_embedding.transpose(2, 1).contiguous(), tgt_embedding) / math.sqrt(d_k)
+        scores = torch.softmax(scores, dim=2)
+
+        src_corr = torch.matmul(tgt, scores.transpose(2, 1).contiguous())
+
+        src_centered = src - src.mean(dim=2, keepdim=True)
+
+        src_corr_centered = src_corr - src_corr.mean(dim=2, keepdim=True)
+
+        H = torch.matmul(src_centered, src_corr_centered.transpose(2, 1).contiguous())
+
+        U, S, V = [], [], []
+        R = []
+
+        for i in range(src.size(0)):
+            u, s, v = torch.svd(H[i])
+            r = torch.matmul(v, u.transpose(1, 0).contiguous())
+            r_det = torch.det(r)
+            if r_det < 0:
+                u, s, v = torch.svd(H[i])
+                v = torch.matmul(v, self.reflect)
+                r = torch.matmul(v, u.transpose(1, 0).contiguous())
+                # r = r * self.reflect
+            R.append(r)
+
+            U.append(u)
+            S.append(s)
+            V.append(v)
+
+        U = torch.stack(U, dim=0)
+        V = torch.stack(V, dim=0)
+        S = torch.stack(S, dim=0)
+        R = torch.stack(R, dim=0)
+
+        t = torch.matmul(-R, src.mean(dim=2, keepdim=True)) + src_corr.mean(dim=2, keepdim=True)
+        return R, t.view(batch_size, 3)