File based header detection, collate_per_PC training.

dataset/shapenet.py

@@ -20,6 +20,7 @@ class CustomShapeNet(InMemoryDataset):
                  headers=True, **kwargs):
         self.has_headers = headers
         self.collate_per_element = collate_per_segment
+        self.train = train
         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)
@@ -70,12 +71,9 @@ class CustomShapeNet(InMemoryDataset):
         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):
-            path_to_clouds = os.path.join(self.raw_dir, setting)
+        idx, data_folder = (0, self.raw_file_names[0]) if self.train else (1, self.raw_file_names[1])
+        path_to_clouds = os.path.join(self.raw_dir, data_folder)
 
         if '.headers' in os.listdir(path_to_clouds):
             self.has_headers = True
@@ -85,12 +83,12 @@ class CustomShapeNet(InMemoryDataset):
             pass
 
         for pointcloud in tqdm(os.scandir(path_to_clouds)):
-
             if not os.path.isdir(pointcloud):
                 continue
             data, paths = None, 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
@@ -131,12 +129,13 @@ class CustomShapeNet(InMemoryDataset):
                     # , points=points, norm=points[:3], )
                     data = self._transform_and_filter(data)
                     if self.collate_per_element:
-                            datasets[setting].append(data)
+                        datasets[data_folder].append(data)
             if not self.collate_per_element:
-                    datasets[setting].append(Data(**{key: torch.cat(data[key]) for key in data.keys()}))
+                datasets[data_folder].append(Data(**{key: torch.cat(data[key]) for key in data.keys()}))
 
+        if datasets[data_folder]:
             os.makedirs(self.processed_dir, exist_ok=True)
-            torch.save(self.collate(datasets[setting]), self.processed_paths[idx])
+            torch.save(self.collate(datasets[data_folder]), self.processed_paths[idx])
 
     def __repr__(self):
         return f'{self.__class__.__name__}({len(self)})'
@@ -291,7 +290,7 @@ class PredictNetPartSegDataset(Dataset):
     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)
+        self.dataset = ShapeNetPartSegDataset(root=root_dir, train=False, transform=transform, headers=headers)
 
     def __getitem__(self, index):
         data = self.dataset[index]

vis/show_seg_res.py

@@ -5,7 +5,7 @@ import sys
 import os
 sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/../')  # add project root directory
 
-from dataset.shapenet import PredictNetPartSegDataset
+from dataset.shapenet import ShapeNetPartSegDataset
 from model.pointnet2_part_seg import PointNet2PartSegmentNet
 import torch_geometric.transforms as GT
 import torch
@@ -16,8 +16,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_custom_8.pth', help='model path')
+parser.add_argument('--npoints', type=int, default=2048, help='resample points number')
+parser.add_argument('--model', type=str, default='./checkpoint/seg_model_custom_249.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)
@@ -28,8 +28,10 @@ if __name__ == '__main__':
     print('Construct dataset ..')
     test_transform = GT.Compose([GT.NormalizeScale(),])
 
-    test_dataset = PredictNetPartSegDataset(
+    test_dataset = ShapeNetPartSegDataset(
         root_dir=opt.dataset,
+        collate_per_segment=False,
+        train=False,
         transform=test_transform,
         npoints=opt.npoints
     )
@@ -121,16 +123,17 @@ if __name__ == '__main__':
 
         print('mIoU: ', compute_mIoU(pred_labels, gt_labels))
 
-
         # View result
+        if True:
+            print('View gt labels ..')
+            view_points_labels(points, gt_labels)
 
-        # print('View gt labels ..')
-        # view_points_labels(points, gt_labels)
-
+        if True:
             print('View diff labels ..')
             print(diff_labels)
             view_points_labels(points, diff_labels)
 
-        # print('View pred labels ..')
-        # print(pred_labels)
-        # view_points_labels(points, pred_labels)
+        if True:
+            print('View pred labels ..')
+            print(pred_labels)
+            view_points_labels(points, pred_labels)