Added normals to prediction DataObject

dataset/shapenet.py

@@ -46,8 +46,9 @@ class CustomShapeNet(InMemoryDataset):
 
     def download(self):
         dir_count = len([name for name in os.listdir(self.raw_dir) if os.path.isdir(os.path.join(self.raw_dir, name))])
-        print(f'{dir_count} folders have been found....')
+
         if dir_count:
+            print(f'{dir_count} folders have been found....')
             return dir_count
         raise IOError("No raw pointclouds have been found.")
 
@@ -179,6 +180,7 @@ class ShapeNetPartSegDataset(Dataset):
     Resample raw point cloud to fixed number of points.
     Map raw label from range [1, N] to [0, N-1].
     """
+
     def __init__(self, root_dir, npoints=1024, mode='train', **kwargs):
         super(ShapeNetPartSegDataset, self).__init__()
         self.mode = mode
@@ -191,7 +193,8 @@ class ShapeNetPartSegDataset(Dataset):
 
         # Resample to fixed number of points
         try:
-            choice = np.random.choice(data.pos.shape[0], self.npoints, replace=True)
+            npoints = self.npoints if self.mode != 'predict' else data.pos.shape[0]
+            choice = np.random.choice(data.pos.shape[0], npoints, replace=False)
         except ValueError:
             choice = []
 
@@ -204,7 +207,7 @@ class ShapeNetPartSegDataset(Dataset):
             'labels': labels  # torch.Tensor (n,)
         }
         if self.mode == 'predict':
-            normals = data.normals[choice]
+            normals = data.normals[choice, :]
             sample.update(normals=normals)
 
         return sample