openpoints 0.1.0__py3-none-any.whl

openpoints/__init__.py

@@ -0,0 +1,3 @@
+__version__ = "0.1.0"
+
+from .transforms import *

openpoints/cpp/__init__.py

@@ -0,0 +1,6 @@
+"""
+Author: PointNeXt
+
+"""
+
+from .pointnet2_batch import pointnet2_cuda

openpoints/cpp/chamfer_dist/__init__.py

@@ -0,0 +1,110 @@
+# -*- coding: utf-8 -*-
+# @Author: Thibault GROUEIX
+# @Date:   2019-08-07 20:54:24
+# @Last Modified by:   Haozhe Xie
+# @Last Modified time: 2019-12-18 15:06:25
+# @Email:  cshzxie@gmail.com
+"""Chamfer distance CUDA extension loader.
+
+The PyPI package ships Python sources and keeps compiled CUDA extensions
+optional. Importing this module should not fail on machines where the
+``chamfer`` extension has not been built yet; users only need the extension
+when they instantiate/use Chamfer distance losses.
+"""
+
+import torch
+
+try:
+    import chamfer as _chamfer
+except ImportError as exc:
+    _chamfer = None
+    _chamfer_import_error = exc
+else:
+    _chamfer_import_error = None
+
+
+def _require_chamfer():
+    if _chamfer is None:
+        raise ImportError(
+            "The `chamfer` CUDA extension is not installed. Build OpenPoints "
+            "CUDA ops from source first, e.g. `cd cpp/chamfer_dist && "
+            "python setup.py install`, or install a wheel that includes the "
+            "compiled extension."
+        ) from _chamfer_import_error
+    return _chamfer
+
+
+class ChamferFunction(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, xyz1, xyz2):
+        chamfer = _require_chamfer()
+        dist1, dist2, idx1, idx2 = chamfer.forward(xyz1, xyz2)
+        ctx.save_for_backward(xyz1, xyz2, idx1, idx2)
+
+        return dist1, dist2
+
+    @staticmethod
+    def backward(ctx, grad_dist1, grad_dist2):
+        chamfer = _require_chamfer()
+        xyz1, xyz2, idx1, idx2 = ctx.saved_tensors
+        grad_xyz1, grad_xyz2 = chamfer.backward(xyz1, xyz2, idx1, idx2, grad_dist1, grad_dist2)
+        return grad_xyz1, grad_xyz2
+
+
+class ChamferDistanceL2(torch.nn.Module):
+    f''' Chamder Distance L2
+    '''
+    def __init__(self, ignore_zeros=False):
+        super().__init__()
+        self.ignore_zeros = ignore_zeros
+
+    def forward(self, xyz1, xyz2):
+        batch_size = xyz1.size(0)
+        if batch_size == 1 and self.ignore_zeros:
+            non_zeros1 = torch.sum(xyz1, dim=2).ne(0)
+            non_zeros2 = torch.sum(xyz2, dim=2).ne(0)
+            xyz1 = xyz1[non_zeros1].unsqueeze(dim=0)
+            xyz2 = xyz2[non_zeros2].unsqueeze(dim=0)
+
+        dist1, dist2 = ChamferFunction.apply(xyz1, xyz2)
+        return torch.mean(dist1) + torch.mean(dist2)
+
+class ChamferDistanceL2_split(torch.nn.Module):
+    f''' Chamder Distance L2
+    '''
+    def __init__(self, ignore_zeros=False):
+        super().__init__()
+        self.ignore_zeros = ignore_zeros
+
+    def forward(self, xyz1, xyz2):
+        batch_size = xyz1.size(0)
+        if batch_size == 1 and self.ignore_zeros:
+            non_zeros1 = torch.sum(xyz1, dim=2).ne(0)
+            non_zeros2 = torch.sum(xyz2, dim=2).ne(0)
+            xyz1 = xyz1[non_zeros1].unsqueeze(dim=0)
+            xyz2 = xyz2[non_zeros2].unsqueeze(dim=0)
+
+        dist1, dist2 = ChamferFunction.apply(xyz1, xyz2)
+        return torch.mean(dist1), torch.mean(dist2)
+
+class ChamferDistanceL1(torch.nn.Module):
+    f''' Chamder Distance L1
+    '''
+    def __init__(self, ignore_zeros=False):
+        super().__init__()
+        self.ignore_zeros = ignore_zeros
+
+    def forward(self, xyz1, xyz2):
+        batch_size = xyz1.size(0)
+        if batch_size == 1 and self.ignore_zeros:
+            non_zeros1 = torch.sum(xyz1, dim=2).ne(0)
+            non_zeros2 = torch.sum(xyz2, dim=2).ne(0)
+            xyz1 = xyz1[non_zeros1].unsqueeze(dim=0)
+            xyz2 = xyz2[non_zeros2].unsqueeze(dim=0)
+
+        dist1, dist2 = ChamferFunction.apply(xyz1, xyz2)
+        # import pdb
+        # pdb.set_trace()
+        dist1 = torch.sqrt(dist1)
+        dist2 = torch.sqrt(dist2)
+        return (torch.mean(dist1) + torch.mean(dist2))/2

openpoints/cpp/chamfer_dist/setup.py

@@ -0,0 +1,19 @@
+# -*- coding: utf-8 -*-
+# @Author: Haozhe Xie
+# @Date:   2019-08-07 20:54:24
+# @Last Modified by:   Haozhe Xie
+# @Last Modified time: 2019-12-10 10:04:25
+# @Email:  cshzxie@gmail.com
+
+from setuptools import setup
+from torch.utils.cpp_extension import BuildExtension, CUDAExtension
+
+setup(name='chamfer',
+      version='2.0.0',
+      ext_modules=[
+          CUDAExtension('chamfer', [
+              'chamfer_cuda.cpp',
+              'chamfer.cu',
+          ]),
+      ],
+      cmdclass={'build_ext': BuildExtension})

openpoints/cpp/chamfer_dist/test.py

@@ -0,0 +1,34 @@
+# -*- coding: utf-8 -*-
+# @Author: Haozhe Xie
+# @Date:   2019-12-10 10:38:01
+# @Last Modified by:   Haozhe Xie
+# @Last Modified time: 2019-12-26 14:21:36
+# @Email:  cshzxie@gmail.com
+#
+# Note:
+# - Replace float -> double, kFloat -> kDouble in chamfer.cu
+
+import os
+import sys
+import torch
+import unittest
+
+from torch.autograd import gradcheck
+
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), os.path.pardir, os.path.pardir)))
+from openpoints.cpp.chamfer_dist import ChamferFunction
+
+
+class ChamferDistanceTestCase(unittest.TestCase):
+    def test_chamfer_dist(self):
+        x = torch.rand(4, 64, 3).float()
+        y = torch.rand(4, 128, 3).float()
+        x.requires_grad = True
+        y.requires_grad = True
+        print(gradcheck(ChamferFunction.apply, [x.cuda(), y.cuda()]))
+
+
+if __name__ == '__main__':
+    x = torch.rand(32, 128, 3)
+    y = torch.rand(32, 128, 3)
+    # print()

openpoints/cpp/emd/__init__.py

@@ -0,0 +1,3 @@
+from .emd import earth_mover_distance as emd
+
+__all__ = ['emd']

openpoints/cpp/emd/emd.py

@@ -0,0 +1,88 @@
+import torch
+
+try:
+    import emd_cuda as _emd_cuda
+except ImportError as exc:
+    _emd_cuda = None
+    _emd_cuda_import_error = exc
+else:
+    _emd_cuda_import_error = None
+
+
+def _require_emd_cuda():
+    if _emd_cuda is None:
+        raise ImportError(
+            "The `emd_cuda` extension is not installed. Build OpenPoints CUDA "
+            "ops from source first, e.g. `cd cpp/emd && python setup.py "
+            "install`, or install a wheel that includes the compiled extension."
+        ) from _emd_cuda_import_error
+    return _emd_cuda
+
+
+class EarthMoverDistanceFunction(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, xyz1, xyz2):
+        emd_cuda = _require_emd_cuda()
+        xyz1 = xyz1.contiguous()
+        xyz2 = xyz2.contiguous()
+        assert xyz1.is_cuda and xyz2.is_cuda, "Only support cuda currently."
+        match = emd_cuda.approxmatch_forward(xyz1, xyz2)
+        cost = emd_cuda.matchcost_forward(xyz1, xyz2, match)
+        ctx.save_for_backward(xyz1, xyz2, match)
+        return cost
+
+    @staticmethod
+    def backward(ctx, grad_cost):
+        emd_cuda = _require_emd_cuda()
+        xyz1, xyz2, match = ctx.saved_tensors
+        grad_cost = grad_cost.contiguous()
+        grad_xyz1, grad_xyz2 = emd_cuda.matchcost_backward(grad_cost, xyz1, xyz2, match)
+        return grad_xyz1, grad_xyz2
+
+
+class earth_mover_distance(torch.nn.Module):
+    f''' emd
+    '''
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, xyz1, xyz2, transpose=False):
+        """Earth Mover Distance (Approx)
+
+        Args:
+            xyz1 (torch.Tensor): (b, n1, 3)
+            xyz2 (torch.Tensor): (b, n2, 3)
+            transpose (bool): whether to transpose inputs as it might be BCN format.
+                Extensions only support BNC format.
+
+        Returns:
+            cost (torch.Tensor): (b)
+
+        """
+
+        cost = EarthMoverDistanceFunction.apply(xyz1, xyz2)
+        cost = cost / xyz1.size(1)
+        
+        return cost.mean()
+# def earth_mover_distance(xyz1, xyz2, transpose=True):
+#     """Earth Mover Distance (Approx)
+
+#     Args:
+#         xyz1 (torch.Tensor): (b, 3, n1)
+#         xyz2 (torch.Tensor): (b, 3, n1)
+#         transpose (bool): whether to transpose inputs as it might be BCN format.
+#             Extensions only support BNC format.
+
+#     Returns:
+#         cost (torch.Tensor): (b)
+
+#     """
+#     if xyz1.dim() == 2:
+#         xyz1 = xyz1.unsqueeze(0)
+#     if xyz2.dim() == 2:
+#         xyz2 = xyz2.unsqueeze(0)
+#     if transpose:
+#         xyz1 = xyz1.transpose(1, 2)
+#         xyz2 = xyz2.transpose(1, 2)
+#     cost = EarthMoverDistanceFunction.apply(xyz1, xyz2)
+#     return cost

openpoints/cpp/emd/setup.py

@@ -0,0 +1,27 @@
+"""Setup extension
+
+Notes:
+    If extra_compile_args is provided, you need to provide different instances for different extensions.
+    Refer to https://github.com/pytorch/pytorch/issues/20169
+
+"""
+
+from setuptools import setup
+from torch.utils.cpp_extension import BuildExtension, CUDAExtension
+
+
+setup(
+    name='emd_ext',
+    ext_modules=[
+        CUDAExtension(
+            name='emd_cuda',
+            sources=[
+                'cuda/emd.cpp',
+                'cuda/emd_kernel.cu',
+            ],
+            extra_compile_args={'cxx': ['-g'], 'nvcc': ['-O2']}
+        ),
+    ],
+    cmdclass={
+        'build_ext': BuildExtension
+    })

openpoints/cpp/emd/test_emd_loss.py

@@ -0,0 +1,45 @@
+import torch
+import numpy as np
+import time
+from emd import earth_mover_distance
+
+# gt
+p1 = torch.from_numpy(np.array([[[1.7, -0.1, 0.1], [0.1, 1.2, 0.3]]], dtype=np.float32)).cuda()
+p1 = p1.repeat(3, 1, 1)
+p2 = torch.from_numpy(np.array([[[0.3, 1.8, 0.2], [1.2, -0.2, 0.3]]], dtype=np.float32)).cuda()
+p2 = p2.repeat(3, 1, 1)
+print(p1)
+print(p2)
+print(p1.shape)
+p1.requires_grad = True
+p2.requires_grad = True
+
+gt_dist = (((p1[0, 0] - p2[0, 1])**2).sum() + ((p1[0, 1] - p2[0, 0])**2).sum()) / 2 +  \
+         (((p1[1, 0] - p2[1, 1])**2).sum() + ((p1[1, 1] - p2[1, 0])**2).sum()) * 2 + \
+         (((p1[2, 0] - p2[2, 1])**2).sum() + ((p1[2, 1] - p2[2, 0])**2).sum()) / 3
+print('gt_dist: ', gt_dist)
+
+gt_dist.backward()
+print(p1.grad)
+print(p2.grad)
+
+# emd
+p1 = torch.from_numpy(np.array([[[1.7, -0.1, 0.1], [0.1, 1.2, 0.3]]], dtype=np.float32)).cuda()
+p1 = p1.repeat(3, 1, 1)
+p2 = torch.from_numpy(np.array([[[0.3, 1.8, 0.2], [1.2, -0.2, 0.3]]], dtype=np.float32)).cuda()
+p2 = p2.repeat(3, 1, 1)
+print(p1)
+print(p2)
+p1.requires_grad = True
+p2.requires_grad = True
+
+d = earth_mover_distance(p1, p2, transpose=False)
+print(d)
+
+loss = d[0] / 2 + d[1] * 2 + d[2] / 3
+print(loss)
+
+loss.backward()
+print(p1.grad)
+print(p2.grad)
+

openpoints/cpp/pointnet2_batch/__init__.py

@@ -0,0 +1,24 @@
+"""pointnet2_batch CUDA extension loader.
+
+The PyPI wheel ships Python sources only. Build this extension from
+``cpp/pointnet2_batch`` for full CUDA training/evaluation.
+"""
+
+
+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/pointnet2_batch && python setup.py install`, "
+            "or install a wheel that includes the compiled CUDA extension."
+        ) from self.error
+
+
+try:
+    import pointnet2_batch_cuda as pointnet2_cuda
+except ImportError as exc:
+    pointnet2_cuda = _MissingCudaExtension("pointnet2_batch_cuda", exc)

openpoints/cpp/pointnet2_batch/setup.py

@@ -0,0 +1,23 @@
+from setuptools import setup
+from torch.utils.cpp_extension import BuildExtension, CUDAExtension
+
+setup(
+    name='pointnet2_cuda',
+    ext_modules=[
+        CUDAExtension('pointnet2_batch_cuda', [
+        'src/pointnet2_api.cpp',
+        'src/ball_query.cpp',
+        'src/ball_query_gpu.cu',
+        'src/group_points.cpp',
+        'src/group_points_gpu.cu',
+        'src/interpolate.cpp',
+        'src/interpolate_gpu.cu',
+        'src/sampling.cpp',
+        'src/sampling_gpu.cu',
+        ],
+        extra_compile_args={'cxx': ['-g'],
+                            'nvcc': ['-O2']})
+    ],
+    cmdclass={'build_ext': BuildExtension}
+)
+