From 47a76dc978d26ba4db170386b43f8b0e0ba5b707 Mon Sep 17 00:00:00 2001 From: Si11ium Date: Wed, 31 Jul 2019 12:55:47 +0200 Subject: [PATCH] File based header detection, collate_per_PC training. --- .gitignore | 1 + .idea/pointnet2-pytorch.iml | 1 + dataset/shapenet.py | 197 +++++++++++++++++++++++++++++++++--- main.py | 7 +- vis/show_seg_res.py | 53 +++++----- vis/view.py | 3 +- 6 files changed, 218 insertions(+), 44 deletions(-) diff --git a/.gitignore b/.gitignore index 87b49df..9c25236 100644 --- a/.gitignore +++ b/.gitignore @@ -127,3 +127,4 @@ dmypy.json /data/ /checkpoint/ /shapenet/ +/vis/ diff --git a/.idea/pointnet2-pytorch.iml b/.idea/pointnet2-pytorch.iml index 7401a39..d1ad953 100644 --- a/.idea/pointnet2-pytorch.iml +++ b/.idea/pointnet2-pytorch.iml @@ -4,6 +4,7 @@ + diff --git a/dataset/shapenet.py b/dataset/shapenet.py index 4ceedb0..99715ea 100644 --- a/dataset/shapenet.py +++ b/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. - 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) + data = self._transform_and_filter(data) + if self.collate_per_element: + datasets[setting].append(data) + if not self.collate_per_element: + datasets[setting].append(Data(**{key: torch.cat(data[key]) for key in data.keys()})) 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] diff --git a/main.py b/main.py index cb44138..ba572bd 100644 --- a/main.py +++ b/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() diff --git a/vis/show_seg_res.py b/vis/show_seg_res.py index 1e20323..9b5a771 100644 --- a/vis/show_seg_res.py +++ b/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 ..') - pred_label, gt_label = eval_sample(net, sample) - print('Eval done ..') + print('Eval test sample ..') + pred_label, gt_label = eval_sample(net, sample) + print('Eval done ..') - # Get sample result - print('Compute mIoU ..') - points = sample['points'].numpy() - pred_labels = pred_label.numpy() - gt_labels = gt_label.numpy() - diff_labels = label_diff(pred_labels, gt_labels) + # Get sample result + print('Compute mIoU ..') + points = sample['points'].numpy() + pred_labels = pred_label.numpy() + gt_labels = gt_label.numpy() + 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 ..') - # view_points_labels(points, gt_labels) + # print('View gt labels ..') + # view_points_labels(points, gt_labels) - # print('View diff labels ..') - # view_points_labels(points, diff_labels) + print('View diff labels ..') + print(diff_labels) + view_points_labels(points, diff_labels) - print('View pred labels ..') - view_points_labels(points, pred_labels) + # print('View pred labels ..') + # print(pred_labels) + # view_points_labels(points, pred_labels) diff --git a/vis/view.py b/vis/view.py index 948fbf1..3a9ac14 100644 --- a/vis/view.py +++ b/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, )