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 14:33:30 +02:00
parent 97e36df1ba
commit 7443be4c40
2 changed files with 21 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
dataset/shapenet.py
View File

@ -111,7 +111,7 @@ class CustomShapeNet(InMemoryDataset):
for element in paths:
# This was build to filter all variations that aregreater then 25
pattern = re.compile('^((6[0-1]|[1-5][0-9])_\w+?\d+?|pc|\d+?_pc)\.(xyz|dat)$')
pattern = re.compile('^((6[0-1]|[1-5][0-9])_\w+?\d+?|\d+?_pc)\.(xyz|dat)$')
if pattern.match(os.path.split(element)[-1]):
continue
else:
@ -134,8 +134,11 @@ class CustomShapeNet(InMemoryDataset):
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]
if self.mode != 'predict':
y_raw = next(i for i, v in enumerate(self.categories.keys()) if v.lower() in element.lower())
y_all = [y_raw] * points.shape[0]
else:
y_all = [-1] * points.shape[0]
y = torch.as_tensor(y_all, dtype=torch.int)
if self.collate_per_element: