File based header detection, collate_per_PC training.

.gitignore

@@ -127,3 +127,4 @@ dmypy.json
 /data/
 /checkpoint/
 /shapenet/
+/vis/

.idea/pointnet2-pytorch.iml

@@ -4,6 +4,7 @@
     <content url="file://$MODULE_DIR$">
       <excludeFolder url="file://$MODULE_DIR$/data" />
       <excludeFolder url="file://$MODULE_DIR$/net" />
+      <excludeFolder url="file://$MODULE_DIR$/vis" />
     </content>
     <orderEntry type="jdk" jdkName="Python 3.7 (torch)" jdkType="Python SDK" />
     <orderEntry type="sourceFolder" forTests="false" />

dataset/shapenet.py

@@ -16,11 +16,12 @@ 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,
+    def __init__(self, root, collate_per_segment=True, train=True, transform=None, pre_filter=None, pre_transform=None,
                  headers=True, **kwargs):
         self.has_headers = headers
+        self.collate_per_element = collate_per_segment
         super(CustomShapeNet, self).__init__(root, transform, pre_transform, pre_filter)
-        path = self.processed_paths[0] if train else self.processed_paths[1]
+        path = self.processed_paths[0] if train else self.processed_paths[-1]
         self.data, self.slices = torch.load(path)
         print("Initialized")
 
@@ -57,16 +58,37 @@ class CustomShapeNet(InMemoryDataset):
                 continue
         return data, slices
 
+    def _transform_and_filter(self, data):
+        # ToDo: ANy filter to apply? Then do it here.
+        if self.pre_filter is not None and not self.pre_filter(data):
+            data = self.pre_filter(data)
+            raise NotImplementedError
+        # ToDo: ANy transformation to apply? Then do it here.
+        if self.pre_transform is not None:
+            data = self.pre_transform(data)
+            raise NotImplementedError
+        return data
+
     def process(self, delimiter=' '):
         # idx = self.categories[self.category]
         # paths = [osp.join(path, idx) for path in self.raw_paths]
 
         datasets = defaultdict(list)
         for idx, setting in enumerate(self.raw_file_names):
-            for pointcloud in tqdm(os.scandir(os.path.join(self.raw_dir, setting))):
+            path_to_clouds = os.path.join(self.raw_dir, setting)
+
+            if '.headers' in os.listdir(path_to_clouds):
+                self.has_headers = True
+            elif 'no.headers' in os.listdir(path_to_clouds):
+                self.has_headers = False
+            else:
+                pass
+
+            for pointcloud in tqdm(os.scandir(path_to_clouds)):
+
                 if not os.path.isdir(pointcloud):
                     continue
-                paths = list()
+                data, paths = None, list()
                 for ext in ['dat', 'xyz']:
                     paths.extend(glob.glob(os.path.join(pointcloud.path, f'*.{ext}')))
                 for element in paths:
@@ -99,17 +121,19 @@ class CustomShapeNet(InMemoryDataset):
                         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)
-                        data = Data(y=y, pos=points[:, :3])
+                        if self.collate_per_element:
+                            data = Data(y=y, pos=points[:, :3])
+                        else:
+                            if not data:
+                                data = defaultdict(list)
+                            for key, val in dict(y=y, pos= points[:, :3]).items():
+                                data[key].append(val)
                         # , points=points, norm=points[:3], )
-                        # ToDo: ANy filter to apply? Then do it here.
+                        data = self._transform_and_filter(data)
-                        if self.pre_filter is not None and not self.pre_filter(data):
+                        if self.collate_per_element:
-                            data = self.pre_filter(data)
+                            datasets[setting].append(data)
-                            raise NotImplementedError
+                if not self.collate_per_element:
-                        # ToDo: ANy transformation to apply? Then do it here.
+                    datasets[setting].append(Data(**{key: torch.cat(data[key]) for key in data.keys()}))
-                        if self.pre_transform is not None:
-                            data = self.pre_transform(data)
-                            raise NotImplementedError
-                        datasets[setting].append(data)
 
             os.makedirs(self.processed_dir, exist_ok=True)
             torch.save(self.collate(datasets[setting]), self.processed_paths[idx])
@@ -123,10 +147,151 @@ 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, headers=True):
+    def __init__(self, root_dir, collate_per_segment=True, 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, headers=headers)
+        self.dataset = CustomShapeNet(root=root_dir, collate_per_segment=collate_per_segment,
+                                      train=train, transform=transform, headers=headers)
+
+    def __getitem__(self, index):
+        data = self.dataset[index]
+        points, labels = data.pos, data.y
+
+        # Resample to fixed number of points
+        try:
+            choice = np.random.choice(points.shape[0], self.npoints, replace=True)
+        except ValueError:
+            choice = []
+
+        points, labels = points[choice, :], labels[choice]
+
+        labels -= 1 if self.num_classes() in labels else 0   # Map label from [1, C] to [0, C-1]
+
+        sample = {
+            'points': points,  # torch.Tensor (n, 3)
+            'labels': labels  # torch.Tensor (n,)
+        }
+
+        return sample
+
+    def __len__(self):
+        return len(self.dataset)
+
+    def num_classes(self):
+        return self.dataset.num_classes
+
+
+class PredictionShapeNet(InMemoryDataset):
+    categories = {key: val for val, key in enumerate(['Box', 'Cone', 'Cylinder', 'Sphere'])}
+
+    def __init__(self, root, transform=None, pre_filter=None, pre_transform=None,
+                 headers=True, **kwargs):
+        self.has_headers = headers
+        super(PredictionShapeNet, self).__init__(root, transform, pre_transform, pre_filter)
+        path = self.processed_paths[0]
+        self.data, self.slices = torch.load(path)
+        print("Initialized")
+
+    @property
+    def raw_file_names(self):
+        # Maybe add more data like validation sets
+        return ['predict']
+
+    @property
+    def processed_file_names(self):
+        return [f'{x}.pt' for x in self.raw_file_names]
+
+    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:
+            return dir_count
+        raise IOError("No raw pointclouds have been found.")
+
+    @property
+    def num_classes(self):
+        return len(self.categories)
+
+    def _load_dataset(self):
+        data, slices = None, None
+        while True:
+            try:
+                filepath = os.path.join(self.root, self.processed_dir, f'{"train" if self.train else "test"}.pt')
+                data, slices = torch.load(filepath)
+                print('Dataset Loaded')
+                break
+            except FileNotFoundError:
+                self.process()
+                continue
+        return data, slices
+
+    def process(self, delimiter=' '):
+
+        datasets = defaultdict(list)
+        for idx, setting in enumerate(self.raw_file_names):
+            path_to_clouds = os.path.join(self.raw_dir, setting)
+
+            if '.headers' in os.listdir(path_to_clouds):
+                self.has_headers = True
+            elif 'no.headers' in os.listdir(path_to_clouds):
+                self.has_headers = False
+            else:
+                pass
+
+            for pointcloud in tqdm(os.scandir(path_to_clouds)):
+                if not os.path.isdir(pointcloud):
+                    continue
+                for extention in ['dat', 'xyz']:
+                    file = os.path.join(pointcloud.path, f'pc.{extention}')
+                    if not os.path.exists(file):
+                        continue
+                    with open(file, 'r') as f:
+                            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
+                                    for x in line.rstrip().split(delimiter)[None:None]] for line in f if line != '']
+                    points = torch.tensor(src, dtype=None).squeeze()
+                    if not len(points.shape) > 1:
+                        continue
+                    # pos = points[:, :3]
+                    # norm = points[:, 3:]
+                    y_fake_all = [-1] * points.shape[0]
+                    y = torch.as_tensor(y_fake_all, dtype=torch.int)
+                    # points = torch.as_tensor(points, dtype=torch.float)
+                    # norm = torch.as_tensor(norm, dtype=torch.float)
+                    data = Data(y=y, pos=points[:, :3])
+                    # , points=points, norm=points[:3], )
+                    # ToDo: ANy filter to apply? Then do it here.
+                    if self.pre_filter is not None and not self.pre_filter(data):
+                        data = self.pre_filter(data)
+                        raise NotImplementedError
+                    # ToDo: ANy transformation to apply? Then do it here.
+                    if self.pre_transform is not None:
+                        data = self.pre_transform(data)
+                        raise NotImplementedError
+                    datasets[setting].append(data)
+
+            os.makedirs(self.processed_dir, exist_ok=True)
+            torch.save(self.collate(datasets[setting]), self.processed_paths[idx])
+
+    def __repr__(self):
+        return f'{self.__class__.__name__}({len(self)})'
+
+
+class PredictNetPartSegDataset(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, transform=None, npoints=2048, headers=True):
+        super(PredictNetPartSegDataset, self).__init__()
+        self.npoints = npoints
+        self.dataset = PredictionShapeNet(root=root_dir, train=False, transform=transform, headers=headers)
 
     def __getitem__(self, index):
         data = self.dataset[index]

main.py

@@ -37,6 +37,7 @@ 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')
+parser.add_argument('--collate_per_segment', type=strtobool, default=True, help='whether to look at pointclouds or sub')
 
 
 opt = parser.parse_args()
@@ -68,10 +69,12 @@ 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, headers=opt.headers)
+    dataset = ShapeNetPartSegDataset(root_dir=opt.dataset, collate_per_segment=opt.collate_per_segment,
+                                     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, headers=opt.headers)
+    test_dataset = ShapeNetPartSegDataset(root_dir=opt.dataset, collate_per_segment=opt.collate_per_segment,
+                                          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()

vis/show_seg_res.py

@@ -1,11 +1,11 @@
 # Warning: import open3d may lead crash, try to to import open3d first here
-from view import view_points_labels
+from vis.view import view_points_labels
 
 import sys
 import os
 sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/../')  # add project root directory
 
-from dataset.shapenet import ShapeNetPartSegDataset
+from dataset.shapenet import PredictNetPartSegDataset
 from model.pointnet2_part_seg import PointNet2PartSegmentNet
 import torch_geometric.transforms as GT
 import torch
@@ -17,9 +17,8 @@ import argparse
 parser = argparse.ArgumentParser()
 parser.add_argument('--dataset', type=str, default='data', help='dataset path')
 parser.add_argument('--npoints', type=int, default=50, help='resample points number')
-parser.add_argument('--model', type=str, default='./checkpoint/seg_model_Airplane_24.pth', help='model path')
+parser.add_argument('--model', type=str, default='./checkpoint/seg_model_custom_8.pth', help='model path')
 parser.add_argument('--sample_idx', type=int, default=0, help='select a sample to segment and view result')
-
 opt = parser.parse_args()
 print(opt)
 
@@ -29,9 +28,8 @@ if __name__ == '__main__':
     print('Construct dataset ..')
     test_transform = GT.Compose([GT.NormalizeScale(),])
 
-    test_dataset = ShapeNetPartSegDataset(
+    test_dataset = PredictNetPartSegDataset(
         root_dir=opt.dataset,
-        train=False,
         transform=test_transform,
         npoints=opt.npoints
     )
@@ -49,7 +47,7 @@ if __name__ == '__main__':
     # net = PointNetPartSegmentNet(num_classes)
     net = PointNet2PartSegmentNet(num_classes)
 
-    net.load_state_dict(torch.load(opt.model))
+    net.load_state_dict(torch.load(opt.model, map_location=device.type))
     net = net.to(device, dtype)
     net.eval()
 
@@ -104,30 +102,35 @@ if __name__ == '__main__':
 
 
     # Get one sample and eval
-    sample = test_dataset[opt.sample_idx]
+    #sample = test_dataset[opt.sample_idx]
+    r_idx = np.random.randint(0, len(test_dataset), 20)
+    for idx in r_idx:
+        sample = test_dataset[int(idx)]
 
-    print('Eval test sample ..')
+        print('Eval test sample ..')
-    pred_label, gt_label = eval_sample(net, sample)
+        pred_label, gt_label = eval_sample(net, sample)
-    print('Eval done ..')
+        print('Eval done ..')
 
 
-    # Get sample result
+        # Get sample result
-    print('Compute mIoU ..')
+        print('Compute mIoU ..')
-    points = sample['points'].numpy()
+        points = sample['points'].numpy()
-    pred_labels = pred_label.numpy()
+        pred_labels = pred_label.numpy()
-    gt_labels = gt_label.numpy()
+        gt_labels = gt_label.numpy()
-    diff_labels = label_diff(pred_labels, gt_labels)
+        diff_labels = label_diff(pred_labels, gt_labels)
 
-    print('mIoU: ', compute_mIoU(pred_labels, gt_labels))
+        print('mIoU: ', compute_mIoU(pred_labels, gt_labels))
 
 
-    # View result
+        # View result
 
-    # print('View gt labels ..')
+        # print('View gt labels ..')
-    # view_points_labels(points, gt_labels)
+        # view_points_labels(points, gt_labels)
 
-    # print('View diff labels ..')
+        print('View diff labels ..')
-    # view_points_labels(points, diff_labels)
+        print(diff_labels)
+        view_points_labels(points, diff_labels)
 
-    print('View pred labels ..')
+        # print('View pred labels ..')
-    view_points_labels(points, pred_labels)
+        # print(pred_labels)
+        # view_points_labels(points, pred_labels)

vis/view.py

@@ -28,7 +28,7 @@ def view_points(points, colors=None):
     '''
     cloud = o3d.PointCloud()
     cloud.points = o3d.Vector3dVector(points) 
-
+    # frame = o3d.create_mesh_coordinate_frame(-1, -1, -1)
     if colors is not None:
         if isinstance(colors, np.ndarray):
             cloud.colors = o3d.Vector3dVector(colors)
@@ -37,6 +37,7 @@ def view_points(points, colors=None):
     o3d.draw_geometries([cloud])
 
 
+
 def label2color(labels):
     '''
     labels: np.ndarray with shape (n, )