Updates and some refactoring

utils/CMakeLists.txt

@@ -1,7 +1,7 @@
 project(PointNet2)
 cmake_minimum_required(VERSION 3.5)
 
-find_package(CUDA)
+find_package(CUDA REQUIRED)
 
 include_directories("${CMAKE_SOURCE_DIR}/cinclude")
 cuda_include_directories("${CMAKE_SOURCE_DIR}/cinclude")
@@ -9,11 +9,13 @@ file(GLOB cuda_kernels_src "csrc/*.cu")
 cuda_compile(cuda_kernels SHARED ${cuda_kernels_src} OPTIONS -O3)
 
 file(GLOB wrapper_headers "cinclude/*wrapper.h")
+file(GLOB wrapper_sources "csrs/*.c")
 add_custom_command(OUTPUT "${CMAKE_SOURCE_DIR}/_ext/__ext.so"
 		   WORKING_DIRECTORY ${CMAKE_SOURCE_DIR}
 		   COMMAND python "${CMAKE_SOURCE_DIR}/build_ffi.py" ${cuda_kernels}
 		   DEPENDS ${cuda_kernels}
 		   DEPENDS ${wrapper_headers}
+		   DEPENDS ${wrapper_sources}
 		   VERBATIM)
 
 add_custom_target(ext ALL

utils/build_ffi.py

@@ -16,7 +16,8 @@ ffi = create_extension(
     relative_to=__file__,
     with_cuda=True,
     extra_objects=extra_objects,
-    include_dirs=[path.join(base_dir, 'cinclude')])
+    include_dirs=[path.join(base_dir, 'cinclude')],
+    verbose=False)
 
 if __name__ == "__main__":
     assert torch.cuda.is_available(), "Needs CUDA!"

utils/data_utils.py

@@ -16,7 +16,7 @@ class PointcloudScale(object):
 class PointcloudRotate(object):
     def __init__(self, x_axis=False, z_axis=True):
         assert x_axis or z_axis
-        self.x, self.y = x_axis, z_axis
+        self.x, self.z = x_axis, z_axis
 
     def _get_angles(self):
         rotation_angle = np.random.uniform() * 2 * np.pi

utils/pointnet2_modules.py

@@ -142,7 +142,7 @@ class PointnetSAModule(nn.Module):
             new_xyz = pointnet2_utils.gather_points(
                 xyz, pointnet2_utils.furthest_point_sample(xyz, self.npoint))
         else:
-            new_xyz = xyz.new([[[0, 0, 0]]]).expand(xyz.size(0), 1, 3)
+            new_xyz = xyz.data.new([[[0, 0, 0]]]).expand(xyz.size(0), 1, 3)
 
         new_points = self.grouper(xyz, new_xyz,
                                   points)  # (B, npoint, nsample, 3 + C)