Labels can now be placed along next to the points within the datasetfile

predict/predict.py

@@ -6,15 +6,22 @@ from dataset.shapenet import ShapeNetPartSegDataset
 from model.pointnet2_part_seg import PointNet2PartSegmentNet
 import torch_geometric.transforms as GT
 import torch
+from distutils.util import strtobool
+
 import numpy as np
 import argparse
 
 ##
 parser = argparse.ArgumentParser()
-parser.add_argument('--dataset', type=str, default='data', help='dataset path')
 parser.add_argument('--npoints', type=int, default=2048, help='resample points number')
-parser.add_argument('--model', type=str, default='./checkpoint/seg_model_custom_246.pth', help='model path')
+parser.add_argument('--model', type=str, default='./checkpoint/seg_model_custom_3.pth', help='model path')
 parser.add_argument('--sample_idx', type=int, default=0, help='select a sample to segment and view result')
+parser.add_argument('--headers', type=strtobool, default=True, help='if raw files come with headers')
+parser.add_argument('--collate_per_segment', type=strtobool, default=True, help='whether to look at pointclouds or sub')
+parser.add_argument('--has_variations', type=strtobool, default=False,
+                    help='whether a single pointcloud has variations '
+                         'named int(id)_pc.(xyz|dat) look at pointclouds or sub')
+
 opt = parser.parse_args()
 print(opt)
 
@@ -26,11 +33,14 @@ if __name__ == '__main__':
 
     test_dataset = ShapeNetPartSegDataset(
         mode='predict',
-        root_dir=opt.dataset,
-        transform=None,
+        root_dir='data',
         npoints=opt.npoints,
-        refresh=True
+        refresh=True,
+        collate_per_segment=opt.collate_per_segment,
+        has_variations=opt.has_variations,
+        headers=opt.headers
     )
+
     num_classes = test_dataset.num_classes()
 
     print('test dataset size: ', len(test_dataset))