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

File based header detection, collate_per_PC training.

Browse Source
This commit is contained in:
Si11ium
2019-08-01 18:16:17 +02:00
parent a9bf053794
commit b7d127e840
2 changed files with 7 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
dataset/shapenet.py
View File

@ -287,10 +287,11 @@ 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):
def __init__(self, root_dir, train=False, transform=None, npoints=2048, headers=True, collate_per_segment=False):
super(PredictNetPartSegDataset, self).__init__()
self.npoints = npoints
self.dataset = ShapeNetPartSegDataset(root=root_dir, train=False, transform=transform, headers=headers)
self.dataset = PredictionShapeNet(root=root_dir, train=train, transform=transform,
headers=headers, collate_per_segment=collate_per_segment)
def __getitem__(self, index):
data = self.dataset[index]