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

Added normals to prediction DataObject

Browse Source
This commit is contained in:
Si11ium
2019-08-09 12:35:55 +02:00
parent 8eb165f76c
commit 39e5d72226
3 changed files with 21 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files
dataset/shapenet.py
+16 -18
View File
@@ -38,11 +38,11 @@ class CustomShapeNet(InMemoryDataset):
@property
def raw_file_names(self):
# Maybe add more data like validation sets
return list(self.modes.keys())
return [self.mode]
@property
def processed_file_names(self):
return [f'{x}.pt' for x in self.raw_file_names]
return [f'{self.mode}.pt']
def download(self):
dir_count = len([name for name in os.listdir(self.raw_dir) if os.path.isdir(os.path.join(self.raw_dir, name))])
@@ -58,7 +58,7 @@ class CustomShapeNet(InMemoryDataset):
def _load_dataset(self):
data, slices = None, None
filepath = self.processed_paths[self.modes[self.mode]]
filepath = self.processed_paths[0]
if self.refresh:
try:
os.remove(filepath)
@@ -91,7 +91,7 @@ class CustomShapeNet(InMemoryDataset):
def process(self, delimiter=' '):
datasets = defaultdict(list)
idx, data_folder = self.modes[self.mode], self.raw_file_names[self.modes[self.mode]]
idx, data_folder = self.modes[self.mode], self.raw_file_names[0]
path_to_clouds = os.path.join(self.raw_dir, data_folder)
if '.headers' in os.listdir(path_to_clouds):
@@ -111,8 +111,8 @@ class CustomShapeNet(InMemoryDataset):
paths.extend(glob.glob(os.path.join(pointcloud.path, f'*.{ext}')))
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+?|\d+?_pc)\.(xyz|dat)$')
# This was build to filter all full clouds
pattern = re.compile('^\d+?_pc\.(xyz|dat)$')
if pattern.match(os.path.split(element)[-1]):
continue
else:
@@ -143,9 +143,8 @@ class CustomShapeNet(InMemoryDataset):
y_all = [-1] * points.shape[0]
y = torch.as_tensor(y_all, dtype=torch.int)
attr_dict = dict(y=y, pos=points[:, :3])
if self.mode == 'predict':
attr_dict.update(normals=points[:, 3:6])
# This is where you define the keys
attr_dict = dict(y=y, pos=points[:, :3]) # , normals=points[:, 3:6])
if self.collate_per_element:
data = Data(**attr_dict)
else:
@@ -162,14 +161,14 @@ class CustomShapeNet(InMemoryDataset):
cloud_variations[int(os.path.split(element)[-1].split('_')[0])].append(data)
if not self.collate_per_element:
if self.has_variations:
for variation in cloud_variations.keys():
for _ in cloud_variations.keys():
datasets[data_folder].append(Data(**{key: torch.cat(data[key]) for key in data.keys()}))
else:
datasets[data_folder].append(Data(**{key: torch.cat(data[key]) for key in data.keys()}))
if datasets[data_folder]:
os.makedirs(self.processed_dir, exist_ok=True)
torch.save(self.collate(datasets[data_folder]), self.processed_paths[idx])
torch.save(self.collate(datasets[data_folder]), self.processed_paths[0])
def __repr__(self):
return f'{self.__class__.__name__}({len(self)})'
@@ -198,18 +197,17 @@ class ShapeNetPartSegDataset(Dataset):
except ValueError:
choice = []
points, labels = data.pos[choice, :], data.y[choice]
# pos, normals, labels = data.pos[choice, :], data.normals[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': points, # torch.Tensor (n, 3)
'labels': labels # torch.Tensor (n,)
'points': torch.cat([pos], dim=1), # torch.Tensor (n, 6)
'labels': labels, # torch.Tensor (n,)
# 'pos': pos, # torch.Tensor (n, 3)
# 'normals': normals # torch.Tensor (n, 3)
}
if self.mode == 'predict':
normals = data.normals[choice, :]
sample.update(normals=normals)
return sample
def __len__(self):