openpoints 0.1.0__py3-none-any.whl

openpoints/models/layers/registry.py

@@ -0,0 +1,168 @@
+
+import inspect
+import warnings
+from functools import partial
+
+
+class Registry:
+    """A registry to map strings to classes.
+    Args:
+        name (str): Registry name.
+    """
+
+    def __init__(self, name):
+        self._name = name
+        self._module_dict = dict()
+
+    def __len__(self):
+        return len(self._module_dict)
+
+    def __contains__(self, key):
+        return self.get(key) is not None
+
+    def __repr__(self):
+        format_str = self.__class__.__name__ + \
+                     f'(name={self._name}, ' \
+                     f'items={self._module_dict})'
+        return format_str
+
+    @property
+    def name(self):
+        return self._name
+
+    @property
+    def module_dict(self):
+        return self._module_dict
+
+    def get(self, key):
+        """Get the registry record.
+        Args:
+            key (str): The class name in string format.
+        Returns:
+            class: The corresponding class.
+        """
+        return self._module_dict.get(key, None)
+
+    def _register_module(self, module_class, module_name=None, force=False):
+        if not inspect.isclass(module_class):
+            raise TypeError('module must be a class, '
+                            f'but got {type(module_class)}')
+
+        if module_name is None:
+            module_name = module_class.__name__
+        if not force and module_name in self._module_dict:
+            raise KeyError(f'{module_name} is already registered '
+                           f'in {self.name}')
+        self._module_dict[module_name] = module_class
+
+    def deprecated_register_module(self, cls=None, force=False):
+        warnings.warn(
+            'The old API of register_module(module, force=False) '
+            'is deprecated and will be removed, please use the new API '
+            'register_module(name=None, force=False, module=None) instead.')
+        if cls is None:
+            return partial(self.deprecated_register_module, force=force)
+        self._register_module(cls, force=force)
+        return cls
+
+    def register_module(self, name=None, force=False, module=None):
+        """Register a module.
+        A record will be added to `self._module_dict`, whose key is the class
+        name or the specified name, and value is the class itself.
+        It can be used as a decorator or a normal function.
+        Example:
+            >>> backbones = Registry('backbone')
+            >>> @backbones.register_module()
+            >>> class ResNet:
+            >>>     pass
+            >>> backbones = Registry('backbone')
+            >>> @backbones.register_module(name='mnet')
+            >>> class MobileNet:
+            >>>     pass
+            >>> backbones = Registry('backbone')
+            >>> class ResNet:
+            >>>     pass
+            >>> backbones.register_module(ResNet)
+        Args:
+            name (str | None): The module name to be registered. If not
+                specified, the class name will be used.
+            force (bool, optional): Whether to override an existing class with
+                the same name. Default: False.
+            module (type): Module class to be registered.
+        """
+        if not isinstance(force, bool):
+            raise TypeError(f'force must be a boolean, but got {type(force)}')
+        # NOTE: This is a walkaround to be compatible with the old api,
+        # while it may introduce unexpected bugs.
+        if isinstance(name, type):
+            return self.deprecated_register_module(name, force=force)
+
+        # use it as a normal method: x.register_module(module=SomeClass)
+        if module is not None:
+            self._register_module(
+                module_class=module, module_name=name, force=force)
+            return module
+
+        # raise the error ahead of time
+        if not (name is None or isinstance(name, str)):
+            raise TypeError(f'name must be a str, but got {type(name)}')
+
+        # use it as a decorator: @x.register_module()
+        def _register(cls):
+            self._register_module(
+                module_class=cls, module_name=name, force=force)
+            return cls
+
+        return _register
+
+
+def build_from_cfg(cfg, registry, default_args=None):
+    """Build a module from config dict.
+    Args:
+        cfg (dict): Config dict. It should at least contain the key "type".
+        registry (:obj:`Registry`): The registry to search the type from.
+        default_args (dict, optional): Default initialization arguments.
+    Returns:
+        object: The constructed object.
+    """
+    if not isinstance(cfg, dict):
+        raise TypeError(f'cfg must be a dict, but got {type(cfg)}')
+    if 'type' not in cfg:
+        if default_args is None or 'type' not in default_args:
+            raise KeyError(
+                '`cfg` or `default_args` must contain the key "type", '
+                f'but got {cfg}\n{default_args}')
+    if not isinstance(registry, Registry):
+        raise TypeError('registry must be an mmcv.Registry object, '
+                        f'but got {type(registry)}')
+    if not (isinstance(default_args, dict) or default_args is None):
+        raise TypeError('default_args must be a dict or None, '
+                        f'but got {type(default_args)}')
+
+    args = cfg.copy()
+
+    if default_args is not None:
+        for name, value in default_args.items():
+            args.setdefault(name, value)
+
+    obj_type = args.pop('type')  # pop the type out, then leave the args for the class.
+    if isinstance(obj_type, str):
+        obj_cls = registry.get(obj_type)
+        if obj_cls is None:
+            raise KeyError(
+                f'{obj_type} is not in the {registry.name} registry. '
+                f'obj_type should be one of {list(registry._module_dict.keys())}. '
+                f'Check if any typo and make sure the capitalization is the same. ')
+    elif inspect.isclass(obj_type):
+        obj_cls = obj_type
+    else:
+        raise TypeError(
+            f'type must be a str or valid type, but got {type(obj_type)}')
+
+    return obj_cls(**args)
+
+
+NORM_LAYERS = Registry('norm layer')
+ACTIVATION_LAYERS = Registry('activation layer')
+CONV_LAYERS = Registry('conv layer')
+DROPOUT_LAYERS = Registry('dropout layer')

openpoints/models/layers/subsample.py

@@ -0,0 +1,185 @@
+# subsample layer for 3d processing.
+from abc import ABC, abstractmethod
+
+import torch
+import torch.nn as nn
+from torch.autograd import Function
+import math
+from openpoints.cpp.pointnet2_batch import pointnet2_cuda
+
+
+class BaseSampler(ABC):
+    """If num_to_sample is provided, sample exactly
+        num_to_sample points. Otherwise sample floor(pos[0] * ratio) points
+    """
+
+    def __init__(self, ratio=None, num_to_sample=None, subsampling_param=None):
+        if num_to_sample is not None:
+            if (ratio is not None) or (subsampling_param is not None):
+                raise ValueError(
+                    "Can only specify ratio or num_to_sample or subsampling_param, not several !")
+            self._num_to_sample = num_to_sample
+
+        elif ratio is not None:
+            self._ratio = ratio
+
+        elif subsampling_param is not None:
+            self._subsampling_param = subsampling_param
+
+        else:
+            raise Exception(
+                'At least ["ratio, num_to_sample, subsampling_param"] should be defined')
+
+    def __call__(self, xyz):
+        return self.sample(xyz)
+
+    def _get_num_to_sample(self, npoints) -> int:
+        if hasattr(self, "_num_to_sample"):
+            return self._num_to_sample
+        else:
+            return math.floor(npoints * self._ratio)
+
+    def _get_ratio_to_sample(self, batch_size) -> float:
+        if hasattr(self, "_ratio"):
+            return self._ratio
+        else:
+            return self._num_to_sample / float(batch_size)
+
+    @abstractmethod
+    def sample(self, xyz, feature=None, batch=None):
+        pass
+
+
+class RandomSample(BaseSampler):
+    """Random Sample for dense data
+        Arguments:
+            xyz -- [B, N, 3]
+    """
+
+    def sample(self, xyz, **kwargs):
+        if len(xyz.shape) != 3:
+            raise ValueError(" Expects the xyz tensor to be of dimension 3")
+        B, N, _ = xyz.shape
+        idx = torch.randint(
+            0, N, (B, self._get_num_to_sample(N)), device=xyz.device)
+        sampled_xyz = torch.gather(xyz, 1, idx.unsqueeze(-1).expand(-1, -1, 3))
+        # sampled_feature = torch.gather(feature, 2, idx.unsqueeze(1).repeat(1, C, 1))
+        return sampled_xyz, idx
+
+
+def random_sample(xyz, npoint):
+    B, N, _ = xyz.shape
+    idx = torch.randint(0, N, (B, npoint), device=xyz.device)
+    return idx
+
+
+class FurthestPointSampling(Function):
+    @staticmethod
+    def forward(ctx, xyz: torch.Tensor, npoint: int) -> torch.Tensor:
+        """
+        Uses iterative furthest point sampling to select a set of npoint features that have the largest
+        minimum distance
+        :param ctx:
+        :param xyz: (B, N, 3) where N > npoint
+        :param npoint: int, number of features in the sampled set
+        :return:
+             output: (B, npoint) tensor containing the set (idx)
+        """
+        assert xyz.is_contiguous()
+
+        B, N, _ = xyz.size()
+        # output = torch.cuda.IntTensor(B, npoint, device=xyz.device)
+        # temp = torch.cuda.FloatTensor(B, N, device=xyz.device).fill_(1e10)
+        output = torch.cuda.IntTensor(B, npoint)
+        temp = torch.cuda.FloatTensor(B, N).fill_(1e10)
+
+        pointnet2_cuda.furthest_point_sampling_wrapper(
+            B, N, npoint, xyz, temp, output)
+        return output
+
+    @staticmethod
+    def backward(xyz, a=None):
+        return None, None
+
+
+furthest_point_sample = FurthestPointSampling.apply
+
+
+class GatherOperation(Function):
+    @staticmethod
+    def forward(ctx, features: torch.Tensor, idx: torch.Tensor) -> torch.Tensor:
+        """
+        :param ctx:
+        :param features: (B, C, N)
+        :param idx: (B, npoint) index tensor of the features to gather
+        :return:
+            output: (B, C, npoint)
+        """
+        assert features.is_contiguous()
+        assert idx.is_contiguous()
+
+        B, npoint = idx.size()
+        _, C, N = features.size()
+        output = torch.cuda.FloatTensor(B, C, npoint, device=features.device)
+
+        pointnet2_cuda.gather_points_wrapper(
+            B, C, N, npoint, features, idx, output)
+
+        ctx.for_backwards = (idx, C, N)
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_out):
+        idx, C, N = ctx.for_backwards
+        B, npoint = idx.size()
+
+        grad_features = torch.zeros(
+            [B, C, N], dtype=torch.float, device=grad_out.device, requires_grad=True)
+        grad_out_data = grad_out.data.contiguous()
+        pointnet2_cuda.gather_points_grad_wrapper(
+            B, C, N, npoint, grad_out_data, idx, grad_features.data)
+        return grad_features, None
+
+
+gather_operation = GatherOperation.apply
+# mark: torch gather is even faster. sampled_xyz = torch.gather(points, 1, idx.unsqueeze(-1).expand(-1, -1, 3))
+
+
+def fps(data, number):
+    '''
+        data B N C
+        number int
+    '''
+    fps_idx = furthest_point_sample(data[:, :, :3].contiguous(), number)
+    fps_data = torch.gather(
+        data, 1, fps_idx.unsqueeze(-1).long().expand(-1, -1, data.shape[-1]))
+    return fps_data
+
+
+if __name__ == '__main__':
+    import time
+
+    B, C, N = 2, 3, 10000
+    K = 16
+    device = 'cuda'
+    points = torch.randn([B, N, 3], device=device, dtype=torch.float)
+    print(points.shape, '\n', points)
+
+    nsample = 4096
+    idx = furthest_point_sample(points, nsample)
+
+    st = time.time()
+    for _ in range(100):
+        query1 = torch.gather(
+            points, 1, idx.long().unsqueeze(-1).expand(-1, -1, 3))
+    print(time.time() - st)
+    print(query1.shape)
+
+    st = time.time()
+    for _ in range(100):
+        query2 = gather_operation(points.transpose(
+            1, 2).contiguous(), idx).transpose(1, 2).contiguous()
+    print(time.time() - st)
+    print(query2.shape)
+
+    print(torch.allclose(query1, query2))

openpoints/models/layers/upsampling.py

@@ -0,0 +1,106 @@
+from typing import List, Tuple
+from torch.autograd import Function
+
+import torch
+import torch.nn as nn
+
+from openpoints.cpp.pointnet2_batch import pointnet2_cuda
+from openpoints.models.layers import create_convblock1d
+
+
+class ThreeNN(Function):
+
+    @staticmethod
+    def forward(ctx, unknown: torch.Tensor, known: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Find the three nearest neighbors of unknown in known
+        :param ctx:
+        :param unknown: (B, N, 3)
+        :param known: (B, M, 3)
+        :return:
+            dist: (B, N, 3) l2 distance to the three nearest neighbors
+            idx: (B, N, 3) index of 3 nearest neighbors
+        """
+        assert unknown.is_contiguous()
+        assert known.is_contiguous()
+
+        B, N, _ = unknown.size()
+        m = known.size(1)
+        dist2 = torch.cuda.FloatTensor(B, N, 3)
+        idx = torch.cuda.IntTensor(B, N, 3)
+
+        pointnet2_cuda.three_nn_wrapper(B, N, m, unknown, known, dist2, idx)
+        return torch.sqrt(dist2), idx
+
+    @staticmethod
+    def backward(ctx, a=None, b=None):
+        return None, None
+
+
+three_nn = ThreeNN.apply
+
+
+class ThreeInterpolate(Function):
+
+    @staticmethod
+    @torch.cuda.amp.custom_fwd(cast_inputs=torch.float32)
+    def forward(ctx, features: torch.Tensor, idx: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
+        """
+        Performs weight linear interpolation on 3 features
+        :param ctx:
+        :param features: (B, C, M) Features descriptors to be interpolated from
+        :param idx: (B, n, 3) three nearest neighbors of the target features in features
+        :param weight: (B, n, 3) weights
+        :return:
+            output: (B, C, N) tensor of the interpolated features
+        """
+        assert features.is_contiguous()
+        assert idx.is_contiguous()
+        assert weight.is_contiguous()
+
+        B, c, m = features.size()
+        n = idx.size(1)
+        ctx.three_interpolate_for_backward = (idx, weight, m)
+        output = torch.cuda.FloatTensor(B, c, n)
+
+        pointnet2_cuda.three_interpolate_wrapper(B, c, m, n, features, idx, weight, output)
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_out: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """
+        :param ctx:
+        :param grad_out: (B, C, N) tensor with gradients of outputs
+        :return:
+            grad_features: (B, C, M) tensor with gradients of features
+            None:
+            None:
+        """
+        idx, weight, m = ctx.three_interpolate_for_backward
+        B, c, n = grad_out.size()
+
+        grad_features = torch.zeros([B, c, m], device='cuda', requires_grad=True)
+        grad_out_data = grad_out.data.contiguous()
+
+        pointnet2_cuda.three_interpolate_grad_wrapper(B, c, n, m, grad_out_data, idx, weight, grad_features.data)
+        return grad_features, None, None
+
+
+three_interpolate = ThreeInterpolate.apply
+
+
+def three_interpolation(unknown_xyz, known_xyz, know_feat):
+    """
+    input: known_xyz: (m, 3), unknown_xyz: (n, 3), feat: (m, c), offset: (b), new_offset: (b)
+    output: (n, c)
+    """
+    dist, idx = three_nn(unknown_xyz, known_xyz)
+    dist_recip = 1.0 / (dist + 1e-8)
+    norm = torch.sum(dist_recip, dim=2, keepdim=True)
+    weight = dist_recip / norm
+    interpolated_feats = three_interpolate(know_feat, idx, weight)
+    return interpolated_feats
+
+
+if __name__ == "__main__":
+    pass

openpoints/models/layers/weight_init.py

@@ -0,0 +1,89 @@
+import torch
+import math
+import warnings
+
+from torch.nn.init import _calculate_fan_in_and_fan_out
+
+
+def _no_grad_trunc_normal_(tensor, mean, std, a, b):
+    # Cut & paste from PyTorch official master until it's in a few official releases - RW
+    # Method based on https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf
+    def norm_cdf(x):
+        # Computes standard normal cumulative distribution function
+        return (1. + math.erf(x / math.sqrt(2.))) / 2.
+
+    if (mean < a - 2 * std) or (mean > b + 2 * std):
+        warnings.warn("mean is more than 2 std from [a, b] in nn.init.trunc_normal_. "
+                      "The distribution of values may be incorrect.",
+                      stacklevel=2)
+
+    with torch.no_grad():
+        # Values are generated by using a truncated uniform distribution and
+        # then using the inverse CDF for the normal distribution.
+        # Get upper and lower cdf values
+        l = norm_cdf((a - mean) / std)
+        u = norm_cdf((b - mean) / std)
+
+        # Uniformly fill tensor with values from [l, u], then translate to
+        # [2l-1, 2u-1].
+        tensor.uniform_(2 * l - 1, 2 * u - 1)
+
+        # Use inverse cdf transform for normal distribution to get truncated
+        # standard normal
+        tensor.erfinv_()
+
+        # Transform to proper mean, std
+        tensor.mul_(std * math.sqrt(2.))
+        tensor.add_(mean)
+
+        # Clamp to ensure it's in the proper range
+        tensor.clamp_(min=a, max=b)
+        return tensor
+
+
+def trunc_normal_(tensor, mean=0., std=1., a=-2., b=2.):
+    # type: (Tensor, float, float, float, float) -> Tensor
+    r"""Fills the input Tensor with values drawn from a truncated
+    normal distribution. The values are effectively drawn from the
+    normal distribution :math:`\mathcal{N}(\text{mean}, \text{std}^2)`
+    with values outside :math:`[a, b]` redrawn until they are within
+    the bounds. The method used for generating the random values works
+    best when :math:`a \leq \text{mean} \leq b`.
+    Args:
+        tensor: an n-dimensional `torch.Tensor`
+        mean: the mean of the normal distribution
+        std: the standard deviation of the normal distribution
+        a: the minimum cutoff value
+        b: the maximum cutoff value
+    Examples:
+        >>> w = torch.empty(3, 5)
+        >>> nn.init.trunc_normal_(w)
+    """
+    return _no_grad_trunc_normal_(tensor, mean, std, a, b)
+
+
+def variance_scaling_(tensor, scale=1.0, mode='fan_in', distribution='normal'):
+    fan_in, fan_out = _calculate_fan_in_and_fan_out(tensor)
+    if mode == 'fan_in':
+        denom = fan_in
+    elif mode == 'fan_out':
+        denom = fan_out
+    elif mode == 'fan_avg':
+        denom = (fan_in + fan_out) / 2
+
+    variance = scale / denom
+
+    if distribution == "truncated_normal":
+        # constant is stddev of standard normal truncated to (-2, 2)
+        trunc_normal_(tensor, std=math.sqrt(variance) / .87962566103423978)
+    elif distribution == "normal":
+        tensor.normal_(std=math.sqrt(variance))
+    elif distribution == "uniform":
+        bound = math.sqrt(3 * variance)
+        tensor.uniform_(-bound, bound)
+    else:
+        raise ValueError(f"invalid distribution {distribution}")
+
+
+def lecun_normal_(tensor):
+    variance_scaling_(tensor, mode='fan_in', distribution='truncated_normal')

openpoints/models/reconstruction/__init__.py

@@ -0,0 +1,8 @@
+"""
+Author: PointNeXt
+
+"""
+from .base_recontruct import MaskedTransformerDecoder, FoldingNet, NodeShuffle
+from .maskedpointvit import MaskedPointViT
+from .maskedpoint import MaskedPoint
+from .maskedpointgroup import MaskedPointGroup