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
Show Stats 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: for element in paths:
# This was build to filter all variations that aregreater then 25 # 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]): if pattern.match(os.path.split(element)[-1]):
continue continue
else: else:
@ -134,8 +134,11 @@ class CustomShapeNet(InMemoryDataset):
points = points[:, :-1] points = points[:, :-1]
else: else:
# Get the y - Label # Get the y - Label
y_raw = next(i for i, v in enumerate(self.categories.keys()) if v.lower() in element.lower()) if self.mode != 'predict':
y_all = ([y_raw] if self.mode != 'predict' else [-1]) * points.shape[0] 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) y = torch.as_tensor(y_all, dtype=torch.int)
if self.collate_per_element: if self.collate_per_element:

20
predict/predict.py
View File

@ -6,15 +6,22 @@ from dataset.shapenet import ShapeNetPartSegDataset
from model.pointnet2_part_seg import PointNet2PartSegmentNet from model.pointnet2_part_seg import PointNet2PartSegmentNet
import torch_geometric.transforms as GT import torch_geometric.transforms as GT
import torch import torch
from distutils.util import strtobool
import numpy as np import numpy as np
import argparse import argparse
## ##
parser = argparse.ArgumentParser() parser = argparse.ArgumentParser()
parser.add_argument('--dataset', type=str, default='data', help='dataset path')
parser.add_argument('--npoints', type=int, default=2048, help='resample points number') parser.add_argument('--npoints', type=int, default=2048, help='resample points number')
parser.add_argument('--model', type=str, default='./checkpoint/seg_model_custom_246.pth', help='model path') parser.add_argument('--model', type=str, default='./checkpoint/seg_model_custom_3.pth', help='model path')
parser.add_argument('--sample_idx', type=int, default=0, help='select a sample to segment and view result') parser.add_argument('--sample_idx', type=int, default=0, help='select a sample to segment and view result')
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')
parser.add_argument('--has_variations', type=strtobool, default=False,
help='whether a single pointcloud has variations '
'named int(id)_pc.(xyz|dat) look at pointclouds or sub')
opt = parser.parse_args() opt = parser.parse_args()
print(opt) print(opt)
@ -26,11 +33,14 @@ if __name__ == '__main__':
test_dataset = ShapeNetPartSegDataset( test_dataset = ShapeNetPartSegDataset(
mode='predict', mode='predict',
root_dir=opt.dataset, root_dir='data',
transform=None,
npoints=opt.npoints, npoints=opt.npoints,
refresh=True refresh=True,
collate_per_segment=opt.collate_per_segment,
has_variations=opt.has_variations,
headers=opt.headers
) )
num_classes = test_dataset.num_classes() num_classes = test_dataset.num_classes()
print('test dataset size: ', len(test_dataset)) print('test dataset size: ', len(test_dataset))