openpoints 0.1.0__py3-none-any.whl

openpoints/models/layers/group_embed.py

@@ -0,0 +1,286 @@
+# Patch embedding for 2D/3D data
+# Reference:
+import math
+import torch
+from torch import nn as nn
+import torch.nn.functional as F
+from .subsample import furthest_point_sample, random_sample
+from .group import KNNGroup, QueryAndGroup, create_grouper, get_aggregation_feautres
+from .conv import create_convblock1d, create_convblock2d, create_linearblock, create_norm, create_act
+from .local_aggregation import CHANNEL_MAP
+from ..build import MODELS
+
+
+class SubsampleGroup(nn.Module):
+    """ Point cloud to subsampled groups
+    """
+
+    def __init__(self,
+                 num_groups=256, group_size=32,
+                 subsample='fps',  # random, FPS
+                 group='ballquery',
+                 radius=0.1,
+                 **kwargs
+                 ):
+        super().__init__()
+        self.num_groups = num_groups
+        self.group_size = group_size
+
+        self.subsample = subsample
+        self.group = group
+
+        if 'ball' in self.group.lower() or 'query' in self.group.lower():
+            self.grouper = QueryAndGroup(radius, self.group_size)
+        elif 'knn' in self.group.lower():
+            self.grouper = KNNGroup(self.group_size)
+        else:
+            raise NotImplementedError(f'{self.group.lower()} is not implemented. Only support ballquery, knn')
+
+    def forward(self, p, x=None):
+        if 'fps' in self.subsample.lower() or 'furthest' in self.subsample.lower() or 'farthest' in self.subsample.lower():
+            idx = furthest_point_sample(p, self.num_groups).to(torch.int64)
+        elif 'random' in self.subsample.lower() or 'rs' in self.subsample.lower():
+            idx = random_sample(p, self.num_groups)
+        else:
+            raise NotImplementedError(f'{self.subsample.lower()} is not implemented. Only support fps, random')
+        center_p = torch.gather(p, 1,
+                                  idx.unsqueeze(-1).expand(-1, -1, 3))  # downsampled point cloud, [B, npoint, 3]
+        if x is not None:
+            B, C, N = x.shape[:3]
+            center_x = torch.gather(x, 2, idx.unsqueeze(1).expand(-1, C, -1)).unsqueeze(-1)
+            grouped_p, fj = self.grouper(center_p, p, x)
+            return grouped_p, center_p, fj, center_x
+        else:
+            grouped_p, _ = self.grouper(center_p, p)
+            return grouped_p, center_p
+
+
+@MODELS.register_module()
+class PointPatchEmbed(nn.Module):
+    """ Point cloud to Group Embedding using GCN
+    Patch Embedding for 3d data (point cloud)
+    A convolution based approach to patchifying a point cloud w/ embedding projection.
+    """
+
+    def __init__(self,
+                 sample_ratio=0.0625, group_size=32,
+                 in_channels=3,
+                 layers=4,
+                 embed_dim=256,
+                 channels=None,
+                 subsample='fps',  # random, FPS
+                 group='ballquery',
+                 normalize_dp=False,
+                 radius=0.1,
+                 feature_type='dp_df',
+                 relative_xyz=True,
+                 norm_args={'norm': 'bn1d'},
+                 act_args={'act': 'relu'},
+                 conv_args={'order': 'conv-norm-act'},
+                 reduction='max',
+                 **kwargs
+                 ):
+        super().__init__()
+        self.sample_ratio = sample_ratio
+        self.group_size = group_size
+
+        self.feature_type = feature_type
+        # subsample layer and group layer
+        if subsample.lower() == 'fps':
+            self.sample_fn = furthest_point_sample
+        elif 'random' in subsample.lower():
+            self.sample_fn = random_sample
+
+        # TODO: make this embedding progressively
+        self.group = group.lower()
+        if 'ball' in self.group or 'query' in self.group:
+            self.grouper = QueryAndGroup(nsample=self.group_size,
+                                         relative_xyz=relative_xyz, normalize_dp=normalize_dp,
+                                         radius=radius)
+        elif 'knn' in self.group.lower():
+            self.grouper = KNNGroup(self.group_size, relative_xyz=relative_xyz, normalize_dp=normalize_dp)
+        else:
+            raise NotImplementedError(f'{self.group.lower()} is not implemented. Only support ballquery, knn')
+
+        # # convolutions
+        if channels is None:
+            channels = [CHANNEL_MAP[feature_type](in_channels)] + [embed_dim] * (layers // 2) + [embed_dim * 2] * (
+                    layers // 2 - 1) + [embed_dim]
+        else:
+            channels = [CHANNEL_MAP[feature_type](in_channels)] + channels + [embed_dim]
+            layers = len(channels) -1
+        conv1 = []
+        for i in range(layers // 2):
+            conv1.append(create_convblock2d(channels[i], channels[i + 1],
+                                            norm_args=norm_args if i!=(layers//2-1) else None,
+                                            act_args=act_args if i!=(layers//2-1) else None,
+                                            **conv_args))
+        self.conv1 = nn.Sequential(*conv1)
+
+        channels[layers // 2] *= 2
+        conv2 = []
+        for i in range(layers // 2, layers):
+            conv2.append(create_convblock2d(channels[i], channels[i + 1],
+                                            norm_args=norm_args if i!=(layers-1) else None,
+                                            act_args=act_args if i!=(layers-1) else None,
+                                            **conv_args
+                                            ))
+        self.conv2 = nn.Sequential(*conv2)
+
+        # reduction layer
+        if reduction in ['mean', 'avg', 'meanpool', 'avgpool']:
+            self.pool = lambda x: torch.mean(x, dim=-1, keepdim=True)
+        else:
+            self.pool = lambda x: torch.max(x, dim=-1, keepdim=True)[0]
+        self.out_channels = channels[-1]
+        self.channel_list = [in_channels, embed_dim]
+
+    def forward(self, p, x=None):
+        # downsample
+        B, N, _ = p.shape[:3]
+        idx = self.sample_fn(p, int(N * self.sample_ratio)).long()
+        center_p = torch.gather(p, 1, idx.unsqueeze(-1).expand(-1, -1, 3))
+
+        # query neighbors.
+        dp, fj = self.grouper(center_p, p, x)
+
+        """visualization
+        from openpoints.dataset.vis3d import vis_multi_points
+        new_p = (dp.permute(0, 2, 3, 1) + center_p.unsqueeze(2)).view(B, -1, 3)
+        vis_multi_points([p[0].cpu().numpy(), new_p[0].cpu().numpy(), center_p[0].cpu().numpy()])
+        """
+
+        # concat neighborhood x
+        # TODO: using a local aggregation layer
+        if self.feature_type == 'dp':
+            fj = dp
+        elif self.feature_type == 'dp_fj':
+            fj = torch.cat([dp, fj], dim=1)
+        elif self.feature_type == 'dp_df':
+            center_x = torch.gather(x, 2, idx.unsqueeze(1).expand(-1, x.shape[1], -1))
+            fj = torch.cat([dp, fj - center_x.unsqueeze(-1)], dim=1)
+        elif self.feature_type == 'df':
+            center_x = torch.gather(x, 2, idx.unsqueeze(1).expand(-1, x.shape[1], -1))
+            fj = fj - center_x.unsqueeze(-1)
+        fj = self.conv1(fj)
+
+        fj = torch.cat(
+            [self.pool(fj).expand(-1, -1, -1, self.group_size),
+             fj],
+            dim=1)
+        out_f = self.pool(self.conv2(fj)).squeeze(-1)
+        return [p, center_p], [x, out_f]
+
+
+@MODELS.register_module()
+class P3Embed(nn.Module):
+    """
+    Progressive Point Patch Embedding for 3d data (point cloud)
+    A convolution based approach to patchifying a point cloud w/ embedding projection.
+    """
+
+    def __init__(self,
+                 sample_ratio=0.0625,
+                 scale=4,
+                 group_size=32,
+                 in_channels=3,
+                 layers=4,
+                 embed_dim=256,
+                 subsample='fps',  # random, FPS
+                 group='ballquery',
+                 normalize_dp=False,
+                 radius=0.1,
+                 feature_type='dp_df',
+                 relative_xyz=True,
+                 norm_args={'norm': 'bn1d'},
+                 act_args={'act': 'relu'},
+                 conv_args={'order': 'conv-norm-act'},
+                 reduction='max',
+                 **kwargs
+                 ):
+        super().__init__()
+        self.sample_ratio = sample_ratio
+        self.group_size = group_size
+
+        self.feature_type = feature_type
+        # subsample layer and group layer
+        if subsample.lower() == 'fps':
+            self.sample_fn = furthest_point_sample
+        elif 'random' in subsample.lower():
+            self.sample_fn = random_sample
+
+        self.group = group.lower()
+        if 'ball' in self.group or 'query' in self.group:
+            self.grouper = QueryAndGroup(nsample=self.group_size,
+                                         relative_xyz=relative_xyz, normalize_dp=normalize_dp,
+                                         radius=radius)
+        elif 'knn' in self.group.lower():
+            self.grouper = KNNGroup(self.group_size, relative_xyz=relative_xyz, normalize_dp=normalize_dp)
+        else:
+            raise NotImplementedError(f'{self.group.lower()} is not implemented. Only support ballquery, knn')
+
+        # stages
+        stages = int(math.log(1/sample_ratio, scale))
+        embed_dim = int(embed_dim // 2 ** (stages-1))
+        self.convs = nn.ModuleList()
+        self.channel_list = [in_channels]
+        for _ in range(int(stages)):
+            # convolutions
+            channels = [CHANNEL_MAP[feature_type](in_channels)] + [embed_dim] * (layers // 2) + [embed_dim * 2] * (
+                    layers // 2 - 1) + [embed_dim]
+            conv1 = []
+            for i in range(layers // 2):
+                conv1.append(create_convblock2d(channels[i], channels[i + 1],
+                                                norm_args=norm_args if i!=(layers//2-1) else None,
+                                                act_args=act_args if i!=(layers//2-1) else None,
+                                                **conv_args))
+            conv1 = nn.Sequential(*conv1)
+
+            channels[layers // 2] *= 2
+            conv2 = []
+            for i in range(layers // 2, layers):
+                conv2.append(create_convblock2d(channels[i], channels[i + 1],
+                                                norm_args=norm_args,
+                                                act_args=act_args,
+                                                **conv_args
+                                                ))
+            conv2 = nn.Sequential(*conv2)
+            self.convs.append(nn.ModuleList([conv1, conv2]))
+
+            self.channel_list.append(embed_dim)
+            in_channels = embed_dim
+            embed_dim *= 2
+
+        # reduction layer
+        if reduction in ['mean', 'avg', 'meanpool', 'avgpool']:
+            self.pool = lambda x: torch.mean(x, dim=-1, keepdim=True)
+        else:
+            self.pool = lambda x: torch.max(x, dim=-1, keepdim=True)[0]
+        self.out_channels = channels[-1]
+
+    def forward(self, p, f=None):
+        B, N, _ = p.shape[:3]
+        out_p, out_f = [p], [f]
+        for convs in self.convs:
+            # Progressive downsampling
+            cur_p, cur_f = out_p[-1], out_f[-1]
+            idx = self.sample_fn(cur_p, int(N //4)).long()
+            N = N // 4
+            center_p = torch.gather(cur_p, 1, idx.unsqueeze(-1).expand(-1, -1, 3))
+            center_f = torch.gather(cur_f, 2, idx.unsqueeze(1).expand(-1, cur_f.shape[1], -1))
+
+            # query neighbors.
+            dp, fj = self.grouper(center_p, cur_p, cur_f)
+            fj = get_aggregation_feautres(center_p, dp, center_f, fj, self.feature_type)
+
+            # graph convolutions
+            fj = convs[0](fj)
+            fj = torch.cat(
+                [self.pool(fj).expand(-1, -1, -1, self.group_size),
+                fj],
+                dim=1)
+
+            # output
+            out_f.append(self.pool(convs[1](fj)).squeeze(-1))
+            out_p.append(center_p)
+        return out_p, out_f

openpoints/models/layers/helpers.py

@@ -0,0 +1,43 @@
+""" Layer/Module Helpers
+
+Hacked together by / Copyright 2020 Ross Wightman
+"""
+from itertools import repeat
+import collections.abc
+from torch import nn
+
+
+# From PyTorch internals
+def _ntuple(n):
+    def parse(x):
+        if isinstance(x, collections.abc.Iterable):
+            return x
+        return tuple(repeat(x, n))
+    return parse
+
+
+to_1tuple = _ntuple(1)
+to_2tuple = _ntuple(2)
+to_3tuple = _ntuple(3)
+to_4tuple = _ntuple(4)
+to_ntuple = _ntuple
+
+
+def make_divisible(v, divisor=8, min_value=None, round_limit=.9):
+    min_value = min_value or divisor
+    new_v = max(min_value, int(v + divisor / 2) // divisor * divisor)
+    # Make sure that round down does not go down by more than 10%.
+    if new_v < round_limit * v:
+        new_v += divisor
+    return new_v
+
+
+
+class MultipleSequential(nn.Sequential):
+    def forward(self, *inputs):
+        for module in self._modules.values():
+            if type(inputs) == tuple:
+                inputs = module(*inputs)
+            else:
+                inputs = module(inputs)
+        return inputs

openpoints/models/layers/kmeans.py

@@ -0,0 +1,119 @@
+import torch
+from torch import nn
+from fast_pytorch_kmeans import KMeans, MultiKMeans
+from torch_scatter import scatter
+from .local_aggregation import CHANNEL_MAP
+
+
+class KMeansEmbed(nn.Module):
+    """ Point cloud to subsampled groups
+    """
+
+    def __init__(self,
+                 in_chans=3, 
+                 num_groups=256,
+                 encoder_dim=256,
+                 feature_type='dp', 
+                 **kwargs
+                 ):
+        super().__init__()
+
+        self.num_groups = num_groups
+
+        self.kmeans = MultiKMeans(n_clusters=num_groups, n_kmeans=32, mode='euclidean', verbose=0)
+
+        channels = CHANNEL_MAP[feature_type](in_chans)
+        self.feature_type = feature_type
+        self.conv1 = nn.Sequential(
+            nn.Linear(channels, 128),   #   TODO: here, can be better, edgeconv.
+            nn.LayerNorm(128),
+            nn.ReLU(inplace=True),
+            nn.Linear(128, 256)
+        )
+        self.conv2 = nn.Sequential(
+            nn.Linear(512, 512),
+            nn.LayerNorm(512),
+            nn.ReLU(inplace=True),
+            nn.Linear(512, encoder_dim)
+        )
+
+    def forward(self, xyz, features=None):
+        B, N, _ = xyz.shape
+        self.kmeans.centroids = None  #   re-init kmeans 
+        self.kmeans.n_kmeans = xyz.shape[0]
+        labels = self.kmeans.fit_predict(xyz)   # B,N
+
+        # TODO: BUG, sometimes the value is even smaller than the number of centroids!! 
+        centroids = self.kmeans.centroids   # B, K, 3
+        idx = labels.unsqueeze(-1)
+
+        # p_j = xyz
+        p_i = torch.gather(centroids, 1, labels.unsqueeze(-1).expand(-1, -1, 3))   # B, N, 3
+        relative_xyz = xyz - p_i  # p_j-p_i, B, N, 3
+
+        if self.feature_type == 'dp':
+            neighborhood_features = relative_xyz
+        elif self.feature_type == 'pj_dp':
+            neighborhood_features = torch.cat([xyz, relative_xyz], -1)
+        elif self.feature_type == 'pi_dp':
+            neighborhood_features = torch.cat([p_i, relative_xyz], -1)
+
+        neighborhood_features = self.conv1(neighborhood_features)  # B, N, C
+        pooled_feat = scatter(neighborhood_features, idx, dim=1, reduce='max')  # B, K, C
+        reapted_feat = torch.gather(pooled_feat, 1, idx.expand(-1, -1, pooled_feat.shape[-1]))
+        neighborhood_features = torch.cat([reapted_feat, neighborhood_features], dim=-1)
+        out_features = scatter(self.conv2(neighborhood_features), idx, dim=1, reduce='max')
+        return centroids, out_features, p_i, labels 
+
+
+if __name__ == "__main__":
+    import torch
+    import os, sys
+    sys.path.insert(0, os.path.join(os.path.abspath(__file__), "../../../../"))
+    from openpoints.dataset import ModelNet, vis_points, vis_multi_points
+    from openpoints.models.layers import fps
+
+    B, C, N = 8, 3, 8196
+    device = 'cuda'
+
+    dataset = ModelNet("data/ModelNet40",
+                        N, 40, split='test')
+    test_datalodaer = torch.utils.data.DataLoader(dataset, batch_size=B, num_workers=1
+                                                 )
+    data = iter(test_datalodaer).next()[0]
+    points = data['pos'].to(device)
+
+    points = fps(points, N)
+    print(points.shape)
+    # debug one batch
+    K = 12
+    print(points.shape, points.device)
+
+    #
+    # kmeans_group = KMeansGroup(K).to(device)
+    # kmeans_group(points)
+
+    kmeans = KMeans(n_clusters=K, mode='euclidean', verbose=1)
+    labels = kmeans.fit_predict(points[0])
+    print(labels.shape)
+    # 0.5207s for 10000 points, too slow!
+    # vis_points(points[0], labels=labels)
+    center_points = kmeans.centroids
+    print(center_points.shape)
+    # vis_multi_points([points.cpu().numpy(), center_points])
+    vis_points(points[0], labels=labels)
+
+    # B, N, 3
+    # B, N, 1  (label index)
+
+    # debug 8 batch
+    # K = 8
+    # print(points.shape, points.device)
+    # kmeans = MultiKMeans(n_clusters=K, n_kmeans=B,
+    #                      mode='euclidean', verbose=1)
+    # labels = kmeans.fit_predict(points)
+    # print(labels.shape)
+    # # 0.5207s for 10000 points, too slow!
+    # # vis_multi_points(points.cpu().numpy()[:4], labels=labels.cpu().numpy()[:4])
+    # vis_points(points[0], labels=labels[0])
+    # vis_points(points[1], labels=labels[1])

openpoints/models/layers/knn.py

@@ -0,0 +1,110 @@
+import torch
+import torch.nn as nn
+# from knn_cuda import KNN as KNNCUDA
+
+
+@torch.no_grad()
+def knn_point(k, query, support=None):
+    """Get the distances and indices to a fixed number of neighbors
+        Args:
+            support ([tensor]): [B, N, C]
+            query ([tensor]): [B, M, C]
+
+        Returns:
+            [int]: neighbor idx. [B, M, K]
+    """
+    if support is None:
+        support = query
+    dist = torch.cdist(query, support)
+    k_dist = dist.topk(k=k, dim=-1, largest=False, sorted=True)
+    return k_dist.values, k_dist.indices
+    
+
+class KNN(nn.Module):
+    """Get the distances and indices to a fixed number of neighbors
+
+    Reference: https://gist.github.com/ModarTensai/60fe0d0e3536adc28778448419908f47
+
+    Args:
+        neighbors: number of neighbors to consider
+        p_norm: distances are computed based on L_p norm
+        farthest: whether to get the farthest or the nearest neighbors
+        ordered: distance sorted (descending if `farthest`)
+
+    Returns:
+        (distances, indices) both of shape [B, N, `num_neighbors`]
+    """
+    
+    def __init__(self, neighbors, 
+                 farthest=False, 
+                 sorted=True, 
+                 **kwargs):
+        super(KNN, self).__init__()
+        self.neighbors = neighbors
+        self.farthest = farthest
+        self.sorted = sorted
+
+    @torch.no_grad()
+    def forward(self, query, support=None):
+        """
+        Args:
+            support ([tensor]): [B, N, C]
+            query ([tensor]): [B, M, C]
+
+        Returns:
+            [int]: neighbor idx. [B, M, K]
+        """
+        if support is None:
+            support = query
+        dist = torch.cdist(query, support)
+        k_dist = dist.topk(k=self.neighbors, dim=-1, largest=self.farthest, sorted=self.sorted)
+        return k_dist.values, k_dist.indices.int()
+
+
+# dilated knn
+class DenseDilated(nn.Module):
+    """
+    Find dilated neighbor from neighbor list
+    index: (B, npoint, nsample)
+    """
+
+    def __init__(self, k=9, dilation=1, stochastic=False, epsilon=0.0):
+        super(DenseDilated, self).__init__()
+        self.dilation = dilation
+        self.stochastic = stochastic
+        self.epsilon = epsilon
+        self.k = k
+
+    def forward(self, edge_index):
+        if self.stochastic:
+            if torch.rand(1) < self.epsilon and self.training:
+                num = self.k * self.dilation
+                randnum = torch.randperm(num)[:self.k]
+                edge_index = edge_index[:, :, randnum]
+            else:
+                edge_index = edge_index[:, :, ::self.dilation]
+        else:
+            edge_index = edge_index[:, :, ::self.dilation]
+        return edge_index.contiguous()
+
+
+class DilatedKNN(nn.Module):
+    """
+    Find the neighbors' indices based on dilated knn
+    """
+
+    def __init__(self, k=9, dilation=1, stochastic=False, epsilon=0.0):
+        super(DilatedKNN, self).__init__()
+        self.dilation = dilation
+        self.stochastic = stochastic
+        self.epsilon = epsilon
+        self.k = k
+        self._dilated = DenseDilated(k, dilation, stochastic, epsilon)
+        # self.knn = KNNCUDA(k * self.dilation, transpose_mode=True)
+        self.knn = KNN(k * self.dilation, transpose_mode=True)
+
+    def forward(self, query):
+        _, idx = self.knn(query, query)
+        return self._dilated(idx)
+
+