steffen/Labeling-Primitives-With-Point2Net
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Labels can now be placed along next to the points within the datasetfile

Browse Source
This commit is contained in:
Si11ium
2019-08-06 08:36:59 +02:00
parent 54a5b48ddc
commit 97e36df1ba
2 changed files with 41 additions and 43 deletions
Download Patch File Download Diff File Expand all files Collapse all files

45
dataset/shapenet.py
View File

@ -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]