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

Can now be trained with normals

Browse Source
This commit is contained in:
Si11ium
2019-08-09 13:32:55 +02:00
parent a501dcd6b0
commit 92117328ad
3 changed files with 37 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files
dataset/shapenet.py
+9 -8
View File
@@ -24,13 +24,15 @@ class CustomShapeNet(InMemoryDataset):
modes = {key: val for val, key in enumerate(['train', 'test', 'predict'])}
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):
pre_transform=None, headers=True, has_variations=False, refresh=False, labels_within=False,
with_normals=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'
#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
self.with_normals = with_normals
super(CustomShapeNet, self).__init__(root_dir, transform, pre_transform, pre_filter)
self.data, self.slices = self._load_dataset()
print("Initialized")
@@ -143,8 +145,10 @@ class CustomShapeNet(InMemoryDataset):
y_all = [-1] * points.shape[0]
y = torch.as_tensor(y_all, dtype=torch.int)
####################################
# This is where you define the keys
attr_dict = dict(y=y, pos=points[:, :3], normals=points[:, 3:6])
attr_dict = dict(y=y, pos=points[:, :3 if not self.with_normals else 6])
####################################
if self.collate_per_element:
data = Data(**attr_dict)
else:
@@ -197,16 +201,13 @@ class ShapeNetPartSegDataset(Dataset):
except ValueError:
choice = []
pos, normals, labels = data.pos[choice, :], data.normals[choice, :], data.y[choice]
# pos, labels = data.pos[choice, :], data.y[choice]
pos, labels = data.pos[choice, :], data.y[choice]
labels -= 1 if self.num_classes() in labels else 0 # Map label from [1, C] to [0, C-1]
sample = {
'points': torch.cat([pos, normals], dim=1), # torch.Tensor (n, 6)
'labels': labels, # torch.Tensor (n,)
'pos': pos, # torch.Tensor (n, 3)
'normals': normals # torch.Tensor (n, 3)
'points': pos, # torch.Tensor (n, 6)
'labels': labels # torch.Tensor (n,)
}
return sample