From 97e36df1bab11e277cd1f9a8f0574889666770c1 Mon Sep 17 00:00:00 2001 From: Si11ium Date: Tue, 6 Aug 2019 08:36:59 +0200 Subject: [PATCH] Labels can now be placed along next to the points within the datasetfile --- dataset/shapenet.py | 45 +++++++++++++++++++-------------------------- main.py | 39 ++++++++++++++++++++++----------------- 2 files changed, 41 insertions(+), 43 deletions(-) diff --git a/dataset/shapenet.py b/dataset/shapenet.py index 4787301..ffe829e 100644 --- a/dataset/shapenet.py +++ b/dataset/shapenet.py @@ -23,16 +23,15 @@ class CustomShapeNet(InMemoryDataset): categories = {key: val for val, key in enumerate(['Box', 'Cone', 'Cylinder', 'Sphere'])} modes = {key: val for val, key in enumerate(['train', 'test', 'predict'])} - def __init__(self, root, collate_per_segment=True, mode='train', transform=None, pre_filter=None, pre_transform=None, - headers=True, has_variations=False, refresh=False): + def __init__(self, root_dir, collate_per_segment=True, mode='train', transform=None, pre_filter=None, + pre_transform=None, headers=True, has_variations=False, refresh=False, labels_within=False): assert mode in self.modes.keys(), f'"mode" must be one of {self.modes.keys()}' assert not (collate_per_segment and has_variations), 'Either use each element or pointclouds - with variations' - self.has_headers = headers - self.has_variations = has_variations - self.collate_per_element = collate_per_segment - self.mode = mode - self.refresh = refresh - super(CustomShapeNet, self).__init__(root, transform, pre_transform, pre_filter) + + #Set the Dataset Parameters + self.has_headers, self.has_variations, self.labels_within = headers, has_variations, labels_within + self.collate_per_element, self.mode, self.refresh = collate_per_segment, mode, refresh + super(CustomShapeNet, self).__init__(root_dir, transform, pre_transform, pre_filter) self.data, self.slices = self._load_dataset() print("Initialized") @@ -116,19 +115,6 @@ class CustomShapeNet(InMemoryDataset): if pattern.match(os.path.split(element)[-1]): continue else: - # The following two lines were intendations - # check if it is "just" a pc.dat or pc.xyz - # 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 - - # 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: if self.has_headers: headers = f.__next__() @@ -142,7 +128,15 @@ class CustomShapeNet(InMemoryDataset): points = torch.tensor(src, dtype=None).squeeze() if not len(points.shape) > 1: continue - y_all = ([y_raw] if self.mode != 'predict' else [-1]) * points.shape[0] + # Place Fake Labels to hold the given structure + if self.labels_within: + y_all = points[:, -1] + points = points[:, :-1] + else: + # Get the y - Label + y_raw = next(i for i, v in enumerate(self.categories.keys()) if v.lower() in element.lower()) + y_all = ([y_raw] if self.mode != 'predict' else [-1]) * points.shape[0] + y = torch.as_tensor(y_all, dtype=torch.int) if self.collate_per_element: data = Data(y=y, pos=points[:, :3]) # , points=points, norm=points[:, 3:]) @@ -178,12 +172,11 @@ 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, collate_per_segment=True, mode='train', transform=None, refresh=False, - has_variations=False, npoints=1024, headers=True): + def __init__(self, root_dir, npoints=1024, **kwargs): super(ShapeNetPartSegDataset, self).__init__() + kwargs.update(dict(root_dir=root_dir)) self.npoints = npoints - self.dataset = CustomShapeNet(root=root_dir, collate_per_segment=collate_per_segment, refresh=refresh, - mode=mode, transform=transform, headers=headers, has_variations=has_variations) + self.dataset = CustomShapeNet(**kwargs) def __getitem__(self, index): data = self.dataset[index] diff --git a/main.py b/main.py index 474a78b..7c05350 100644 --- a/main.py +++ b/main.py @@ -33,6 +33,7 @@ parser.add_argument('--npoints', type=int, default=1024, help='resample points n parser.add_argument('--model', type=str, default='', help='model path') parser.add_argument('--nepoch', type=int, default=250, help='number of epochs to train for') parser.add_argument('--outf', type=str, default='checkpoint', help='output folder') +parser.add_argument('--labels_within', type=strtobool, default=False, help='defines the label location') 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') @@ -56,30 +57,34 @@ torch.manual_seed(opt.manual_seed) torch.cuda.manual_seed(opt.manual_seed) if __name__ == '__main__': - # Dataset and transform print('Construct dataset ..') - if True: - rot_max_angle = 15 - trans_max_distance = 0.01 - RotTransform = GT.Compose([GT.RandomRotate(rot_max_angle, 0), - GT.RandomRotate(rot_max_angle, 1), - GT.RandomRotate(rot_max_angle, 2)] - ) - TransTransform = GT.RandomTranslate(trans_max_distance) + rot_max_angle = 15 + trans_max_distance = 0.01 - train_transform = GT.Compose([GT.NormalizeScale(), RotTransform, TransTransform]) - test_transform = GT.Compose([GT.NormalizeScale(), ]) + RotTransform = GT.Compose([GT.RandomRotate(rot_max_angle, 0), + GT.RandomRotate(rot_max_angle, 1), + GT.RandomRotate(rot_max_angle, 2)] + ) - dataset = ShapeNetPartSegDataset(root_dir=opt.dataset, collate_per_segment=opt.collate_per_segment, - mode='train', transform=train_transform, npoints=opt.npoints, - has_variations=opt.has_variations, headers=opt.headers) + TransTransform = GT.RandomTranslate(trans_max_distance) + train_transform = GT.Compose([GT.NormalizeScale(), RotTransform, TransTransform]) + test_transform = GT.Compose([GT.NormalizeScale(), ]) + + params = dict(root_dir=opt.dataset, + collate_per_segment=opt.collate_per_segment, + transform=train_transform, + npoints=opt.npoints, + labels_within=opt.labels_within, + has_variations=opt.has_variations, + headers=opt.headers + ) + + dataset = ShapeNetPartSegDataset(mode='train', **params) dataLoader = DataLoader(dataset, batch_size=opt.batch_size, shuffle=True, num_workers=opt.num_workers) - test_dataset = ShapeNetPartSegDataset(root_dir=opt.dataset, collate_per_segment=opt.collate_per_segment, - mode='test', transform=test_transform, npoints=opt.npoints, - has_variations=opt.has_variations, headers=opt.headers) + test_dataset = ShapeNetPartSegDataset(mode='test', **params) test_dataLoader = DataLoader(test_dataset, batch_size=opt.batch_size, shuffle=True, num_workers=opt.num_workers) num_classes = dataset.num_classes()