New File Types, automatic detection and header parameters

main.py

@@ -7,6 +7,7 @@ https://github.com/dragonbook/pointnet2-pytorch/blob/master/main.py
 
 import os
 import sys
+from distutils.util import strtobool
 import random
 import numpy as np
 import argparse
@@ -35,12 +36,15 @@ parser.add_argument('--outf', type=str, default='checkpoint', help='output folde
 parser.add_argument('--batch_size', type=int, default=8, help='input batch size')
 parser.add_argument('--test_per_batches', type=int, default=1000, help='run a test batch per training batches number')
 parser.add_argument('--num_workers', type=int, default=4, help='number of data loading workers')
+parser.add_argument('--headers', type=strtobool, default=True, help='if raw files come with headers')
+
 
 opt = parser.parse_args()
 print(opt)
 
 # Random seed
 opt.manual_seed = 123
+opt.headers = bool(opt.headers)
 print('Random seed: ', opt.manual_seed)
 random.seed(opt.manual_seed)
 np.random.seed(opt.manual_seed)
@@ -64,10 +68,10 @@ if __name__ == '__main__':
         train_transform = GT.Compose([GT.NormalizeScale(), RotTransform, TransTransform])
         test_transform = GT.Compose([GT.NormalizeScale(), ])
 
-    dataset = ShapeNetPartSegDataset(root_dir=opt.dataset, train=True, transform=train_transform, npoints=opt.npoints)
+    dataset = ShapeNetPartSegDataset(root_dir=opt.dataset, train=True, transform=train_transform, npoints=opt.npoints, headers=opt.headers)
     dataLoader = DataLoader(dataset, batch_size=opt.batch_size, shuffle=True, num_workers=opt.num_workers)
 
-    test_dataset = ShapeNetPartSegDataset(root_dir=opt.dataset, train=False, transform=test_transform, npoints=opt.npoints)
+    test_dataset = ShapeNetPartSegDataset(root_dir=opt.dataset, train=False, transform=test_transform, npoints=opt.npoints, headers=opt.headers)
     test_dataLoader = DataLoader(test_dataset, batch_size=opt.batch_size, shuffle=True, num_workers=opt.num_workers)
 
     num_classes = dataset.num_classes()