fix #20

Author: chrischoy

config.py

@@ -103,8 +103,8 @@ def add_argument_group(name):
 data_arg.add_argument('--test_batch_size', type=int, default=1)
 data_arg.add_argument('--cache_data', type=str2bool, default=False)
 data_arg.add_argument(
-    '--threads', type=int, default=1, help='num threads for train/test dataloader')
-data_arg.add_argument('--val_threads', type=int, default=1, help='num threads for val dataloader')
+    '--num_workers', type=int, default=1, help='num workers for train/test dataloader')
+data_arg.add_argument('--num_val_workers', type=int, default=1, help='num workers for val dataloader')
 data_arg.add_argument('--ignore_label', type=int, default=255)
 data_arg.add_argument('--return_transformation', type=str2bool, default=False)
 data_arg.add_argument('--ignore_duplicate_class', type=str2bool, default=False)

lib/dataset.py

@@ -11,7 +11,7 @@
 
 import MinkowskiEngine as ME
 
-from lib.pc_utils import read_plyfile
+from plyfile import PlyData
 import lib.transforms as t
 from lib.dataloader import InfSampler
 from lib.voxelizer import Voxelizer
@@ -173,7 +173,12 @@ def __getitem__(self, index):
 
   def load_ply(self, index):
     filepath = self.data_root / self.data_paths[index]
-    return read_plyfile(filepath), None
+    plydata = PlyData.read(filepath)
+    data = plydata.elements[0].data
+    coords = np.array([data['x'], data['y'], data['z']], dtype=np.float32).T
+    feats = np.array([data['red'], data['green'], data['blue']], dtype=np.float32).T
+    labels = np.array(data['label'], dtype=np.int32)
+    return coords, feats, labels, None
 
   def __len__(self):
     num_data = len(self.data_paths)
@@ -412,7 +417,7 @@ def __len__(self):
 def initialize_data_loader(DatasetClass,
                            config,
                            phase,
-                           threads,
+                           num_workers,
                            shuffle,
                            repeat,
                            augment_data,
@@ -467,7 +472,7 @@ def initialize_data_loader(DatasetClass,
 
   data_args = {
       'dataset': dataset,
-      'num_workers': threads,
+      'num_workers': num_workers,
       'batch_size': batch_size,
       'collate_fn': collate_fn,
   }

lib/datasets/stanford.py

@@ -4,11 +4,8 @@
 import numpy as np
 from collections import defaultdict
 from scipy import spatial
-from plyfile import PlyData, PlyElement
+from plyfile import PlyData
 
-import torch
-
-from lib.pc_utils import read_plyfile
 from lib.utils import read_txt, fast_hist, per_class_iu
 from lib.dataset import VoxelizationDataset, DatasetPhase, str2datasetphase_type, cache
 import lib.transforms as t

lib/train.py

@@ -72,10 +72,7 @@ def train(model, data_loader, val_data_loader, config, transform_data_fn=None):
       for sub_iter in range(config.iter_size):
         # Get training data
         data_timer.tic()
-        if config.return_transformation:
-          coords, input, target, pointcloud, transformation = data_iter.next()
-        else:
-          coords, input, target = data_iter.next()
+        coords, input, target = data_iter.next()
 
         # For some networks, making the network invariant to even, odd coords is important
         coords[:, :3] += (torch.rand(3) * 100).type_as(coords)

lib/transforms.py

@@ -287,19 +287,16 @@ def __init__(self, limit_numpoints):
     self.limit_numpoints = limit_numpoints
 
   def __call__(self, list_data):
-    coords, feats, labels, pointclouds, transformations = list(zip(*list_data))
+    coords, feats, labels, transformations = list(zip(*list_data))
     cfl_collate_fn = cfl_collate_fn_factory(limit_numpoints=self.limit_numpoints)
     coords_batch, feats_batch, labels_batch = cfl_collate_fn(list(zip(coords, feats, labels)))
     num_truncated_batch = coords_batch[:, -1].max().item() + 1
 
     batch_id = 0
     pointclouds_batch, transformations_batch = [], []
-    for pointcloud, transformation in zip(pointclouds, transformations):
+    for transformation in transformations:
       if batch_id >= num_truncated_batch:
         break
-      pointclouds_batch.append(
-          torch.cat((torch.from_numpy(pointcloud), torch.ones(pointcloud.shape[0], 1) * batch_id),
-                    1))
       transformations_batch.append(
           torch.cat(
               (torch.from_numpy(transformation), torch.ones(transformation.shape[0], 1) * batch_id),
@@ -308,4 +305,4 @@ def __call__(self, list_data):
 
     pointclouds_batch = torch.cat(pointclouds_batch, 0).float()
     transformations_batch = torch.cat(transformations_batch, 0).float()
-    return coords_batch, feats_batch, labels_batch, pointclouds_batch, transformations_batch
+    return coords_batch, feats_batch, labels_batch, transformations_batch

main.py

@@ -70,7 +70,7 @@ def main():
         DatasetClass,
         config,
         phase=config.train_phase,
-        threads=config.threads,
+        num_workers=config.num_workers,
         augment_data=True,
         shuffle=True,
         repeat=True,
@@ -80,7 +80,7 @@ def main():
     val_data_loader = initialize_data_loader(
         DatasetClass,
         config,
-        threads=config.val_threads,
+        num_workers=config.num_val_workers,
         phase=config.val_phase,
         augment_data=False,
         shuffle=True,
@@ -97,7 +97,7 @@ def main():
     test_data_loader = initialize_data_loader(
         DatasetClass,
         config,
-        threads=config.threads,
+        num_workers=config.num_workers,
         phase=config.test_phase,
         augment_data=False,
         shuffle=False,