Can now be trained with normals

model/pointnet2_part_seg.py

@@ -8,15 +8,15 @@ from torch_geometric.utils.num_nodes import maybe_num_nodes
 from torch_geometric.data.data import Data
 from torch_scatter import scatter_add, scatter_max
 
-GLOBAL_POINT_FEATURES = 6
 
 class PointNet2SAModule(torch.nn.Module):
-    def __init__(self, sample_radio, radius, max_num_neighbors, mlp):
+    def __init__(self, sample_radio, radius, max_num_neighbors, mlp, features=3):
         super(PointNet2SAModule, self).__init__()
         self.sample_ratio = sample_radio
         self.radius = radius
         self.max_num_neighbors = max_num_neighbors
         self.point_conv = PointConv(mlp)
+        self.features=features
 
     def forward(self, data):
         x, pos, batch = data
@@ -40,9 +40,10 @@ class PointNet2GlobalSAModule(torch.nn.Module):
     One group with all input points, can be viewed as a simple PointNet module.
     It also return the only one output point(set as origin point).
     '''
-    def __init__(self, mlp):
+    def __init__(self, mlp, features=3):
         super(PointNet2GlobalSAModule, self).__init__()
         self.mlp = mlp
+        self.features = features
 
     def forward(self, data):
         x, pos, batch = data
@@ -52,7 +53,7 @@ class PointNet2GlobalSAModule(torch.nn.Module):
         x1 = scatter_max(x1, batch, dim=0)[0]  # (batch_size, C1)
 
         batch_size = x1.shape[0]
-        pos1 = x1.new_zeros((batch_size, GLOBAL_POINT_FEATURES))  # set the output point as origin
+        pos1 = x1.new_zeros((batch_size, self.features))  # set the output point as origin
         batch1 = torch.arange(batch_size).to(batch.device, batch.dtype)
 
         return x1, pos1, batch1
@@ -158,44 +159,47 @@ class PointNet2PartSegmentNet(torch.nn.Module):
         - https://github.com/charlesq34/pointnet2/blob/master/models/pointnet2_part_seg.py
         - https://github.com/rusty1s/pytorch_geometric/blob/master/examples/pointnet++.py
     '''
-    def __init__(self, num_classes):
+    def __init__(self, num_classes, with_normals=False):
         super(PointNet2PartSegmentNet, self).__init__()
         self.num_classes = num_classes
+        self.features = 3 if not with_normals else 6
 
         # SA1
         sa1_sample_ratio = 0.5
         sa1_radius = 0.2
         sa1_max_num_neighbours = 64
-        sa1_mlp = make_mlp(GLOBAL_POINT_FEATURES, [64, 64, 128])
-        self.sa1_module = PointNet2SAModule(sa1_sample_ratio, sa1_radius, sa1_max_num_neighbours, sa1_mlp)
+        sa1_mlp = make_mlp(self.features, [64, 64, 128])
+        self.sa1_module = PointNet2SAModule(sa1_sample_ratio, sa1_radius, sa1_max_num_neighbours, sa1_mlp,
+                                            features=self.features)
 
         # SA2
         sa2_sample_ratio = 0.25
         sa2_radius = 0.4
         sa2_max_num_neighbours = 64
-        sa2_mlp = make_mlp(128+GLOBAL_POINT_FEATURES, [128, 128, 256])
-        self.sa2_module = PointNet2SAModule(sa2_sample_ratio, sa2_radius, sa2_max_num_neighbours, sa2_mlp)
+        sa2_mlp = make_mlp(128+self.features, [128, 128, 256])
+        self.sa2_module = PointNet2SAModule(sa2_sample_ratio, sa2_radius, sa2_max_num_neighbours, sa2_mlp,
+                                            features=self.features)
 
         # SA3
-        sa3_mlp = make_mlp(256+GLOBAL_POINT_FEATURES, [256, 512, 1024])
-        self.sa3_module = PointNet2GlobalSAModule(sa3_mlp)
+        sa3_mlp = make_mlp(256+self.features, [256, 512, 1024])
+        self.sa3_module = PointNet2GlobalSAModule(sa3_mlp, self.features)
 
         ##
-        knn_num = GLOBAL_POINT_FEATURES
+        knn_num = self.features
 
         # FP3, reverse of sa3
         fp3_knn_num = 1  # After global sa module, there is only one point in point cloud
-        fp3_mlp = make_mlp(1024+256+GLOBAL_POINT_FEATURES, [256, 256])
+        fp3_mlp = make_mlp(1024+256+self.features, [256, 256])
         self.fp3_module = PointNet2FPModule(fp3_knn_num, fp3_mlp)
 
         # FP2, reverse of sa2
         fp2_knn_num = knn_num
-        fp2_mlp = make_mlp(256+128+GLOBAL_POINT_FEATURES, [256, 128])
+        fp2_mlp = make_mlp(256+128+self.features, [256, 128])
         self.fp2_module = PointNet2FPModule(fp2_knn_num, fp2_mlp)
 
         # FP1, reverse of sa1
         fp1_knn_num = knn_num
-        fp1_mlp = make_mlp(128+GLOBAL_POINT_FEATURES, [128, 128, 128])
+        fp1_mlp = make_mlp(128+self.features, [128, 128, 128])
         self.fp1_module = PointNet2FPModule(fp1_knn_num, fp1_mlp)
 
         self.fc1 = Lin(128, 128)
@@ -252,11 +256,12 @@ class PointNet2PartSegmentNet(torch.nn.Module):
 
 if __name__ == '__main__':
     num_classes = 10
-    net = PointNet2PartSegmentNet(num_classes)
+    num_features = 6
+    net = PointNet2PartSegmentNet(num_classes, features=num_features)
 
     #
     print('Test dense input ..')
-    data1 = torch.rand((2, GLOBAL_POINT_FEATURES, 1024))  # (batch_size, 3, num_points)
+    data1 = torch.rand((2, num_features, 1024))  # (batch_size, 3, num_points)
     print('data1: ', data1.shape)
 
     out1 = net(data1)
@@ -272,7 +277,7 @@ if __name__ == '__main__':
         data_batch = Data()
 
         # data_batch.x = None
-        data_batch.pos = torch.cat([torch.rand(pos_num1, GLOBAL_POINT_FEATURES), torch.rand(pos_num2, GLOBAL_POINT_FEATURES)], dim=0)
+        data_batch.pos = torch.cat([torch.rand(pos_num1, num_features), torch.rand(pos_num2, num_features)], dim=0)
         data_batch.batch = torch.cat([torch.zeros(pos_num1, dtype=torch.long), torch.ones(pos_num2, dtype=torch.long)])
 
         return data_batch