Reformat code with black.

Author: hzxie

torch_points_kernels/cubic_feature_sampling.py

@@ -9,16 +9,12 @@ class CubicFeatureSamplingFunction(torch.autograd.Function):
     def forward(ctx, ptcloud, cubic_features, neighborhood_size=1):
         scale = cubic_features.size(2)
         if not torch.cuda.is_available():
-            raise NotImplementedError(
-                "CPU version is not available for Cubic Feature Sampling"
-            )
+            raise NotImplementedError("CPU version is not available for Cubic Feature Sampling")
 
         point_features, grid_pt_indexes = tpcuda.cubic_feature_sampling(
             scale, neighborhood_size, ptcloud, cubic_features
         )
-        ctx.save_for_backward(
-            torch.Tensor([scale]), torch.Tensor([neighborhood_size]), grid_pt_indexes
-        )
+        ctx.save_for_backward(torch.Tensor([scale]), torch.Tensor([neighborhood_size]), grid_pt_indexes)
         return point_features
 
     @staticmethod
@@ -56,10 +52,8 @@ def cubic_feature_sampling(ptcloud, cubic_features, neighborhood_size=1):
         (B, n_pts, n_vertices, c), where n_vertices = (neighborhood_size * 2)^3
     """
     if len(ptcloud.shape) != 3 or ptcloud.shape[2] != 3:
-        raise ValueError('The input point cloud should be of size (B, n_pts, 3).')
+        raise ValueError("The input point cloud should be of size (B, n_pts, 3).")
 
     h_scale = cubic_features.size(2) / 2
     ptcloud = ptcloud * h_scale + h_scale
-    return CubicFeatureSamplingFunction.apply(
-        ptcloud, cubic_features, neighborhood_size
-    )
+    return CubicFeatureSamplingFunction.apply(ptcloud, cubic_features, neighborhood_size)

torch_points_kernels/torchpoints.py

@@ -30,10 +30,7 @@ def furthest_point_sample(xyz, npoint):
         (B, npoint) tensor containing the set
     """
     if npoint > xyz.shape[1]:
-        raise ValueError(
-            "caanot sample %i points from an input set of %i points"
-            % (npoint, xyz.shape[1])
-        )
+        raise ValueError("caanot sample %i points from an input set of %i points" % (npoint, xyz.shape[1]))
     if xyz.is_cuda:
         return tpcuda.furthest_point_sampling(xyz, npoint)
     else:
@@ -102,13 +99,9 @@ def backward(ctx, grad_out):
         idx, weight, m = ctx.three_interpolate_for_backward
 
         if grad_out.is_cuda:
-            grad_features = tpcuda.three_interpolate_grad(
-                grad_out.contiguous(), idx, weight, m
-            )
+            grad_features = tpcuda.three_interpolate_grad(grad_out.contiguous(), idx, weight, m)
         else:
-            grad_features = tpcpu.knn_interpolate_grad(
-                grad_out.contiguous(), idx, weight, m
-            )
+            grad_features = tpcpu.knn_interpolate_grad(grad_out.contiguous(), idx, weight, m)
 
         return grad_features, None, None
 
@@ -150,23 +143,17 @@ def grouping_operation(features, idx):
     all_idx = idx.reshape(idx.shape[0], -1)
     all_idx = all_idx.unsqueeze(1).repeat(1, features.shape[1], 1)
     grouped_features = features.gather(2, all_idx)
-    return grouped_features.reshape(
-        idx.shape[0], features.shape[1], idx.shape[1], idx.shape[2]
-    )
+    return grouped_features.reshape(idx.shape[0], features.shape[1], idx.shape[1], idx.shape[2])
 
 
-def ball_query_dense(
-    radius, nsample, xyz, new_xyz, batch_xyz=None, batch_new_xyz=None, sort=False
-):
+def ball_query_dense(radius, nsample, xyz, new_xyz, batch_xyz=None, batch_new_xyz=None, sort=False):
     # type: (Any, float, int, torch.Tensor, torch.Tensor) -> torch.Tensor
     if new_xyz.is_cuda:
         if sort:
             raise NotImplementedError("CUDA version does not sort the neighbors")
         ind, dist = tpcuda.ball_query_dense(new_xyz, xyz, radius, nsample)
     else:
-        ind, dist = tpcpu.dense_ball_query(
-            new_xyz, xyz, radius, nsample, mode=0, sorted=sort
-        )
+        ind, dist = tpcpu.dense_ball_query(new_xyz, xyz, radius, nsample, mode=0, sorted=sort)
     return ind, dist
 
 
@@ -175,13 +162,9 @@ def ball_query_partial_dense(radius, nsample, x, y, batch_x, batch_y, sort=False
     if x.is_cuda:
         if sort:
             raise NotImplementedError("CUDA version does not sort the neighbors")
-        ind, dist = tpcuda.ball_query_partial_dense(
-            x, y, batch_x, batch_y, radius, nsample
-        )
+        ind, dist = tpcuda.ball_query_partial_dense(x, y, batch_x, batch_y, radius, nsample)
     else:
-        ind, dist = tpcpu.batch_ball_query(
-            x, y, batch_x, batch_y, radius, nsample, mode=0, sorted=sort
-        )
+        ind, dist = tpcpu.batch_ball_query(x, y, batch_x, batch_y, radius, nsample, mode=0, sorted=sort)
     return ind, dist
 
 
@@ -224,9 +207,7 @@ def ball_query(
         assert x.size(0) == batch_x.size(0)
         assert y.size(0) == batch_y.size(0)
         assert x.dim() == 2
-        return ball_query_partial_dense(
-            radius, nsample, x, y, batch_x, batch_y, sort=sort
-        )
+        return ball_query_partial_dense(radius, nsample, x, y, batch_x, batch_y, sort=sort)
 
     elif mode.lower() == "dense":
         if (batch_x is not None) or (batch_y is not None):
@@ -248,9 +229,7 @@ def forward(ctx, xyz1, xyz2):
     @staticmethod
     def backward(ctx, grad_dist1, grad_dist2):
         xyz1, xyz2, idx1, idx2 = ctx.saved_tensors
-        grad_xyz1, grad_xyz2 = tpcuda.chamfer_dist_grad(
-            xyz1, xyz2, idx1, idx2, grad_dist1, grad_dist2
-        )
+        grad_xyz1, grad_xyz2 = tpcuda.chamfer_dist_grad(xyz1, xyz2, idx1, idx2, grad_dist1, grad_dist2)
         return grad_xyz1, grad_xyz2
 
 
@@ -281,4 +260,3 @@ def chamfer_dist(xyz1, xyz2, ignore_zeros=False):
 
     dist1, dist2 = ChamferFunction.apply(xyz1, xyz2)
     return torch.mean(dist1) + torch.mean(dist2)
-