diff --git a/dataset/shapenet.py b/dataset/shapenet.py
index 33f4dbf..4ceedb0 100644
--- a/dataset/shapenet.py
+++ b/dataset/shapenet.py
@@ -16,7 +16,9 @@ class CustomShapeNet(InMemoryDataset):
 
     categories = {key: val for val, key in enumerate(['Box', 'Cone', 'Cylinder', 'Sphere'])}
 
-    def __init__(self, root, train=True, transform=None, pre_filter=None, pre_transform=None, **kwargs):
+    def __init__(self, root, train=True, transform=None, pre_filter=None, pre_transform=None,
+                 headers=True, **kwargs):
+        self.has_headers = headers
         super(CustomShapeNet, self).__init__(root, transform, pre_transform, pre_filter)
         path = self.processed_paths[0] if train else self.processed_paths[1]
         self.data, self.slices = torch.load(path)
@@ -64,20 +66,26 @@ class CustomShapeNet(InMemoryDataset):
             for pointcloud in tqdm(os.scandir(os.path.join(self.raw_dir, setting))):
                 if not os.path.isdir(pointcloud):
                     continue
-                for element in glob.glob(os.path.join(pointcloud.path, '*.dat')):
-                    if os.path.split(element)[-1] not in ['pc.dat']:
+                paths = list()
+                for ext in ['dat', 'xyz']:
+                    paths.extend(glob.glob(os.path.join(pointcloud.path, f'*.{ext}')))
+                for element in paths:
+                    if all([x not in os.path.split(element)[-1] for x in ['pc.dat', 'pc.xyz']]):
                         # Assign training data to the data container
                         # Following the original logic;
                         # y should be the label;
                         # pos should be the six dimensional vector describing: !its pos not points!!
                         # x,y,z,x_rot,y_rot,z_rot
-                        y_raw = os.path.splitext(element)[0].split('_')[-2]
+
+                        # Get the y - Label
+                        y_raw = next(i for i, v in enumerate(self.categories.keys()) if v.lower() in element.lower())
+                        # y_raw = os.path.splitext(element)[0].split('_')[-2]
                         with open(element,'r') as f:
-                            headers = f.__next__()
-                            # Check if there are no useable nodes in this file, header says 0.
-                            if not int(headers.rstrip().split(delimiter)[0]):
-                                continue
-                            # Get the y - Label
+                            if self.has_headers:
+                                headers = f.__next__()
+                                # Check if there are no useable nodes in this file, header says 0.
+                                if not int(headers.rstrip().split(delimiter)[0]):
+                                    continue
 
                             # Iterate over all rows
                             src = [[float(x) if x not in ['-nan(ind)', 'nan(ind)'] else 0
@@ -87,7 +95,7 @@ class CustomShapeNet(InMemoryDataset):
                             continue
                         # pos = points[:, :3]
                         # norm = points[:, 3:]
-                        y_all = [self.categories[y_raw]] * points.shape[0]
+                        y_all = [y_raw] * points.shape[0]
                         y = torch.as_tensor(y_all, dtype=torch.int)
                         # points = torch.as_tensor(points, dtype=torch.float)
                         # norm = torch.as_tensor(norm, dtype=torch.float)
@@ -115,10 +123,10 @@ 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, train=True, transform=None, npoints=1024):
+    def __init__(self, root_dir, train=True, transform=None, npoints=1024, headers=True):
         super(ShapeNetPartSegDataset, self).__init__()
         self.npoints = npoints
-        self.dataset = CustomShapeNet(root=root_dir, train=train, transform=transform)
+        self.dataset = CustomShapeNet(root=root_dir, train=train, transform=transform, headers=headers)
 
     def __getitem__(self, index):
         data = self.dataset[index]
diff --git a/main.py b/main.py
index 4c58913..cb44138 100644
--- a/main.py
+++ b/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()