openpoints 0.1.0__py3-none-any.whl

openpoints/cpp/pointops/functions/pointops.py

@@ -0,0 +1,314 @@
+from typing import Tuple
+
+import torch
+from torch.autograd import Function
+import torch.nn as nn
+
+try:
+    import pointops_cuda
+except ImportError as exc:
+    class _MissingCudaExtension:
+        def __init__(self, module_name, error):
+            self.module_name = module_name
+            self.error = error
+
+        def __getattr__(self, name):
+            raise ImportError(
+                f"{self.module_name} is not installed. Build OpenPoints CUDA ops from "
+                "source first, e.g. `cd cpp/pointops && python setup.py install`, "
+                "or install a wheel that includes the compiled CUDA extension."
+            ) from self.error
+
+    pointops_cuda = _MissingCudaExtension("pointops_cuda", exc)
+
+
+class FurthestSampling(Function):
+    @staticmethod
+    def forward(ctx, xyz, offset, new_offset):
+        """
+        input: xyz: (n, 3), offset: (b), new_offset: (b)
+        output: idx: (m)
+        """
+        assert xyz.is_contiguous()
+        n, b, n_max = xyz.shape[0], offset.shape[0], offset[0]
+        for i in range(1, b):
+            n_max = max(offset[i] - offset[i - 1], n_max)
+        idx = torch.cuda.IntTensor(new_offset[b - 1].item()).zero_()
+        tmp = torch.cuda.FloatTensor(n).fill_(1e10)
+        pointops_cuda.furthestsampling_cuda(
+            b, n_max, xyz, offset, new_offset, tmp, idx)
+        del tmp
+        return idx
+
+
+furthestsampling = FurthestSampling.apply
+
+
+class KNNQuery(Function):
+    @staticmethod
+    def forward(ctx, nsample, xyz, new_xyz, offset, new_offset):
+        """
+        input: xyz: (n, 3), new_xyz: (m, 3), offset: (b), new_offset: (b)
+        output: idx: (m, nsample), dist2: (m, nsample)
+        """
+        if new_xyz is None:
+            new_xyz = xyz
+        assert xyz.is_contiguous() and new_xyz.is_contiguous()
+        m = new_xyz.shape[0]
+        idx = torch.cuda.IntTensor(m, nsample).zero_()
+        dist2 = torch.cuda.FloatTensor(m, nsample).zero_()
+        pointops_cuda.knnquery_cuda(
+            m, nsample, xyz, new_xyz, offset, new_offset, idx, dist2)
+        return idx, torch.sqrt(dist2)
+
+
+knnquery = KNNQuery.apply
+
+
+class BallQuery(Function):
+    @staticmethod
+    def forward(ctx, radius, nsample, xyz, new_xyz, offset, new_offset):
+        """
+        input: xyz: (n, 3), new_xyz: (m, 3), offset: (b), new_offset: (b)
+        output: idx: (m, nsample), dist2: (m, nsample)
+        """
+        if new_xyz is None:
+            new_xyz = xyz
+        assert xyz.is_contiguous() and new_xyz.is_contiguous()
+        m = new_xyz.shape[0]
+        idx = torch.cuda.IntTensor(m, nsample).zero_()
+        pointops_cuda.ballquery_cuda(
+            m, radius, nsample, xyz, new_xyz, offset, new_offset, idx)
+        return idx
+
+
+ballquery = BallQuery.apply
+
+
+class Grouping(Function):
+    @staticmethod
+    def forward(ctx, input, idx):
+        """
+        input: input: (n, c), idx : (m, nsample)
+        output: (m, nsample, c)
+        """
+        assert input.is_contiguous() and idx.is_contiguous()
+        m, nsample, n, c = idx.shape[0], idx.shape[1], input.shape[0], input.shape[1]
+        output = torch.cuda.FloatTensor(m, nsample, c)
+        pointops_cuda.grouping_forward_cuda(m, nsample, c, input, idx, output)
+        ctx.n = n
+        ctx.save_for_backward(idx)
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        """
+        input: grad_out: (m, c, nsample)
+        output: (n, c), None
+        """
+        n = ctx.n
+        idx, = ctx.saved_tensors
+        m, nsample, c = grad_output.shape
+        grad_input = torch.cuda.FloatTensor(n, c).zero_()
+        pointops_cuda.grouping_backward_cuda(
+            m, nsample, c, grad_output, idx, grad_input)
+        return grad_input, None
+
+
+grouping = Grouping.apply
+
+
+def querygroup(nsample, xyz, new_xyz, feat, offset, new_offset,
+               radius=None, query_method='knn',
+               normalize_dp=False, idx=None
+               ):
+    """
+    the query and group function 
+    support KNN and Ball Query
+    input: xyz: (n, 3), new_xyz: (m, 3), feat: (n, c), offset: (b), new_offset: (b), idx: (m, nsample) 
+    output: 
+    """
+    assert xyz.is_contiguous() and new_xyz.is_contiguous() and feat.is_contiguous()
+    if new_xyz is None:
+        new_xyz = xyz
+    if idx is None:
+        if nsample is not None:
+            if query_method in ['knn', 'knnquery']:
+                idx, _ = knnquery(nsample, xyz, new_xyz,
+                                  offset, new_offset)  # (m, nsample)
+            else:
+                idx = ballquery(radius, nsample, xyz,
+                                new_xyz, offset, new_offset)  # (m, nsample)
+            idx = idx.flatten().long()
+            n, m, c = xyz.shape[0], new_xyz.shape[0], feat.shape[1]
+            grouped_xyz = xyz[idx, :].view(
+                m, nsample, 3)  # (m, nsample, 3)
+            grouped_xyz -= new_xyz.unsqueeze(1)  # (m, nsample, 3)
+            if normalize_dp:
+                max_dist = grouped_xyz.norm(dim=-1, p=2, keepdim=True).max(dim=-1, keepdim=True)[
+                    0] + 1.0e-8 if query_method == 'knn' else radius
+                grouped_xyz /= max_dist
+
+            if feat is not None:
+                grouped_feat = feat[idx, :].view(
+                    m, nsample, c)  # (m, nsample, c)
+            else:
+                grouped_feat = None
+        else:
+            grouped_xyz = xyz.transpose(1, 2).unsqueeze(2)
+            if feat is not None:
+                grouped_feat = feat.unsqueeze(2)
+            else:
+                grouped_feat = None
+        return grouped_xyz, grouped_feat
+
+
+def queryandgroup(nsample, xyz, new_xyz, feat, idx, offset, new_offset, use_xyz=True):
+    """
+    input: xyz: (n, 3), new_xyz: (m, 3), feat: (n, c), idx: (m, nsample), offset: (b), new_offset: (b)
+    output: new_feat: (m, c+3, nsample), grouped_idx: (m, nsample)
+    """
+    assert xyz.is_contiguous() and new_xyz.is_contiguous() and feat.is_contiguous()
+    if new_xyz is None:
+        new_xyz = xyz
+    if idx is None:
+        idx, _ = knnquery(nsample, xyz, new_xyz, offset,
+                          new_offset)  # (m, nsample)
+
+    n, m, c = xyz.shape[0], new_xyz.shape[0], feat.shape[1]
+    grouped_xyz = xyz[idx.view(-1).long(), :].view(m,
+                                                   nsample, 3)  # (m, nsample, 3)
+    grouped_xyz -= new_xyz.unsqueeze(1)  # (m, nsample, 3)
+    grouped_feat = feat[idx.view(-1).long(), :].view(m, nsample, c)
+
+    if use_xyz:
+        return torch.cat((grouped_xyz, grouped_feat), -1)  # (m, nsample, 3+c)
+    else:
+        return grouped_feat
+
+
+class Subtraction(Function):
+    @staticmethod
+    def forward(ctx, input1, input2, idx):
+        """
+        input: input1: (n, c), input2: (n, c), idx: (n, nsample)
+        output:  (n, nsample, c)
+        """
+        assert input1.is_contiguous() and input2.is_contiguous()
+        n, c = input1.shape
+        nsample = idx.shape[-1]
+        output = torch.cuda.FloatTensor(n, nsample, c).zero_()
+        pointops_cuda.subtraction_forward_cuda(
+            n, nsample, c, input1, input2, idx, output)
+        ctx.save_for_backward(idx)
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        """
+        input: grad_out: (n, nsample, c)
+        output: grad_input1: (n, c), grad_input2: (n, c)
+        """
+        idx, = ctx.saved_tensors
+        n, nsample, c = grad_output.shape
+        grad_input1 = torch.cuda.FloatTensor(n, c).zero_()
+        grad_input2 = torch.cuda.FloatTensor(n, c).zero_()
+        pointops_cuda.subtraction_backward_cuda(
+            n, nsample, c, idx, grad_output, grad_input1, grad_input2)
+        return grad_input1, grad_input2, None
+
+
+subtraction = Subtraction.apply
+
+
+class Aggregation(Function):
+    @staticmethod
+    def forward(ctx, input, position, weight, idx):
+        """
+        input: input: (n, c), position: (n, nsample, c), weight : (n, nsample, c'), idx: (n, nsample)
+        output: (n, c)
+        """
+        assert input.is_contiguous() and position.is_contiguous() and weight.is_contiguous()
+        n, nsample, c = position.shape
+        w_c = weight.shape[-1]
+        output = torch.cuda.FloatTensor(n, c).zero_()
+        pointops_cuda.aggregation_forward_cuda(
+            n, nsample, c, w_c, input, position, weight, idx, output)
+        ctx.save_for_backward(input, position, weight, idx)
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        """
+        input: grad_out: (n, c)
+        output: grad_input: (n, c), grad_position: (n, nsample, c), grad_weight : (n, nsample, c')
+        """
+        input, position, weight, idx = ctx.saved_tensors
+        n, nsample, c = position.shape
+        w_c = weight.shape[-1]
+        grad_input = torch.cuda.FloatTensor(n, c).zero_()
+        grad_position = torch.cuda.FloatTensor(n, nsample, c).zero_()
+        grad_weight = torch.cuda.FloatTensor(n, nsample, w_c).zero_()
+        pointops_cuda.aggregation_backward_cuda(n, nsample, c, w_c, input, position, weight, idx, grad_output,
+                                                grad_input, grad_position, grad_weight)
+        return grad_input, grad_position, grad_weight, None
+
+
+aggregation = Aggregation.apply
+
+
+def interpolation(xyz, new_xyz, feat, offset, new_offset, k=3):
+    """
+    input: xyz: (m, 3), new_xyz: (n, 3), feat: (m, c), offset: (b), new_offset: (b)
+    output: (n, c)
+    """
+    assert xyz.is_contiguous() and new_xyz.is_contiguous() and feat.is_contiguous()
+    idx, dist = knnquery(k, xyz, new_xyz, offset, new_offset)  # (n, 3), (n, 3)
+    dist_recip = 1.0 / (dist + 1e-8)  # (n, 3)
+    norm = torch.sum(dist_recip, dim=1, keepdim=True)
+    weight = dist_recip / norm  # (n, 3)
+
+    new_feat = torch.cuda.FloatTensor(new_xyz.shape[0], feat.shape[1]).zero_()
+    for i in range(k):
+        new_feat += feat[idx[:, i].long(), :] * weight[:, i].unsqueeze(-1)
+    return new_feat
+
+
+class Interpolation(Function):
+    @staticmethod
+    def forward(ctx, xyz, new_xyz, input, offset, new_offset, k=3):
+        """
+        input: xyz: (m, 3), new_xyz: (n, 3), input: (m, c), offset: (b), new_offset: (b)
+        output: (n, c)
+        """
+        assert xyz.is_contiguous() and new_xyz.is_contiguous() and input.is_contiguous()
+        idx, dist = knnquery(k, xyz, new_xyz, offset,
+                             new_offset)  # (n, k), (n, k)
+        dist_recip = 1.0 / (dist + 1e-8)  # (n, k)
+        norm = torch.sum(dist_recip, dim=1, keepdim=True)
+        weight = dist_recip / norm  # (n, k)
+
+        n, c, m = new_xyz.shape[0], input.shape[1], input.shape[0]
+        output = torch.cuda.FloatTensor(n, c).zero_()
+        pointops_cuda.interpolation_forward_cuda(
+            n, c, k, input, idx, weight, output)
+        ctx.m, ctx.k = m, k
+        ctx.save_for_backward(idx, weight)
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        """
+        input: xyz: (m, 3), new_xyz: (n, 3), input: (m, c), offset: (b), new_offset: (b)
+        output: (n, c)
+        """
+        m, k = ctx.m, ctx.k
+        idx, weight = ctx.saved_tensors
+        n, c = grad_output.shape
+        grad_input = torch.cuda.FloatTensor(m, c).zero_()
+        pointops_cuda.interpolation_backward_cuda(
+            n, c, k, grad_output, idx, weight, grad_input)
+        return None, None, grad_input, None, None, None
+
+
+interpolation2 = Interpolation.apply

openpoints/cpp/pointops/setup.py

@@ -0,0 +1,37 @@
+#python3 setup.py install
+from setuptools import setup
+from torch.utils.cpp_extension import BuildExtension, CUDAExtension
+import os
+from distutils.sysconfig import get_config_vars
+
+(opt,) = get_config_vars('OPT')
+os.environ['OPT'] = " ".join(
+    flag for flag in opt.split() if flag != '-Wstrict-prototypes'
+)
+
+setup(
+    name='pointops',
+    author='Hengshuang Zhao',
+    ext_modules=[
+        CUDAExtension('pointops_cuda', [
+            'src/pointops_api.cpp',
+            'src/knnquery/knnquery_cuda.cpp',
+            'src/knnquery/knnquery_cuda_kernel.cu',
+            'src/ballquery/ballquery_cuda.cpp',
+            'src/ballquery/ballquery_cuda_kernel.cu',
+            'src/sampling/sampling_cuda.cpp',
+            'src/sampling/sampling_cuda_kernel.cu',
+            'src/grouping/grouping_cuda.cpp',
+            'src/grouping/grouping_cuda_kernel.cu',
+            'src/interpolation/interpolation_cuda.cpp',
+            'src/interpolation/interpolation_cuda_kernel.cu',
+            'src/subtraction/subtraction_cuda.cpp',
+            'src/subtraction/subtraction_cuda_kernel.cu',
+            'src/aggregation/aggregation_cuda.cpp',
+            'src/aggregation/aggregation_cuda_kernel.cu',
+            ],
+        extra_compile_args={'cxx': ['-g'], 'nvcc': ['-O2']}
+        )
+    ],
+    cmdclass={'build_ext': BuildExtension}
+)

openpoints/dataset/__init__.py

@@ -0,0 +1,10 @@
+from .data_util import get_features_by_keys, crop_pc, get_class_weights
+from .build import build_dataloader_from_cfg, build_dataset_from_cfg
+from .vis3d import vis_multi_points, vis_points
+from .modelnet import *
+from .s3dis import S3DIS, S3DISSphere
+from .shapenet import *
+from .semantic_kitti import *
+from .scanobjectnn import *
+from .shapenetpart import *
+from .scannetv2 import *

openpoints/dataset/atom3d/__init__.py

@@ -0,0 +1 @@
+from .psr import Atom2Points, AtomPSR

openpoints/dataset/atom3d/psr.py

@@ -0,0 +1,38 @@
+import os
+import numpy as np
+from atom3d.datasets import LMDBDataset
+from ..build import DATASETS
+from ...transforms import  DataTransforms
+
+prot_atoms = ['C', 'H', 'O', 'N', 'S', 'P', 'ZN', 'NA', 'FE', 'CA', 'MN', 'NI', 'CO', 'MG', 'CU', 'CL', 'SE', 'F']
+
+
+def one_of_k_encoding_unk(x, allowable_set):
+    """Converts input to 1-hot encoding given a set of allowable values. Additionally maps inputs not in the allowable set to the last element."""
+    if x not in allowable_set:
+        x = allowable_set[-1]
+    return list(map(lambda s: x == s, allowable_set))
+
+
+@DataTransforms.register_module()
+class Atom2Points(object):
+    def __call__(self, item):
+        # Transform protein into voxel grids.
+        # Apply random rotation matrix.
+        id = eval(item['id'])
+        transformed = {
+            'pos': item['atoms'][['x', 'y', 'z']].to_numpy().astype(np.float32),
+            'features': np.array([one_of_k_encoding_unk(e, prot_atoms) for e in item['atoms']['element']]).astype(np.float32).transpose(1, 0),
+            'label': np.float32(item['scores']['gdt_ts']),
+            'target': id[0],
+            'decoy': id[1],
+        }
+        return transformed
+
+
+@DATASETS.register_module()
+class AtomPSR(LMDBDataset):
+    def __init__(self, data_dir, split, transform=Atom2Points()):
+        assert split in ['train', 'val', 'test']
+        super().__init__(os.path.join(data_dir, split), transform)
+

openpoints/dataset/build.py

@@ -0,0 +1,98 @@
+"""
+Author: PointNeXt
+"""
+import numpy as np
+import torch
+from easydict import EasyDict as edict
+from openpoints.utils import registry
+from openpoints.transforms import build_transforms_from_cfg
+
+DATASETS = registry.Registry('dataset')
+
+
+def concat_collate_fn(datas):
+    """collate fn for point transformer
+    """
+    pts, feats, labels, offset, count, batches = [], [], [], [], 0, []
+    for i, data in enumerate(datas):
+        count += len(data['pos'])
+        offset.append(count)
+        pts.append(data['pos'])
+        feats.append(data['x'])
+        labels.append(data['y'])
+        batches += [i] *len(data['pos'])
+        
+    data = {'pos': torch.cat(pts), 'x': torch.cat(feats), 'y': torch.cat(labels),
+            'o': torch.IntTensor(offset), 'batch': torch.LongTensor(batches)}
+    return data
+
+
+def build_dataset_from_cfg(cfg, default_args=None):
+    """
+    Build a dataset, defined by `dataset_name`.
+    Args:
+        cfg (eDICT):
+    Returns:
+        Dataset: a constructed dataset specified by dataset_name.
+    """
+    return DATASETS.build(cfg, default_args=default_args)
+
+def worker_init_fn(worker_id):
+    np.random.seed(np.random.get_state()[1][0] + worker_id)
+
+
+def build_dataloader_from_cfg(batch_size,
+                              dataset_cfg=None,
+                              dataloader_cfg=None,
+                              datatransforms_cfg=None,
+                              split='train',
+                              distributed=True,
+                              dataset=None
+                              ):
+    if dataset is None:
+        if datatransforms_cfg is not None:
+            # in case only val or test transforms are provided. 
+            if split not in datatransforms_cfg.keys() and split in ['val', 'test']:
+                trans_split = 'val'
+            else:
+                trans_split = split
+            data_transform = build_transforms_from_cfg(trans_split, datatransforms_cfg)
+        else:
+            data_transform = None
+
+        if split not in dataset_cfg.keys() and split in ['val', 'test']:
+            dataset_split = 'test' if split == 'val' else 'val'
+        else:
+            dataset_split = split
+        split_cfg = dataset_cfg.get(dataset_split, edict())
+        if split_cfg.get('split', None) is None:    # add 'split' in dataset_split_cfg
+            split_cfg.split = split
+        split_cfg.transform = data_transform
+        dataset = build_dataset_from_cfg(dataset_cfg.common, split_cfg)
+
+    collate_fn = dataset.collate_fn if hasattr(dataset, 'collate_fn') else None
+    collate_fn = dataloader_cfg.collate_fn if dataloader_cfg.get('collate_fn', None) is not None else collate_fn
+    collate_fn = eval(collate_fn) if isinstance(collate_fn, str) else collate_fn
+
+    shuffle = split == 'train'
+    if distributed:
+        sampler = torch.utils.data.distributed.DistributedSampler(dataset, shuffle=shuffle)
+        dataloader = torch.utils.data.DataLoader(dataset,
+                                                 batch_size=batch_size,
+                                                 num_workers=int(dataloader_cfg.num_workers),
+                                                 worker_init_fn=worker_init_fn,
+                                                 drop_last=split == 'train',
+                                                 sampler=sampler,
+                                                 collate_fn=collate_fn, 
+                                                 pin_memory=True
+                                                 )
+    else:
+        dataloader = torch.utils.data.DataLoader(dataset,
+                                                 batch_size=batch_size,
+                                                 num_workers=int(dataloader_cfg.num_workers),
+                                                 worker_init_fn=worker_init_fn,
+                                                 drop_last=split == 'train',
+                                                 shuffle=shuffle,
+                                                 collate_fn=collate_fn,
+                                                 pin_memory=True)
+    return dataloader